fetch_ml/internal/api/routes.go

package api

import (
	"net/http"
	"os"

	"github.com/jfraeys/fetch_ml/internal/api/audit"
	"github.com/jfraeys/fetch_ml/internal/api/datasets"
	"github.com/jfraeys/fetch_ml/internal/api/jobs"
	"github.com/jfraeys/fetch_ml/internal/api/jupyter"
	"github.com/jfraeys/fetch_ml/internal/api/plugins"
	sch "github.com/jfraeys/fetch_ml/internal/api/scheduler"
	"github.com/jfraeys/fetch_ml/internal/api/ws"
	"github.com/jfraeys/fetch_ml/internal/prommetrics"
	"github.com/labstack/echo/v4"
)

// registerRoutes sets up all HTTP routes and handlers
func (s *Server) registerRoutes(mux *http.ServeMux) {
	// Register Prometheus metrics endpoint (if enabled)
	if s.config.Monitoring.Prometheus.Enabled {
		s.promMetrics = prommetrics.New()
		s.logger.Info("prometheus metrics initialized")

		// Register metrics endpoint
		metricsPath := s.config.Monitoring.Prometheus.Path
		if metricsPath == "" {
			metricsPath = "/metrics"
		}
		mux.Handle(metricsPath, s.promMetrics.Handler())
		s.logger.Info("metrics endpoint registered", "path", metricsPath)
	}

	// Register health check endpoints (if enabled)
	if s.config.Monitoring.HealthChecks.Enabled {
		s.registerHealthRoutes(mux)
	}

	// Register WebSocket endpoint
	s.registerWebSocketRoutes(mux)

	// Register new REST API endpoints for TUI
	jobsHandler := jobs.NewHandler(
		s.expManager,
		s.logger,
		s.taskQueue,
		s.db,
		s.config.BuildAuthConfig(),
		nil,
	)

	// Experiment history endpoint: GET /api/experiments/:id/history
	mux.HandleFunc("GET /api/experiments/{id}/history", jobsHandler.GetExperimentHistoryHTTP)

	// Team jobs endpoint: GET /api/jobs?all_users=true
	mux.HandleFunc("GET /api/jobs", jobsHandler.ListAllJobsHTTP)

	// Register OpenAPI-generated routes with Echo router
	s.registerOpenAPIRoutes(mux, jobsHandler)

	// Register scheduler API endpoints
	schedulerHandler := sch.NewHandler(s.schedulerHub, s.logger, s.config.BuildAuthConfig())
	mux.HandleFunc("GET /api/v1/workers", schedulerHandler.GetV1SchedulerWorkers)
	mux.HandleFunc("GET /api/v1/workers/{workerId}", schedulerHandler.GetV1SchedulerWorkersWorkerID)
	mux.HandleFunc("DELETE /api/v1/workers/{workerId}", schedulerHandler.DeleteV1SchedulerWorkersWorkerID)

	// Register API documentation endpoint
	s.registerDocsRoutes(mux)

	// Register OpenAPI spec endpoint
	mux.HandleFunc("GET /api/openapi.yaml", ServeOpenAPISpec)
	s.logger.Info("OpenAPI spec endpoint registered", "path", "/api/openapi.yaml")
}

// registerDocsRoutes sets up API documentation serving
func (s *Server) registerDocsRoutes(mux *http.ServeMux) {
	// Check if docs directory exists
	if _, err := os.Stat("docs/api"); err == nil {
		docsFS := http.FileServer(http.Dir("docs/api"))
		mux.Handle("/docs/", http.StripPrefix("/docs/", docsFS))
		s.logger.Info("API documentation endpoint registered", "path", "/docs/")
	} else {
		s.logger.Debug("API documentation not available", "path", "docs/api")
	}
}

// registerOpenAPIRoutes sets up Echo router with generated OpenAPI handlers
func (s *Server) registerOpenAPIRoutes(mux *http.ServeMux, jobsHandler *jobs.Handler) {
	// Create Echo instance for OpenAPI-generated routes
	e := echo.New()

	// Create jupyter handler
	jupyterHandler := jupyter.NewHandler(
		s.logger,
		s.jupyterServiceMgr,
		s.config.BuildAuthConfig(),
	)

	// Create datasets handler
	datasetsHandler := datasets.NewHandler(
		s.logger,
		s.db,
		s.config.DataDir,
	)

	// Create plugins handler
	pluginsHandler := plugins.NewHandler(
		s.logger,
		s.trackingRegistry, // Need to add this to Server
		nil,                // Plugin config - can be loaded from config
	)

	// Create scheduler handler
	var schedulerHandler *sch.APIHandler
	if s.schedulerHub != nil { // Need to add this to Server
		schedulerHandler = sch.NewAPIHandler(s.logger, s.schedulerHub)
	}

	// Create audit handler
	auditHandler := audit.NewHandler(s.logger, nil) // Audit store can be added later

	// Create adapter implementing ServerInterface
	handlerAdapter := NewHandlerAdapter(
		jobsHandler,
		jupyterHandler,
		datasetsHandler,
		pluginsHandler,
		schedulerHandler,
		auditHandler,
	)

	// Register generated OpenAPI routes
	RegisterHandlers(e, handlerAdapter)

	// Add scheduler workers endpoint directly to main mux (not Echo)
	if schedulerHandler != nil {
		mux.HandleFunc("GET /api/v1/workers", schedulerHandler.GetV1SchedulerWorkers)
		mux.HandleFunc("GET /api/v1/workers/{workerId}", schedulerHandler.GetV1SchedulerWorkersWorkerID)
		mux.HandleFunc("DELETE /api/v1/workers/{workerId}", schedulerHandler.DeleteV1SchedulerWorkersWorkerID)
	}

	// Wrap Echo router to work with existing middleware chain
	echoHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		e.ServeHTTP(w, r)
	})

	// Register Echo router at /v1/ prefix
	// These paths take precedence over legacy routes
	mux.Handle("/v1/", echoHandler)

	s.logger.Info("OpenAPI-generated routes registered with Echo router")
}

// registerHealthRoutes sets up health check endpoints
func (s *Server) registerHealthRoutes(mux *http.ServeMux) {
	healthHandler := NewHealthHandler(s)
	healthHandler.RegisterRoutes(mux)
	s.logger.Info("health check endpoints registered")
}

// registerWebSocketRoutes sets up WebSocket endpoint
func (s *Server) registerWebSocketRoutes(mux *http.ServeMux) {
	// Initialize audit logger for WebSocket connections
	auditLogger := s.initAuditLogger()

	// Register WebSocket handler with security config and audit logger
	securityCfg := getSecurityConfig(s.config)

	// Create jobs handler
	jobsHandler := jobs.NewHandler(
		s.expManager,
		s.logger,
		s.taskQueue,
		s.db,
		s.config.BuildAuthConfig(),
		nil, // privacyEnforcer - not enabled for now
	)

	// Create jupyter handler
	jupyterHandler := jupyter.NewHandler(
		s.logger,
		s.jupyterServiceMgr,
		s.config.BuildAuthConfig(),
	)

	// Create datasets handler
	datasetsHandler := datasets.NewHandler(
		s.logger,
		s.db,
		s.config.DataDir,
	)

	wsHandler := ws.NewHandler(
		s.config.BuildAuthConfig(),
		s.logger,
		s.expManager,
		s.config.DataDir,
		s.taskQueue,
		s.db,
		s.jupyterServiceMgr,
		securityCfg,
		auditLogger,
		jobsHandler,
		jupyterHandler,
		datasetsHandler,
	)

	mux.Handle("/ws", wsHandler)
	s.logger.Info("websocket endpoint registered")
}