api.go

package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	log2 "mist/multilogger"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/redis/go-redis/v9"
)

type App struct {
	redisClient    *redis.Client
	scheduler      *Scheduler
	supervisor     *Supervisor
	httpServer     *http.Server
	wg             sync.WaitGroup
	log            *slog.Logger
	statusRegistry *StatusRegistry
	authToken      string // if set, Bearer token required for protected endpoints
}

func NewApp(redisAddr, gpuType string, log *slog.Logger) *App {
	client := redis.NewClient(&redis.Options{Addr: redisAddr})
	scheduler := NewScheduler(redisAddr, log)
	statusRegistry := NewStatusRegistry(client, log)

	consumerID := fmt.Sprintf("worker_%d", os.Getpid())
	supervisor := NewSupervisor(redisAddr, consumerID, gpuType, log)

	authToken := os.Getenv("AUTH_TOKEN")

	mux := http.NewServeMux()
	a := &App{
		redisClient:    client,
		scheduler:      scheduler,
		supervisor:     supervisor,
		httpServer:     &http.Server{Addr: ":3000", Handler: mux},
		log:            log,
		statusRegistry: statusRegistry,
		authToken:      authToken,
	}

	mux.HandleFunc("/auth/login", a.login)
	mux.HandleFunc("/auth/refresh", a.refresh)
	mux.HandleFunc("/jobs", a.handleJobs)
	mux.HandleFunc("/jobs/status", a.getJobStatus)
	mux.HandleFunc("/jobs/logs/", a.requireAuth(a.getJobLogs))
	mux.HandleFunc("/supervisors/status", a.getSupervisorStatus)
	mux.HandleFunc("/supervisors/status/", a.getSupervisorStatusByID)
	mux.HandleFunc("/supervisors", a.getAllSupervisors)

	a.log.Info("new app initialized", "redis_address", redisAddr,
		"gpu_type", gpuType, "http_address", a.httpServer.Addr)

	return a
}

func (a *App) Start() error {
	// Connect to redis
	if err := a.redisClient.Ping(context.Background()).Err(); err != nil {
		a.log.Error("redis ping failed", "err", err)
		return err
	}

	// Start supervisor
	if err := a.supervisor.Start(); err != nil {
		a.log.Error("supervisor start failed", "err", err)
		return err
	}

	// Launch HTTP server
	a.wg.Add(1)
	go func() {
		defer a.wg.Done()
		slog.Info("http server started", "address", a.httpServer.Addr)
		if err := a.httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			a.log.Error("HTTP server error", "err", err)
		}
	}()

	return nil
}

func (a *App) Shutdown(ctx context.Context) error {
	if err := a.httpServer.Shutdown(ctx); err != nil {
		a.log.Error("error shutting down HTTP server", "err", err)
	}

	// Wait for ListenAndServe goroutine to finish
	a.wg.Wait()

	a.supervisor.Stop()

	if err := a.scheduler.Close(); err != nil {
		a.log.Error("error closing scheduler", "err", err)

	} else {
		a.log.Info("scheduler closed successfully")
	}

	if err := a.redisClient.Close(); err != nil {
		a.log.Error("error closing redis client", "err", err)
	} else {
		a.log.Info("redis client closed successfully")
	}

	a.log.Info("shutdown completed")

	return nil
}

func main() {
	cfg, err := log2.GetLogConfig()
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to get log config: %v\n", err)
	}
	log, err := log2.CreateLogger("app", &cfg)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to create logger: %v\n", err)
		os.Exit(1)
	}
	app := NewApp("localhost:6379", "AMD", log)

	if err := app.Start(); err != nil {
		log.Error("failed to start app", "err", err)
		os.Exit(1)
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	<-ctx.Done()
	log.Info("shutdown signal received")

	shutdownCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	if err := app.Shutdown(shutdownCtx); err != nil {
		log.Error("shutdown error", "err", err)
	}

	log.Info("all services stopped cleanly")
}

func (a *App) login(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	a.log.Info("login handler accessed", "remote_address", r.RemoteAddr)
	val, err := a.redisClient.Get(ctx, "some:key").Result()
	if errors.Is(err, redis.Nil) {
		a.log.Info("redis key not found")
		http.Error(w, "redis key not found", http.StatusNotFound)
		return
	}
	if err != nil {
		a.log.Error("redis error on login", "err", err)
		http.Error(w, "redis error", http.StatusInternalServerError)
		return
	}
	a.log.Info("login success", "remote_address", r.RemoteAddr)
	fmt.Fprintf(w, "login page; redis says: %q\n", val)
}

func (a *App) refresh(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Hello, world!\n")
}

type CreateJobRequest struct {
	Type        string                 `json:"type"`
	Payload     map[string]interface{} `json:"payload"`
	RequiredGPU string                 `json:"gpu,omitempty"`
}

type CreateJobResponse struct {
	JobID string `json:"job_id"`
}

func (a *App) handleJobs(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodPost {
		a.createJob(w, r)
		return
	}
	http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
}

func (a *App) createJob(w http.ResponseWriter, r *http.Request) {

	a.log.Info("createJob handler accessed", "remote_address", r.RemoteAddr)

	var req CreateJobRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		a.log.Error("failed to decode request body", "err", err)
		http.Error(w, "Invalid request body", http.StatusBadRequest)
		return
	}

	if req.Type == "" {
		http.Error(w, "Job type is required", http.StatusBadRequest)
		return
	}
	jobID, err := a.scheduler.Enqueue(req.Type, req.RequiredGPU, req.Payload)
	if err != nil {
		a.log.Error("enqueue failed", "err", err, "payload", req.Payload)
		http.Error(w, "enqueue failed", http.StatusInternalServerError)
		return
	}

	a.log.Info("job created", "job_id", jobID, "type", req.Type, "gpu", req.RequiredGPU)

	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusCreated)

	response := CreateJobResponse{JobID: jobID}
	if err := json.NewEncoder(w).Encode(response); err != nil {
		a.log.Error("failed to encode response", "err", err)
	}
}

func (a *App) getJobStatus(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	jobID := r.URL.Query().Get("id")
	if jobID == "" {
		// Try to get from path if query param not provided
		path := strings.TrimPrefix(r.URL.Path, "/jobs/status/")
		if path != "" && path != "/jobs/status" {
			jobID = path
		}
	}

	if jobID == "" {
		http.Error(w, "Job ID is required", http.StatusBadRequest)
		return
	}

	a.log.Info("getJobStatus handler accessed", "job_id", jobID, "remote_address", r.RemoteAddr)

	job, err := a.statusRegistry.GetJobStatus(jobID)
	if err != nil {
		a.log.Error("failed to get job status", "job_id", jobID, "error", err)
		http.Error(w, fmt.Sprintf("Job not found: %s", jobID), http.StatusNotFound)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(job); err != nil {
		a.log.Error("failed to encode job status response", "error", err)
		http.Error(w, "Failed to encode response", http.StatusInternalServerError)
		return
	}
}

func (a *App) getSupervisorStatus(w http.ResponseWriter, r *http.Request) {
	supervisors, err := a.statusRegistry.GetAllSupervisors()
	if err != nil {
		a.log.Error("failed to get supervisor status", "error", err)
		http.Error(w, "failed to get supervisor status", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(map[string]interface{}{
		"supervisors": supervisors,
		"count":       len(supervisors),
	}); err != nil {
		a.log.Error("failed to encode supervisor status response", "error", err)
		http.Error(w, "failed to encode response", http.StatusInternalServerError)
		return
	}
}

func (a *App) getSupervisorStatusByID(w http.ResponseWriter, r *http.Request) {
	// extract consumer ID from URL path
	path := strings.TrimPrefix(r.URL.Path, "/supervisors/status/")
	if path == "" {
		http.Error(w, "consumer ID required", http.StatusBadRequest)
		return
	}

	supervisor, err := a.statusRegistry.GetSupervisor(path)
	if err != nil {
		a.log.Error("failed to get supervisor status", "consumer_id", path, "error", err)
		http.Error(w, "supervisor not found", http.StatusNotFound)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(supervisor); err != nil {
		a.log.Error("failed to encode supervisor status response", "error", err)
		http.Error(w, "failed to encode response", http.StatusInternalServerError)
		return
	}
}

// requireAuth wraps a handler and enforces Bearer token authentication.
// If AUTH_TOKEN is not set, returns 503 (logs feature not configured).
// If Authorization header is missing or invalid, returns 401.
func (a *App) requireAuth(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if a.authToken == "" {
			a.log.Warn("job logs requested but AUTH_TOKEN not configured")
			http.Error(w, "Logs require authentication to be configured", http.StatusServiceUnavailable)
			return
		}
		auth := r.Header.Get("Authorization")
		if auth == "" {
			http.Error(w, "Authorization header required", http.StatusUnauthorized)
			return
		}
		const prefix = "Bearer "
		if !strings.HasPrefix(auth, prefix) || strings.TrimSpace(strings.TrimPrefix(auth, prefix)) != a.authToken {
			http.Error(w, "Invalid or expired token", http.StatusUnauthorized)
			return
		}
		next(w, r)
	}
}

func (a *App) getJobLogs(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}

	path := strings.TrimPrefix(r.URL.Path, "/jobs/logs/")
	jobID := strings.Trim(path, "/")
	if jobID == "" {
		jobID = r.URL.Query().Get("id")
	}
	if jobID == "" {
		http.Error(w, "Job ID is required", http.StatusBadRequest)
		return
	}

	a.log.Info("getJobLogs handler accessed", "job_id", jobID, "remote_address", r.RemoteAddr)

	logs, err := a.supervisor.GetContainerLogsForJob(jobID)
	if err != nil {
		a.log.Error("failed to get job logs", "job_id", jobID, "error", err)
		http.Error(w, fmt.Sprintf("Logs not available for job: %s (container must be running)", jobID), http.StatusNotFound)
		return
	}

	w.Header().Set("Content-Type", "text/plain; charset=utf-8")
	w.WriteHeader(http.StatusOK)
	if _, err := w.Write(logs); err != nil {
		a.log.Error("failed to write job logs response", "job_id", jobID, "error", err)
	}
}

func (a *App) getAllSupervisors(w http.ResponseWriter, r *http.Request) {
	activeOnly := r.URL.Query().Get("active") == "true"

	var supervisors []SupervisorStatus
	var err error

	if activeOnly {
		supervisors, err = a.statusRegistry.GetActiveSupervisors()
	} else {
		supervisors, err = a.statusRegistry.GetAllSupervisors()
	}

	if err != nil {
		a.log.Error("failed to get supervisors", "active_only", activeOnly, "error", err)
		http.Error(w, "failed to get supervisors", http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if err := json.NewEncoder(w).Encode(map[string]interface{}{
		"supervisors": supervisors,
		"count":       len(supervisors),
		"active_only": activeOnly,
	}); err != nil {
		a.log.Error("failed to encode supervisors response", "error", err)
		http.Error(w, "failed to encode response", http.StatusInternalServerError)
		return
	}
}