package main

import (
	"context"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"

	"github.com/seifghazi/claude-code-monitor/internal/config"
	"github.com/seifghazi/claude-code-monitor/internal/handler"
	"github.com/seifghazi/claude-code-monitor/internal/middleware"
	"github.com/seifghazi/claude-code-monitor/internal/provider"
	"github.com/seifghazi/claude-code-monitor/internal/runtime"
	"github.com/seifghazi/claude-code-monitor/internal/service"
)

func main() {
	logger := log.New(os.Stdout, "proxy: ", log.LstdFlags|log.Lshortfile)

	cfg, err := config.Load()
	if err != nil {
		logger.Fatalf("❌ Failed to load configuration: %v", err)
	}

	// Initialize providers
	providers := make(map[string]provider.Provider)
	providers["anthropic"] = provider.NewAnthropicProvider(&cfg.Providers.Anthropic)
	providers["openai"] = provider.NewOpenAIProvider(&cfg.Providers.OpenAI)

	// Initialize model router
	modelRouter := service.NewModelRouter(cfg, providers, logger)

	// Initialize storage based on DB_TYPE
	var storageService service.StorageService
	switch cfg.Storage.DBType {
	case "postgres", "postgresql":
		if cfg.Storage.DatabaseURL == "" {
			logger.Fatalf("❌ DATABASE_URL is required when DB_TYPE=postgres")
		}
		storageService, err = service.NewPostgresStorageService(&cfg.Storage)
		if err != nil {
			logger.Fatalf("❌ Failed to initialize PostgreSQL storage: %v", err)
		}
		logger.Println("🐘 PostgreSQL database ready")
	default:
		storageService, err = service.NewSQLiteStorageService(&cfg.Storage)
		if err != nil {
			logger.Fatalf("❌ Failed to initialize SQLite storage: %v", err)
		}
		logger.Println("🗿 SQLite database ready")
	}

	h := handler.New(storageService, logger, modelRouter, providers["anthropic"], cfg.Providers.Anthropic.DemoteNonstreaming)

	srv := &http.Server{
		Addr:         net.JoinHostPort(cfg.Server.Host, cfg.Server.Port),
		Handler:      buildHandler(cfg, h),
		ReadTimeout:  cfg.Server.ReadTimeout,
		WriteTimeout: cfg.Server.WriteTimeout,
		IdleTimeout:  cfg.Server.IdleTimeout,
	}

	go func() {
		logger.Println("🚀 Claude Code Proxy started")
		logger.Printf("   upstream: %s", cfg.Providers.Anthropic.BaseURL)
		logger.Printf("   storage:  %s", cfg.Storage.DBType)
		if cfg.Auth.Enabled {
			logger.Printf("   auth:     bearer / %s", cfg.Auth.APIKeyHeader)
		}
		if cfg.Auth.DashboardPassword != "" {
			logger.Printf("   dashboard: password protected")
		}

		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logger.Fatalf("❌ Server failed to start: %v", err)
		}
	}()

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	sig := <-quit

	// Drain ceiling: must accommodate the 30-min forward context in
	// handlers.Messages plus a small buffer. Container-level
	// stop_grace_period must be set just above this in docker-compose so
	// docker doesn't SIGKILL us mid-drain.
	const drainCeiling = 32 * time.Minute
	const drainPoll = 1 * time.Second

	logger.Printf("🛑 Received %s — draining (ceiling %s)...", sig, drainCeiling)
	runtime.SetDraining(true)

	drainStart := time.Now()
	lastLog := drainStart
	for {
		n := runtime.InFlight()
		if n == 0 {
			logger.Printf("✅ Drain complete (%s elapsed, in-flight=0)", time.Since(drainStart).Round(time.Second))
			break
		}
		if time.Since(drainStart) >= drainCeiling {
			logger.Printf("⏰ Drain ceiling hit — proceeding with shutdown (in_flight=%d)", n)
			break
		}
		if time.Since(lastLog) >= 10*time.Second {
			logger.Printf("⏳ Draining: in_flight=%d elapsed=%s", n, time.Since(drainStart).Round(time.Second))
			lastLog = time.Now()
		}
		time.Sleep(drainPoll)
	}

	// Server.Shutdown stops accepting new connections and waits for active
	// handlers to return. Use a short ceiling here — the long wait already
	// happened in the drain loop above. If a handler is somehow still
	// running past the drain ceiling, give it 30s before forced close.
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	if err := srv.Shutdown(ctx); err != nil {
		logger.Printf("❌ Server forced to shutdown: %v", err)
	}

	// Close storage service (checkpoints WAL, closes prepared statements)
	logger.Println("🗄️ Closing database...")
	if err := storageService.Close(); err != nil {
		logger.Printf("❌ Error closing storage: %v", err)
	}

	logger.Println("✅ Server exited")
}

func buildHandler(cfg *config.Config, h *handler.Handler) http.Handler {
	r := mux.NewRouter()

	corsHandler := handlers.CORS(
		handlers.AllowedOrigins(cfg.CORS.AllowedOrigins),
		handlers.AllowedMethods(cfg.CORS.AllowedMethods),
		handlers.AllowedHeaders(cfg.CORS.AllowedHeaders),
	)

	r.Use(middleware.Logging)
	r.Use(middleware.Auth(cfg.Auth))

	// Liveness/readiness — must remain unauthenticated and uncounted.
	r.HandleFunc("/health", h.Health).Methods("GET")
	r.HandleFunc("/livez", h.Livez).Methods("GET")
	r.HandleFunc("/openapi.json", h.OpenAPIJSON).Methods("GET")
	r.HandleFunc("/openapi.yaml", h.OpenAPIYAML).Methods("GET")

	// Proxy endpoints — no dashboard auth, but tracked by the in-flight
	// gauge so graceful shutdown can wait for them to drain.
	// Trailing slash on the prefix keeps the boundary tight: /v10 etc. won't match.
	v1 := r.PathPrefix("/v1/").Subrouter()
	v1.Use(middleware.InFlight)
	v1.HandleFunc("/chat/completions", h.ChatCompletions).Methods("POST")
	v1.HandleFunc("/messages", h.Messages).Methods("POST")
	// Catch-all: proxy any other /v1/* requests to Anthropic (config, quota, batches, etc.)
	v1.PathPrefix("/").HandlerFunc(h.ProxyPassthrough)

	// Dashboard/data endpoints — protected by DASHBOARD_PASSWORD basic auth
	dashboard := r.PathPrefix("/").Subrouter()
	dashboard.Use(middleware.DashboardAuth(cfg.Auth.DashboardPassword))
	dashboard.HandleFunc("/", h.UI).Methods("GET")
	dashboard.HandleFunc("/ui", h.UI).Methods("GET")
	dashboard.HandleFunc("/api/requests", h.GetRequests).Methods("GET")
	dashboard.HandleFunc("/api/requests", h.DeleteRequests).Methods("DELETE")
	dashboard.HandleFunc("/api/requests/summary", h.GetRequestsSummary).Methods("GET")
	dashboard.HandleFunc("/api/requests/latest-date", h.GetLatestRequestDate).Methods("GET")
	dashboard.HandleFunc("/api/requests/{id}", h.GetRequestByID).Methods("GET")
	dashboard.HandleFunc("/api/stats", h.GetStats).Methods("GET")
	dashboard.HandleFunc("/api/stats/dashboard", h.GetDashboardStats).Methods("GET")
	dashboard.HandleFunc("/api/stats/hourly", h.GetHourlyStats).Methods("GET")
	dashboard.HandleFunc("/api/stats/models", h.GetModelStats).Methods("GET")
	dashboard.HandleFunc("/api/stats/organizations", h.GetOrganizations).Methods("GET")
	dashboard.HandleFunc("/api/conversations", h.GetConversations).Methods("GET")
	dashboard.HandleFunc("/api/conversations/{id}", h.GetConversationByID).Methods("GET")
	dashboard.HandleFunc("/api/conversations/project", h.GetConversationsByProject).Methods("GET")
	dashboard.HandleFunc("/api/settings", h.GetSettings).Methods("GET")
	dashboard.HandleFunc("/api/settings", h.SaveSettings).Methods("PUT")

	r.NotFoundHandler = http.HandlerFunc(h.NotFound)

	return corsHandler(r)
}
