silo/internal/api/solver_handlers.go

package api

import (
	"context"
	"encoding/json"
	"net/http"
	"strconv"
	"strings"

	"github.com/go-chi/chi/v5"
	"github.com/kindredsystems/silo/internal/auth"
	"github.com/kindredsystems/silo/internal/db"
)

// SubmitSolveRequest is the JSON body for POST /api/solver/jobs.
type SubmitSolveRequest struct {
	Solver         string          `json:"solver"`
	Operation      string          `json:"operation"`
	Context        json.RawMessage `json:"context"`
	Priority       *int            `json:"priority,omitempty"`
	ItemPartNumber string          `json:"item_part_number,omitempty"`
	RevisionNumber *int            `json:"revision_number,omitempty"`
}

// SolverJobResponse is the JSON response for solver job creation.
type SolverJobResponse struct {
	JobID     string `json:"job_id"`
	Status    string `json:"status"`
	CreatedAt string `json:"created_at"`
}

// SolverResultResponse is the JSON response for cached solver results.
type SolverResultResponse struct {
	ID             string          `json:"id"`
	RevisionNumber int             `json:"revision_number"`
	JobID          *string         `json:"job_id,omitempty"`
	Operation      string          `json:"operation"`
	SolverName     string          `json:"solver_name"`
	Status         string          `json:"status"`
	DOF            *int            `json:"dof,omitempty"`
	Diagnostics    json.RawMessage `json:"diagnostics"`
	Placements     json.RawMessage `json:"placements"`
	NumFrames      int             `json:"num_frames"`
	SolveTimeMS    *float64        `json:"solve_time_ms,omitempty"`
	CreatedAt      string          `json:"created_at"`
}

// operationToDefinition maps solve operations to job definition names.
var operationToDefinition = map[string]string{
	"solve":     "assembly-solve",
	"diagnose":  "assembly-validate",
	"kinematic": "assembly-kinematic",
}

// HandleSubmitSolverJob creates a solver job via the existing job queue.
// POST /api/solver/jobs
func (s *Server) HandleSubmitSolverJob(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()

	// Enforce max context size at the HTTP boundary.
	maxBytes := int64(s.cfg.Solver.MaxContextSizeMB) * 1024 * 1024
	r.Body = http.MaxBytesReader(w, r.Body, maxBytes)

	var req SubmitSolveRequest
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		if err.Error() == "http: request body too large" {
			writeError(w, http.StatusRequestEntityTooLarge, "context_too_large",
				"SolveContext exceeds maximum size")
			return
		}
		writeError(w, http.StatusBadRequest, "invalid_body", "Invalid JSON body")
		return
	}

	// Validate operation.
	if req.Operation == "" {
		req.Operation = "solve"
	}
	defName, ok := operationToDefinition[req.Operation]
	if !ok {
		writeError(w, http.StatusBadRequest, "invalid_operation",
			"Operation must be 'solve', 'diagnose', or 'kinematic'")
		return
	}

	// Context is required.
	if len(req.Context) == 0 {
		writeError(w, http.StatusBadRequest, "missing_context", "SolveContext is required")
		return
	}

	// Look up job definition.
	def, err := s.jobs.GetDefinition(ctx, defName)
	if err != nil {
		s.logger.Error().Err(err).Str("definition", defName).Msg("failed to look up solver job definition")
		writeError(w, http.StatusInternalServerError, "internal_error", "Failed to look up job definition")
		return
	}
	if def == nil {
		writeError(w, http.StatusNotFound, "definition_not_found",
			"Solver job definition '"+defName+"' not found; ensure job definition YAML is loaded")
		return
	}

	// Resolve item_part_number → item_id (optional).
	var itemID *string
	if req.ItemPartNumber != "" {
		item, err := s.items.GetByPartNumber(ctx, req.ItemPartNumber)
		if err != nil {
			s.logger.Error().Err(err).Msg("failed to get item for solver job")
			writeError(w, http.StatusInternalServerError, "internal_error", "Failed to resolve item")
			return
		}
		if item == nil {
			writeError(w, http.StatusNotFound, "item_not_found",
				"Item '"+req.ItemPartNumber+"' not found")
			return
		}
		itemID = &item.ID
	}

	// Pack solver-specific data into scope_metadata.
	scopeMeta := map[string]any{
		"solver":    req.Solver,
		"operation": req.Operation,
		"context":   req.Context,
	}
	if req.RevisionNumber != nil {
		scopeMeta["revision_number"] = *req.RevisionNumber
	}
	if req.ItemPartNumber != "" {
		scopeMeta["item_part_number"] = req.ItemPartNumber
	}

	priority := def.Priority
	if req.Priority != nil {
		priority = *req.Priority
	}

	username := ""
	if user := auth.UserFromContext(ctx); user != nil {
		username = user.Username
	}

	job := &db.Job{
		JobDefinitionID: &def.ID,
		DefinitionName:  def.Name,
		Priority:        priority,
		ItemID:          itemID,
		ScopeMetadata:   scopeMeta,
		RunnerTags:      def.RunnerTags,
		TimeoutSeconds:  def.TimeoutSeconds,
		MaxRetries:      def.MaxRetries,
		CreatedBy:       &username,
	}

	// Use solver default timeout if the definition has none.
	if job.TimeoutSeconds == 0 {
		job.TimeoutSeconds = s.cfg.Solver.DefaultTimeout
	}

	if err := s.jobs.CreateJob(ctx, job); err != nil {
		s.logger.Error().Err(err).Msg("failed to create solver job")
		writeError(w, http.StatusInternalServerError, "internal_error", "Failed to create solver job")
		return
	}

	s.broker.Publish("job.created", mustMarshal(map[string]any{
		"job_id":          job.ID,
		"definition_name": job.DefinitionName,
		"trigger":         "manual",
		"item_id":         job.ItemID,
	}))

	writeJSON(w, http.StatusCreated, SolverJobResponse{
		JobID:     job.ID,
		Status:    job.Status,
		CreatedAt: job.CreatedAt.UTC().Format("2006-01-02T15:04:05Z"),
	})
}

// HandleGetSolverJob returns a single solver job.
// GET /api/solver/jobs/{jobID}
func (s *Server) HandleGetSolverJob(w http.ResponseWriter, r *http.Request) {
	jobID := chi.URLParam(r, "jobID")

	job, err := s.jobs.GetJob(r.Context(), jobID)
	if err != nil {
		s.logger.Error().Err(err).Msg("failed to get solver job")
		writeError(w, http.StatusInternalServerError, "internal_error", "Failed to get job")
		return
	}
	if job == nil {
		writeError(w, http.StatusNotFound, "not_found", "Job not found")
		return
	}

	writeJSON(w, http.StatusOK, job)
}

// HandleListSolverJobs lists solver jobs with optional filters.
// GET /api/solver/jobs
func (s *Server) HandleListSolverJobs(w http.ResponseWriter, r *http.Request) {
	status := r.URL.Query().Get("status")
	itemPartNumber := r.URL.Query().Get("item")
	operation := r.URL.Query().Get("operation")

	limit := 20
	if v := r.URL.Query().Get("limit"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 100 {
			limit = n
		}
	}
	offset := 0
	if v := r.URL.Query().Get("offset"); v != "" {
		if n, err := strconv.Atoi(v); err == nil && n >= 0 {
			offset = n
		}
	}

	// Resolve item part number to ID if provided.
	var itemID string
	if itemPartNumber != "" {
		item, err := s.items.GetByPartNumber(r.Context(), itemPartNumber)
		if err != nil {
			s.logger.Error().Err(err).Msg("failed to resolve item for solver job list")
			writeError(w, http.StatusInternalServerError, "internal_error", "Failed to resolve item")
			return
		}
		if item == nil {
			writeJSON(w, http.StatusOK, map[string]any{
				"jobs":   []*db.Job{},
				"total":  0,
				"limit":  limit,
				"offset": offset,
			})
			return
		}
		itemID = item.ID
	}

	jobs, err := s.jobs.ListSolverJobs(r.Context(), status, itemID, operation, limit, offset)
	if err != nil {
		s.logger.Error().Err(err).Msg("failed to list solver jobs")
		writeError(w, http.StatusInternalServerError, "internal_error", "Failed to list solver jobs")
		return
	}

	writeJSON(w, http.StatusOK, map[string]any{
		"jobs":   jobs,
		"limit":  limit,
		"offset": offset,
	})
}

// HandleCancelSolverJob cancels a solver job.
// POST /api/solver/jobs/{jobID}/cancel
func (s *Server) HandleCancelSolverJob(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	jobID := chi.URLParam(r, "jobID")
	user := auth.UserFromContext(ctx)

	cancelledBy := "system"
	if user != nil {
		cancelledBy = user.Username
	}

	if err := s.jobs.CancelJob(ctx, jobID, cancelledBy); err != nil {
		writeError(w, http.StatusBadRequest, "cancel_failed", err.Error())
		return
	}

	s.broker.Publish("job.cancelled", mustMarshal(map[string]any{
		"job_id":       jobID,
		"cancelled_by": cancelledBy,
	}))

	writeJSON(w, http.StatusOK, map[string]string{
		"job_id": jobID,
		"status": "cancelled",
	})
}

// HandleGetSolverRegistry returns available solvers from online runners.
// GET /api/solver/solvers
func (s *Server) HandleGetSolverRegistry(w http.ResponseWriter, r *http.Request) {
	runners, err := s.jobs.ListRunners(r.Context())
	if err != nil {
		s.logger.Error().Err(err).Msg("failed to list runners for solver registry")
		writeError(w, http.StatusInternalServerError, "internal_error", "Failed to list runners")
		return
	}

	type solverInfo struct {
		Name            string   `json:"name"`
		DisplayName     string   `json:"display_name,omitempty"`
		Deterministic   bool     `json:"deterministic,omitempty"`
		SupportedJoints []string `json:"supported_joints,omitempty"`
		RunnerCount     int      `json:"runner_count"`
	}

	solverMap := make(map[string]*solverInfo)

	for _, runner := range runners {
		if runner.Status != "online" {
			continue
		}
		// Check runner has the solver tag.
		hasSolverTag := false
		for _, tag := range runner.Tags {
			if tag == "solver" {
				hasSolverTag = true
				break
			}
		}
		if !hasSolverTag {
			continue
		}

		// Extract solver capabilities from runner metadata.
		if runner.Metadata == nil {
			continue
		}
		solvers, ok := runner.Metadata["solvers"]
		if !ok {
			continue
		}

		// solvers can be []any (array of solver objects or strings).
		solverList, ok := solvers.([]any)
		if !ok {
			continue
		}

		for _, entry := range solverList {
			switch v := entry.(type) {
			case string:
				// Simple string entry: just the solver name.
				if _, exists := solverMap[v]; !exists {
					solverMap[v] = &solverInfo{Name: v}
				}
				solverMap[v].RunnerCount++
			case map[string]any:
				// Rich entry with name, display_name, supported_joints, etc.
				name, _ := v["name"].(string)
				if name == "" {
					continue
				}
				if _, exists := solverMap[name]; !exists {
					info := &solverInfo{Name: name}
					if dn, ok := v["display_name"].(string); ok {
						info.DisplayName = dn
					}
					if det, ok := v["deterministic"].(bool); ok {
						info.Deterministic = det
					}
					if joints, ok := v["supported_joints"].([]any); ok {
						for _, j := range joints {
							if js, ok := j.(string); ok {
								info.SupportedJoints = append(info.SupportedJoints, js)
							}
						}
					}
					solverMap[name] = info
				}
				solverMap[name].RunnerCount++
			}
		}
	}

	solverList := make([]*solverInfo, 0, len(solverMap))
	for _, info := range solverMap {
		solverList = append(solverList, info)
	}

	writeJSON(w, http.StatusOK, map[string]any{
		"solvers":        solverList,
		"default_solver": s.cfg.Solver.DefaultSolver,
	})
}

// HandleGetSolverResults returns cached solver results for an item.
// GET /api/items/{partNumber}/solver/results
func (s *Server) HandleGetSolverResults(w http.ResponseWriter, r *http.Request) {
	ctx := r.Context()
	partNumber := chi.URLParam(r, "partNumber")

	item, err := s.items.GetByPartNumber(ctx, partNumber)
	if err != nil {
		s.logger.Error().Err(err).Msg("failed to get item for solver results")
		writeError(w, http.StatusInternalServerError, "internal_error", "Failed to get item")
		return
	}
	if item == nil {
		writeError(w, http.StatusNotFound, "not_found", "Item not found")
		return
	}

	results, err := s.solverResults.GetByItem(ctx, item.ID)
	if err != nil {
		s.logger.Error().Err(err).Msg("failed to list solver results")
		writeError(w, http.StatusInternalServerError, "internal_error", "Failed to list solver results")
		return
	}

	resp := make([]SolverResultResponse, len(results))
	for i, r := range results {
		diag := json.RawMessage(r.Diagnostics)
		if diag == nil {
			diag = json.RawMessage("[]")
		}
		place := json.RawMessage(r.Placements)
		if place == nil {
			place = json.RawMessage("[]")
		}
		resp[i] = SolverResultResponse{
			ID:             r.ID,
			RevisionNumber: r.RevisionNumber,
			JobID:          r.JobID,
			Operation:      r.Operation,
			SolverName:     r.SolverName,
			Status:         r.Status,
			DOF:            r.DOF,
			Diagnostics:    diag,
			Placements:     place,
			NumFrames:      r.NumFrames,
			SolveTimeMS:    r.SolveTimeMS,
			CreatedAt:      r.CreatedAt.UTC().Format("2006-01-02T15:04:05Z"),
		}
	}

	writeJSON(w, http.StatusOK, resp)
}

// maybeCacheSolverResult is called asynchronously after a job completes.
// It checks if the job is a solver job and upserts the result into solver_results.
func (s *Server) maybeCacheSolverResult(ctx context.Context, jobID string) {
	job, err := s.jobs.GetJob(ctx, jobID)
	if err != nil || job == nil {
		s.logger.Warn().Err(err).Str("job_id", jobID).Msg("solver result cache: failed to get job")
		return
	}

	if !strings.HasPrefix(job.DefinitionName, "assembly-") {
		return
	}
	if !s.modules.IsEnabled("solver") {
		return
	}
	if job.ItemID == nil {
		return
	}

	// Extract fields from scope_metadata.
	operation, _ := job.ScopeMetadata["operation"].(string)
	if operation == "" {
		operation = "solve"
	}
	solverName, _ := job.ScopeMetadata["solver"].(string)

	var revisionNumber int
	if rn, ok := job.ScopeMetadata["revision_number"].(float64); ok {
		revisionNumber = int(rn)
	}

	// Extract fields from result.
	if job.Result == nil {
		return
	}

	status, _ := job.Result["status"].(string)
	if status == "" {
		// Try nested result object.
		if inner, ok := job.Result["result"].(map[string]any); ok {
			status, _ = inner["status"].(string)
		}
	}
	if status == "" {
		status = "Unknown"
	}

	// Solver name from result takes precedence.
	if sn, ok := job.Result["solver_name"].(string); ok && sn != "" {
		solverName = sn
	}
	if solverName == "" {
		solverName = "unknown"
	}

	var dof *int
	if d, ok := job.Result["dof"].(float64); ok {
		v := int(d)
		dof = &v
	} else if inner, ok := job.Result["result"].(map[string]any); ok {
		if d, ok := inner["dof"].(float64); ok {
			v := int(d)
			dof = &v
		}
	}

	var solveTimeMS *float64
	if t, ok := job.Result["solve_time_ms"].(float64); ok {
		solveTimeMS = &t
	}

	// Marshal diagnostics and placements as raw JSONB.
	var diagnostics, placements []byte
	if d, ok := job.Result["diagnostics"]; ok {
		diagnostics, _ = json.Marshal(d)
	} else if inner, ok := job.Result["result"].(map[string]any); ok {
		if d, ok := inner["diagnostics"]; ok {
			diagnostics, _ = json.Marshal(d)
		}
	}
	if p, ok := job.Result["placements"]; ok {
		placements, _ = json.Marshal(p)
	} else if inner, ok := job.Result["result"].(map[string]any); ok {
		if p, ok := inner["placements"]; ok {
			placements, _ = json.Marshal(p)
		}
	}

	numFrames := 0
	if nf, ok := job.Result["num_frames"].(float64); ok {
		numFrames = int(nf)
	} else if inner, ok := job.Result["result"].(map[string]any); ok {
		if nf, ok := inner["num_frames"].(float64); ok {
			numFrames = int(nf)
		}
	}

	result := &db.SolverResult{
		ItemID:         *job.ItemID,
		RevisionNumber: revisionNumber,
		JobID:          &job.ID,
		Operation:      operation,
		SolverName:     solverName,
		Status:         status,
		DOF:            dof,
		Diagnostics:    diagnostics,
		Placements:     placements,
		NumFrames:      numFrames,
		SolveTimeMS:    solveTimeMS,
	}

	if err := s.solverResults.Upsert(ctx, result); err != nil {
		s.logger.Warn().Err(err).Str("job_id", jobID).Msg("solver result cache: failed to upsert")
	} else {
		s.logger.Info().Str("job_id", jobID).Str("operation", operation).Msg("cached solver result")
	}
}