package main

import (
	"context"
	"log"
	"net"
	"net/http"
	"strings"

	"git.tornberg.me/go-otel/internal/collector"
	"git.tornberg.me/go-otel/internal/storage"
	"git.tornberg.me/go-otel/internal/websocket"
	"go.opentelemetry.io/collector/pdata/plog"
	"go.opentelemetry.io/collector/pdata/pmetric"
	"go.opentelemetry.io/collector/pdata/ptrace"
	logsgrpc "go.opentelemetry.io/proto/otlp/collector/logs/v1"
	metricsgrpc "go.opentelemetry.io/proto/otlp/collector/metrics/v1"
	tracegrpc "go.opentelemetry.io/proto/otlp/collector/trace/v1"
	"google.golang.org/grpc"
	"google.golang.org/grpc/stats"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/proto"
)

// gRPC Logging Interceptor
func loggingInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
	resp, err := handler(ctx, req)
	if err != nil {
		st, _ := status.FromError(err)
		log.Printf("gRPC ERROR [%s] code=%s message=%s", info.FullMethod, st.Code(), st.Message())
	}
	//else {
	// log.Printf("gRPC SUCCESS [%s]", info.FullMethod)
	//}
	return resp, err
}

// gRPC Connection Stats Handler
type connStatsHandler struct {
	stats.Handler
}

func (h *connStatsHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
	log.Printf("gRPC CONN [OPEN] from=%s", info.RemoteAddr)
	return ctx
}

func (h *connStatsHandler) HandleConn(ctx context.Context, s stats.ConnStats) {
	switch s.(type) {
	case *stats.ConnEnd:
		log.Printf("gRPC CONN [CLOSED]")
	}
}

func (h *connStatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
	return ctx
}
func (h *connStatsHandler) HandleRPC(ctx context.Context, s stats.RPCStats) {}

type traceService struct {
	tracegrpc.UnimplementedTraceServiceServer
	collector *collector.Collector
}

func (s *traceService) Export(ctx context.Context, req *tracegrpc.ExportTraceServiceRequest) (*tracegrpc.ExportTraceServiceResponse, error) {
	// log.Printf("Received gRPC traces %d", len(req.ResourceSpans))

	// Convert protobuf to bytes then to pdata
	data, err := proto.Marshal(req)
	if err != nil {
		log.Printf("failed to marshal trace request: %v", err)
		return nil, err
	}

	unmarshaler := &ptrace.ProtoUnmarshaler{}
	traces, err := unmarshaler.UnmarshalTraces(data)
	if err != nil {
		log.Printf("failed to unmarshal traces: %v", err)
		return nil, err
	}

	if err := s.collector.ConsumeTraces(ctx, traces); err != nil {
		log.Printf("failed to consume traces: %v", err)
		return nil, err
	}

	return &tracegrpc.ExportTraceServiceResponse{}, nil
}

type metricsService struct {
	metricsgrpc.UnimplementedMetricsServiceServer
	collector *collector.Collector
}

func (s *metricsService) Export(ctx context.Context, req *metricsgrpc.ExportMetricsServiceRequest) (*metricsgrpc.ExportMetricsServiceResponse, error) {
	// log.Printf("Received gRPC metrics %d", len(req.ResourceMetrics))

	// Convert protobuf to bytes then to pdata
	data, err := proto.Marshal(req)
	if err != nil {
		log.Printf("failed to marshal metrics request: %v", err)
		return nil, err
	}

	unmarshaler := &pmetric.ProtoUnmarshaler{}
	metrics, err := unmarshaler.UnmarshalMetrics(data)
	if err != nil {
		log.Printf("failed to unmarshal metrics: %v", err)
		return nil, err
	}

	if err := s.collector.ConsumeMetrics(ctx, metrics); err != nil {
		log.Printf("failed to consume metrics: %v", err)
		return nil, err
	}

	return &metricsgrpc.ExportMetricsServiceResponse{}, nil
}

type logsService struct {
	logsgrpc.UnimplementedLogsServiceServer
	collector *collector.Collector
}

func (s *logsService) Export(ctx context.Context, req *logsgrpc.ExportLogsServiceRequest) (*logsgrpc.ExportLogsServiceResponse, error) {
	// log.Printf("Received gRPC logs %d", len(req.ResourceLogs))

	// Convert protobuf to bytes then to pdata
	data, err := proto.Marshal(req)
	if err != nil {
		log.Printf("failed to marshal logs request: %v", err)
		return nil, err
	}

	unmarshaler := &plog.ProtoUnmarshaler{}
	logs, err := unmarshaler.UnmarshalLogs(data)
	if err != nil {
		log.Printf("failed to unmarshal logs: %v", err)
		return nil, err
	}

	if err := s.collector.ConsumeLogs(ctx, logs); err != nil {
		log.Printf("failed to consume logs: %v", err)
		return nil, err
	}

	return &logsgrpc.ExportLogsServiceResponse{}, nil
}

func main() {

	// Configure history size (default 100, can be made configurable)
	historySize := 100
	store := storage.NewFileStore("otel-agent.log")
	coll := collector.New(nil, historySize, store) // Temporary nil for wsServer
	wsServer := websocket.New(coll)
	// Update collector with wsServer
	coll.SetBroadcaster(wsServer)

	// Start gRPC server
	go func() {
		lis, err := net.Listen("tcp", ":4317")
		if err != nil {
			log.Fatal("failed to listen for gRPC", err)
		}

		s := grpc.NewServer(
			grpc.UnaryInterceptor(loggingInterceptor),
			grpc.StatsHandler(&connStatsHandler{}),
		)
		tracegrpc.RegisterTraceServiceServer(s, &traceService{collector: coll})
		metricsgrpc.RegisterMetricsServiceServer(s, &metricsService{collector: coll})
		logsgrpc.RegisterLogsServiceServer(s, &logsService{collector: coll})

		log.Print("Starting gRPC server on port 4317")
		if err := s.Serve(lis); err != nil {
			log.Fatalf("failed to serve gRPC: %v", err)
		}
	}()

	// OTLP HTTP mux (standard OTel port 4318)
	otelMux := http.NewServeMux()
	otelMux.HandleFunc("/v1/traces", coll.HandleTraces)
	otelMux.HandleFunc("/v1/metrics", coll.HandleMetrics)
	otelMux.HandleFunc("/v1/logs", coll.HandleLogs)
	otelMux.HandleFunc("/health", coll.HandleHealth)

	// UI and API mux (port 8080)
	uiMux := http.NewServeMux()
	// Also register OTLP endpoints on 8080 for backwards compatibility
	uiMux.HandleFunc("/v1/traces", coll.HandleTraces)
	uiMux.HandleFunc("/v1/metrics", coll.HandleMetrics)
	uiMux.HandleFunc("/v1/logs", coll.HandleLogs)
	uiMux.HandleFunc("/health", coll.HandleHealth)

	uiMux.HandleFunc("/api/history", coll.HandleHistory)
	uiMux.HandleFunc("/api/history/traces", coll.HandleHistoryByType("traces"))
	uiMux.HandleFunc("/api/history/metrics", coll.HandleHistoryByType("metrics"))
	uiMux.HandleFunc("/api/history/logs", coll.HandleHistoryByType("logs"))
	uiMux.HandleFunc("/api/history/search", coll.HandleSearch)
	uiMux.HandleFunc("/api/history/wait", coll.HandleWait)
	uiMux.HandleFunc("/ws", wsServer.HandleConnections)
	uiMux.Handle("/", spaHandler(http.Dir("./dist")))

	log.Printf("Starting HTTP server on port 8080 (UI/API)")
	log.Printf("OTLP HTTP endpoints (also on port 8080):")
	log.Printf("  Traces:  http://localhost:8080/v1/traces")
	log.Printf("  Metrics: http://localhost:8080/v1/metrics")
	log.Printf("  Logs:    http://localhost:8080/v1/logs")
	log.Printf("OTLP gRPC endpoint:")
	log.Printf("  Traces/Metrics/Logs: localhost:4317")
	log.Printf("WebSocket: ws://localhost:8080/ws")

	// Start OTLP HTTP server on port 4318
	go func() {
		log.Printf("Starting OTLP HTTP server on port 4318")
		if err := http.ListenAndServe(":4318", otelMux); err != nil {
			log.Printf("failed to start OTLP HTTP server: %v", err)
		}
	}()

	if err := http.ListenAndServe(":8080", uiMux); err != nil {
		log.Fatalf("failed to start HTTP server: %v", err)
	}
}
func spaHandler(root http.FileSystem) http.Handler {
	fileServer := http.FileServer(root)
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		path := r.URL.Path
		// Skip API, WebSocket, and OTLP endpoints
		if !strings.HasPrefix(path, "/api") && !strings.HasPrefix(path, "/ws") && !strings.HasPrefix(path, "/v1") {
			if _, err := root.Open(path); err != nil {
				// File does not exist, serve index.html
				r.URL.Path = "/"
			}
		}
		fileServer.ServeHTTP(w, r)
	})
}