// Package main implements the Go OTel Debugger server.
package main

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

	"git.tornberg.me/go-otel/internal/auth"
	"git.tornberg.me/go-otel/internal/collector"
	"git.tornberg.me/go-otel/internal/honeypot"
	kube "git.tornberg.me/go-otel/internal/kubernetes"
	"git.tornberg.me/go-otel/internal/prometheus"
	"git.tornberg.me/go-otel/internal/sourcemaps"
	"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) {
	if _, ok := s.(*stats.ConnEnd); ok {
		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
}

// withCORS allows cross-origin ingestion (e.g. from the otel-rum.js browser
// script running on a third-party site like event-agent.k6n.net) into the
// OTLP HTTP endpoints. It is deliberately permissive on origin and headers,
// scoped by the ingest key/rate limiting instead of same-origin policy.
func withCORS(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS")
		reqHeaders := r.Header.Get("Access-Control-Request-Headers")
		if reqHeaders != "" {
			w.Header().Set("Access-Control-Allow-Headers", reqHeaders)
		} else {
			w.Header().Set("Access-Control-Allow-Headers", "Content-Type, X-Otel-Ingest-Key, Authorization, Baggage, Traceparent, Tracestate, X-Requested-With")
		}
		w.Header().Set("Access-Control-Max-Age", "86400")
		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusNoContent)
			return
		}
		next(w, r)
	}
}

// setupStore returns the telemetry store: ClickHouse when OTEL_CLICKHOUSE_ADDR
// is set (host:port, native protocol), otherwise the JSONL file store.
func setupStore(ctx context.Context) storage.Store {
	addr := os.Getenv("OTEL_CLICKHOUSE_ADDR")
	if addr == "" {
		logPath := os.Getenv("OTELLOG_PATH")
		if logPath == "" {
			logPath = "otel-agent.log"
		}
		log.Printf("storage: file store at %s (set OTEL_CLICKHOUSE_ADDR for ClickHouse)", logPath)
		return storage.NewFileStore(logPath)
	}

	chStore, err := storage.NewClickHouseStore(ctx, storage.ClickHouseConfig{
		Addr:     addr,
		Database: os.Getenv("OTEL_CLICKHOUSE_DB"),
		Username: os.Getenv("OTEL_CLICKHOUSE_USER"),
		Password: os.Getenv("OTEL_CLICKHOUSE_PASSWORD"),
	})
	if err != nil {
		log.Printf("storage: ClickHouse unavailable (%v), falling back to file store", err)
		logPath := os.Getenv("OTELLOG_PATH")
		if logPath == "" {
			logPath = "otel-agent.log"
		}
		return storage.NewFileStore(logPath)
	}
	log.Printf("storage: ClickHouse at %s", addr)
	return chStore
}

func main() {

	// Configure history size (default 100, can be made configurable)
	historySize := 100
	store := setupStore(context.Background())
	coll := collector.New(nil, historySize, store) // Temporary nil for wsServer
	kubeWatcher, kubeErr := kube.NewFromEnv(context.Background())
	if kubeErr != nil {
		log.Printf("kubernetes: rollout watcher disabled: %v", kubeErr)
	} else if kubeWatcher != nil {
		coll.SetKubernetesWatcher(kubeWatcher)
		log.Printf("kubernetes: rollout watcher enabled")
	}
	coll.SetPrometheusClient(prometheus.NewFromEnv())
	sourceMaps, err := sourcemaps.New(os.Getenv("OTEL_SOURCEMAP_DIR"))
	if err != nil {
		log.Fatalf("failed to initialize source maps: %v", err)
	}
	coll.SetSourceMaps(sourceMaps)
	// Aggregate stats endpoints are enabled only for stores that support
	// them (ClickHouse); they 501 politely on the file store.
	if sq, ok := store.(storage.StatsQuerier); ok {
		coll.SetStatsQuerier(sq)
	}
	wsServer := websocket.New(coll)
	honeypotTracker := honeypot.New(honeypot.ConfigFromEnv())
	// Update collector with wsServer
	coll.SetBroadcaster(wsServer)
	ingestKey := os.Getenv("OTEL_INGEST_KEY")
	coll.SetIngestKey(ingestKey)
	oidcAuth := auth.NewOidcAuth(ingestKey)
	log.Printf("auth: OIDC issuer=%s client=%s disabled=%v", oidcAuth.Issuer(), oidcAuth.ClientID(), oidcAuth.IsDisabled())

	// 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", withCORS(coll.HandleTraces))
	otelMux.HandleFunc("/v1/traces/", withCORS(coll.HandleTraces))
	otelMux.HandleFunc("/v1/metrics", withCORS(coll.HandleMetrics))
	otelMux.HandleFunc("/v1/metrics/", withCORS(coll.HandleMetrics))
	otelMux.HandleFunc("/v1/logs", withCORS(coll.HandleLogs))
	otelMux.HandleFunc("/v1/logs/", withCORS(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.
	// CORS is enabled here since this is the publicly exposed port (via the
	// ingress) that a browser script such as static/otel-rum.js posts to.
	uiMux.HandleFunc("/v1/traces", withCORS(coll.HandleTraces))
	uiMux.HandleFunc("/v1/traces/", withCORS(coll.HandleTraces))
	uiMux.HandleFunc("/v1/metrics", withCORS(coll.HandleMetrics))
	uiMux.HandleFunc("/v1/metrics/", withCORS(coll.HandleMetrics))
	uiMux.HandleFunc("/v1/logs", withCORS(coll.HandleLogs))
	uiMux.HandleFunc("/v1/logs/", withCORS(coll.HandleLogs))
	uiMux.HandleFunc("/health", coll.HandleHealth)
	uiMux.HandleFunc("GET /internal/honeypot/ips", honeypotTracker.HandleIPs)

	// Public OIDC endpoints
	uiMux.HandleFunc("POST /api/auth/token", oidcAuth.HandleTokenExchange)
	uiMux.HandleFunc("GET /api/auth/config", oidcAuth.HandleConfig)
	uiMux.Handle("/api/honeypot/ips", oidcAuth.RequireAuth(http.HandlerFunc(honeypotTracker.HandleIPsForUser)))
	// Source-map uploads are intentionally public so the release build script
	// can upload maps without an OIDC session or bearer token.
	uiMux.HandleFunc("POST /api/sourcemaps", coll.HandleSourceMap)
	uiMux.HandleFunc("GET /api/needle/status", coll.HandleNeedleStatus)
	uiMux.HandleFunc("POST /api/needle/classify", coll.HandleNeedleClassify)

	// Protected API & WebSocket routes
	uiMux.Handle("/api/history", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleHistory)))
	uiMux.Handle("/api/history/traces", oidcAuth.RequireAuth(coll.HandleHistoryByType("traces")))
	uiMux.Handle("/api/history/metrics", oidcAuth.RequireAuth(coll.HandleHistoryByType("metrics")))
	uiMux.Handle("/api/history/logs", oidcAuth.RequireAuth(coll.HandleHistoryByType("logs")))
	uiMux.Handle("/api/history/search", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleSearch)))
	uiMux.Handle("/api/history/wait", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleWait)))
	uiMux.Handle("GET /api/traces/{traceId}", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleTraceByID)))
	uiMux.Handle("/api/issues", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleIssues)))
	uiMux.Handle("POST /api/ml/test-spike", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleMLDetectorTest)))
	uiMux.Handle("GET /api/issues/{id}", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleIssueByID)))
	uiMux.Handle("GET /api/issues/{id}/similar", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleIssueSimilar)))
	uiMux.Handle("POST /api/issues/{id}/similar", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleSimilarityFeedback)))
	uiMux.Handle("GET /api/patterns", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandlePatterns)))
	uiMux.Handle("GET /api/rollouts", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleRollouts)))
	uiMux.Handle("GET /api/rollouts/{id}", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleRolloutByID)))
	uiMux.Handle("GET /api/rollouts/{id}/impact", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleRolloutImpact)))
	// Aggregate stats (ClickHouse-backed); return 501 with an explanatory
	// JSON body when only the file store is configured.
	uiMux.Handle("/api/stats/errors", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleStatsErrors)))
	uiMux.Handle("/api/stats/releases", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleStatsReleases)))
	uiMux.Handle("/api/stats/services", oidcAuth.RequireAuth(http.HandlerFunc(coll.HandleStatsServices)))
	uiMux.HandleFunc("/ws", wsServer.HandleConnections)
	uiMux.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.Dir("./static"))))
	// Serve the agent instructions doc as a static route so an AI agent can
	// fetch it directly from the base URL before making any API calls.
	uiMux.HandleFunc("GET /agents_example.md", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "agents_example.md")
	})
	uiMux.HandleFunc("GET /instrumentation.md", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "instrumentation.md")
	})
	uiMux.HandleFunc("GET /static/instrumentation.md", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "instrumentation.md")
	})
	distDir := "./dist"
	if _, err := os.Stat(distDir); os.IsNotExist(err) {
		if _, err := os.Stat("./otel-ui/dist"); err == nil {
			distDir = "./otel-ui/dist"
		} else if _, err := os.Stat("./index.html"); err == nil {
			// Development fallback: serve the checked-in visualizer when no
			// separately built frontend exists.
			distDir = "."
		}
	}
	uiMux.Handle("/", spaHandler(http.Dir(distDir)))

	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)
		}
	}()

	httpPort := os.Getenv("PORT")
	if httpPort == "" {
		httpPort = "8080"
	}
	uiHandler := honeypotTracker.Middleware(uiMux)
	if err := http.ListenAndServe(":"+httpPort, uiHandler); err != nil {
		log.Fatalf("failed to start HTTP server on :%s: %v", httpPort, 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)
	})
}