ez-api/cmd/server/main.go

package main

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

	"github.com/ez-api/ez-api/internal/api"
	"github.com/ez-api/ez-api/internal/config"
	"github.com/ez-api/ez-api/internal/middleware"
	"github.com/ez-api/ez-api/internal/model"
	"github.com/ez-api/ez-api/internal/service"
	"github.com/gin-gonic/gin"
	"github.com/redis/go-redis/v9"
	"gorm.io/driver/postgres"
	"gorm.io/gorm"
)

func main() {
	// 1. Load Configuration
	cfg, err := config.Load()
	if err != nil {
		log.Fatalf("Failed to load config: %v", err)
	}

	// 2. Initialize Redis Client
	rdb := redis.NewClient(&redis.Options{
		Addr:     cfg.Redis.Addr,
		Password: cfg.Redis.Password,
		DB:       cfg.Redis.DB,
	})

	// Verify Redis connection
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := rdb.Ping(ctx).Err(); err != nil {
		log.Fatalf("Failed to connect to Redis: %v", err)
	}
	log.Println("Connected to Redis successfully")

	// 3. Initialize GORM (PostgreSQL)
	db, err := gorm.Open(postgres.Open(cfg.Postgres.DSN), &gorm.Config{})
	if err != nil {
		log.Fatalf("Failed to connect to PostgreSQL: %v", err)
	}
	sqlDB, err := db.DB()
	if err != nil {
		log.Fatalf("Failed to get generic database object: %v", err)
	}
	// Verify DB connection
	if err := sqlDB.Ping(); err != nil {
		log.Fatalf("Failed to ping PostgreSQL: %v", err)
	}
	log.Println("Connected to PostgreSQL successfully")

	// Auto Migrate
	if err := db.AutoMigrate(&model.Master{}, &model.Key{}, &model.Provider{}, &model.Model{}, &model.LogRecord{}); err != nil {
		log.Fatalf("Failed to auto migrate: %v", err)
	}

	// 4. Setup Services and Handlers
	syncService := service.NewSyncService(rdb)
	logWriter := service.NewLogWriter(db, cfg.Log.QueueCapacity, cfg.Log.BatchSize, cfg.Log.FlushInterval)
	logCtx, cancelLogs := context.WithCancel(context.Background())
	defer cancelLogs()
	logWriter.Start(logCtx)

	adminService, err := service.NewAdminService()
	if err != nil {
		log.Fatalf("Failed to create admin service: %v", err)
	}
	masterService := service.NewMasterService(db)
	healthService := service.NewHealthCheckService(db, rdb)

	handler := api.NewHandler(db, syncService, logWriter)
	adminHandler := api.NewAdminHandler(masterService)
	masterHandler := api.NewMasterHandler(masterService)

	// 4.1 Prime Redis snapshots so DP can start with data
	if err := syncService.SyncAll(db); err != nil {
		log.Printf("Initial sync warning: %v", err)
	}

	// 5. Setup Gin Router
	r := gin.Default()

	// CORS Middleware
	r.Use(func(c *gin.Context) {
		c.Writer.Header().Set("Access-Control-Allow-Origin", "*") // TODO: Restrict this in production
		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
		c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, accept, origin, Cache-Control, X-Requested-With")
		c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE")

		if c.Request.Method == "OPTIONS" {
			c.AbortWithStatus(204)
			return
		}

		c.Next()
	})

	// Health Check Endpoint
	r.GET("/health", func(c *gin.Context) {
		status := healthService.Check(c.Request.Context())
		httpStatus := http.StatusOK
		if status.Status == "down" {
			httpStatus = http.StatusServiceUnavailable
		}
		c.JSON(httpStatus, status)
	})

	// API Routes
	// Admin Routes
	adminGroup := r.Group("/admin")
	adminGroup.Use(middleware.AdminAuthMiddleware(adminService))
	{
		adminGroup.POST("/masters", adminHandler.CreateMaster)
		// Other admin routes for managing providers, models, etc.
		adminGroup.POST("/providers", handler.CreateProvider)
		adminGroup.POST("/models", handler.CreateModel)
		adminGroup.GET("/models", handler.ListModels)
		adminGroup.POST("/sync/snapshot", handler.SyncSnapshot)
	}

	// Master Routes
	masterGroup := r.Group("/v1")
	masterGroup.Use(middleware.MasterAuthMiddleware(masterService))
	{
		masterGroup.POST("/tokens", masterHandler.IssueChildKey)
	}

	// Public/General Routes (if any)
	r.POST("/logs", handler.IngestLog)

	srv := &http.Server{
		Addr:    ":" + cfg.Server.Port,
		Handler: r,
	}

	// 6. Start Server with Graceful Shutdown
	go func() {
		log.Printf("Starting ez-api on port %s", cfg.Server.Port)
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			log.Fatalf("Server failed: %v", err)
		}
	}()

	// Wait for interrupt signal
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
	log.Println("Shutting down server...")

	// Shutdown with timeout
	ctx, cancel = context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := srv.Shutdown(ctx); err != nil {
		log.Fatalf("Server forced to shutdown: %v", err)
	}

	log.Println("Server exited properly")
}