AIRouter/internal/server/server.go

package server

import (
	"context"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/edgeai/gateway/internal/adapter"
	"github.com/edgeai/gateway/internal/auth"
	"github.com/edgeai/gateway/internal/config"
	"github.com/edgeai/gateway/internal/handler"
	"github.com/edgeai/gateway/internal/middleware"
	"github.com/edgeai/gateway/internal/observability"
	"github.com/edgeai/gateway/internal/router"
	"github.com/edgeai/gateway/internal/scheduler"
	"github.com/edgeai/gateway/internal/session"
)

// Server is the main HTTP server for the AI gateway.
type Server struct {
	cfg       *config.Config
	logger    *observability.Logger
	metrics   *observability.Metrics
	HTTPSrv   *http.Server
	auth      *auth.Authenticator
	registry  *adapter.Registry
	modelMap  *router.LogicalModelMapping
	scheduler *scheduler.Scheduler
	sessions  *session.Store
}

// New creates a new Server instance with all components wired.
func New(cfg *config.Config, logger *observability.Logger) (*Server, error) {
	// Ensure data directory exists
	dbPath := extractDBPath(cfg.Storage.SessionDB)
	if dbPath != "" {
		os.MkdirAll(filepath.Dir(dbPath), 0755)
	}

	// Initialize auth
	authPath := filepath.Join(filepath.Dir(dbPath), "auth.db")
	authenticator, err := auth.NewAuthenticator(authPath, logger)
	if err != nil {
		return nil, fmt.Errorf("init auth: %w", err)
	}

	// Initialize session store
	sessionStore, err := session.NewStore(dbPath)
	if err != nil {
		return nil, fmt.Errorf("init session store: %w", err)
	}

	// Initialize adapter registry
	registry := adapter.NewRegistry()

	// Initialize logical model mapping
	modelMap := router.NewLogicalModelMapping(cfg)

	// Register adapters for each unique endpoint
	registered := make(map[string]bool)
	for _, mc := range cfg.Models {
		key := mc.Provider + "|" + mc.Endpoint
		if !registered[key] {
			switch mc.Provider {
			case "ollama":
				registry.Register(mc.Provider, adapter.NewOllamaAdapter(mc.Endpoint))
			case "vllm":
				registry.Register(mc.Provider, adapter.NewVLLMAdapter(mc.Endpoint))
			}
			registered[key] = true
		}
	}

	// Initialize scheduler
	sched := scheduler.NewScheduler(&cfg.Scheduler, logger)

	// Initialize metrics
	metrics := observability.NewMetrics()

	s := &Server{
		cfg:       cfg,
		logger:    logger,
		metrics:   metrics,
		auth:      authenticator,
		registry:  registry,
		modelMap:  modelMap,
		scheduler: sched,
		sessions:  sessionStore,
	}

	mux := http.NewServeMux()
	s.registerRoutes(mux)

	// Apply middleware chain (order: Recovery → Logging → RequestID → BodyLimit → Auth → handler)
	h := middleware.RequestID(mux)
	h = middleware.BodyLimit(cfg.Server.MaxRequestBodyMB)(h)
	h = s.auth.Middleware(h)
	h = middleware.Logging(logger)(h)
	h = middleware.Recovery(logger)(h)

	s.HTTPSrv = &http.Server{
		Addr:         fmt.Sprintf("%s:%d", cfg.Server.Host, cfg.Server.Port),
		Handler:      h,
		ReadTimeout:  30 * time.Second,
		WriteTimeout: 0, // no write timeout for SSE
		IdleTimeout:  120 * time.Second,
	}

	return s, nil
}

func (s *Server) registerRoutes(mux *http.ServeMux) {
	// Health and readiness
	mux.HandleFunc("/health", s.handleHealth)
	mux.HandleFunc("/ready", s.handleReady)

	// Metrics
	mux.HandleFunc(s.cfg.Observability.MetricsPath, s.metrics.Handler())

	// OpenAI-compatible API
	mux.HandleFunc("/v1/chat/completions", s.handleChatCompletions)
	mux.HandleFunc("/v1/models", s.handleModels)

	// Session management
	mux.HandleFunc("/v1/sessions", s.handleSessions)
	mux.HandleFunc("/v1/sessions/", s.handleSessionByID)
}

// Authenticator returns the authenticator instance (for testing/management).
func (s *Server) Authenticator() *auth.Authenticator {
	return s.auth
}

// Start begins listening for HTTP requests.
func (s *Server) Start() error {
	s.logger.Info("http server starting", observability.F().
		Event("server_start").
		Set("addr", s.HTTPSrv.Addr))
	return s.HTTPSrv.ListenAndServe()
}

// Shutdown gracefully shuts down the server.
func (s *Server) Shutdown() error {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	s.scheduler.Stop()
	if s.sessions != nil {
		s.sessions.Close()
	}
	if s.auth != nil {
		s.auth.Close()
	}

	return s.HTTPSrv.Shutdown(ctx)
}

func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
	handler.WriteJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}

func (s *Server) handleReady(w http.ResponseWriter, r *http.Request) {
	w.Header().Set("Content-Type", "application/json")

	ready := true
	reasons := []string{}

	for _, name := range s.registry.Names() {
		a, _ := s.registry.Get(name)
		if err := a.HealthCheck(r.Context()); err != nil {
			ready = false
			reasons = append(reasons, fmt.Sprintf("%s: %v", name, err))
		}
	}

	if ready {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(`{"status":"ready"}`))
	} else {
		w.WriteHeader(http.StatusServiceUnavailable)
		fmt.Fprintf(w, `{"status":"not_ready","reasons":["%s"]}`, strings.Join(reasons, `","`))
	}
}

func extractDBPath(connStr string) string {
	if strings.HasPrefix(connStr, "sqlite://") {
		return strings.TrimPrefix(connStr, "sqlite://")
	}
	return connStr
}