breakpilot-compliance/ai-compliance-sdk/internal/api/handlers/usecase_handler.go

package handlers

import (
	"net/http"

	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/breakpilot/ai-compliance-sdk/internal/llm"
	"github.com/breakpilot/ai-compliance-sdk/internal/usecase"
)

// UseCaseHandler handles use-case compiler endpoints.
type UseCaseHandler struct {
	store       *usecase.Store
	compiler    *usecase.Compiler
	gapDetector *usecase.GapDetector
	llmGen      *usecase.LLMQuestionGenerator
}

// NewUseCaseHandler creates a new UseCaseHandler.
func NewUseCaseHandler(pool *pgxpool.Pool, registry *llm.ProviderRegistry) *UseCaseHandler {
	store := usecase.NewStore(pool)
	return &UseCaseHandler{
		store:       store,
		compiler:    usecase.NewCompiler(store),
		gapDetector: usecase.NewGapDetector(store),
		llmGen:      usecase.NewLLMQuestionGenerator(registry),
	}
}

// GetTemplates returns all available use-case templates.
// GET /sdk/v1/use-case/templates
func (h *UseCaseHandler) GetTemplates(c *gin.Context) {
	templates := usecase.TemplateList()
	c.JSON(http.StatusOK, gin.H{"templates": templates, "total": len(templates)})
}

// GetTemplate returns a specific template with compiled questions.
// GET /sdk/v1/use-case/templates/:id
func (h *UseCaseHandler) GetTemplate(c *gin.Context) {
	id := c.Param("id")
	tmpl, ok := usecase.Templates[id]
	if !ok {
		c.JSON(http.StatusNotFound, gin.H{"error": "template not found"})
		return
	}

	questions, err := h.compiler.Compile(&tmpl)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	tmpl.Questions = questions

	c.JSON(http.StatusOK, gin.H{"template": tmpl})
}

// Compile generates questions from MC filters ad-hoc.
// POST /sdk/v1/use-case/compile
// Optional: "mode": "llm" to use LLM-based generation
func (h *UseCaseHandler) Compile(c *gin.Context) {
	var req struct {
		MCFilters   []string `json:"mc_filters" binding:"required"`
		Regulations []string `json:"regulations"`
		Mode        string   `json:"mode"` // "deterministic" (default) or "llm"
	}
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	tmpl := &usecase.Template{
		ID:          "custom",
		MCFilters:   req.MCFilters,
		Regulations: req.Regulations,
	}

	if req.Mode == "llm" && h.llmGen != nil {
		// Fetch MCs first, then generate via LLM
		mcs, err := h.store.FetchMCsByFilters(req.MCFilters)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		questions, err := h.llmGen.GenerateQuestions(mcs, req.Regulations)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		c.JSON(http.StatusOK, gin.H{"questions": questions, "total": len(questions), "mode": "llm"})
		return
	}

	questions, err := h.compiler.Compile(tmpl)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	c.JSON(http.StatusOK, gin.H{"questions": questions, "total": len(questions), "mode": "deterministic"})
}

// CreateAudit starts a new audit from a template.
// POST /sdk/v1/use-case/audits
func (h *UseCaseHandler) CreateAudit(c *gin.Context) {
	var input usecase.CreateAuditInput
	if err := c.ShouldBindJSON(&input); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	tenantID, err := uuid.Parse(c.GetHeader("X-Tenant-ID"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "X-Tenant-ID required"})
		return
	}

	tmpl, ok := usecase.Templates[input.TemplateID]
	if !ok {
		c.JSON(http.StatusBadRequest, gin.H{"error": "unknown template_id"})
		return
	}

	questions, err := h.compiler.Compile(&tmpl)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	audit := &usecase.Audit{
		TenantID:       tenantID,
		TemplateID:     input.TemplateID,
		Name:           input.Name,
		TargetName:     input.TargetName,
		TotalQuestions: len(questions),
		Questions:      questions,
	}

	if err := h.store.CreateAudit(audit); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create audit"})
		return
	}

	c.JSON(http.StatusCreated, gin.H{"audit": audit})
}

// ListAudits returns all audits for a tenant.
// GET /sdk/v1/use-case/audits
func (h *UseCaseHandler) ListAudits(c *gin.Context) {
	tenantID, err := uuid.Parse(c.GetHeader("X-Tenant-ID"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "X-Tenant-ID required"})
		return
	}

	audits, err := h.store.ListAudits(tenantID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to list audits"})
		return
	}

	c.JSON(http.StatusOK, gin.H{"audits": audits, "total": len(audits)})
}

// GetAudit returns an audit with questions and answers.
// GET /sdk/v1/use-case/audits/:id
func (h *UseCaseHandler) GetAudit(c *gin.Context) {
	id, err := uuid.Parse(c.Param("id"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid audit ID"})
		return
	}

	audit, err := h.store.GetAudit(id)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "audit not found"})
		return
	}

	answers, err := h.store.ListAnswers(id)
	if err != nil {
		answers = nil
	}

	c.JSON(http.StatusOK, gin.H{"audit": audit, "answers": answers})
}

// AnswerQuestion saves an answer for a question in an audit.
// POST /sdk/v1/use-case/audits/:id/answer
func (h *UseCaseHandler) AnswerQuestion(c *gin.Context) {
	auditID, err := uuid.Parse(c.Param("id"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid audit ID"})
		return
	}

	var input usecase.AnswerInput
	if err := c.ShouldBindJSON(&input); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	// Find MC ID from the question
	audit, err := h.store.GetAudit(auditID)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "audit not found"})
		return
	}

	var mcID string
	for _, q := range audit.Questions {
		if q.ID == input.QuestionID {
			mcID = q.MCID
			break
		}
	}

	status := usecase.AnswerStatusAnswered
	if input.Status == "skipped" {
		status = usecase.AnswerStatusSkipped
	} else if input.Status == "escalated" {
		status = usecase.AnswerStatusEscalated
	}

	answer := &usecase.Answer{
		AuditID:     auditID,
		QuestionID:  input.QuestionID,
		MCID:        mcID,
		Value:       input.Value,
		Comment:     input.Comment,
		EvidenceIDs: input.EvidenceIDs,
		Status:      status,
	}

	if err := h.store.SaveAnswer(answer); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to save answer"})
		return
	}

	// Update audit counters
	answers, _ := h.store.ListAnswers(auditID)
	score := usecase.Score(audit, answers)

	auditStatus := usecase.StatusInProgress
	if score.Answered >= audit.TotalQuestions {
		auditStatus = usecase.StatusCompleted
	}
	h.store.UpdateAuditScore(auditID, score.Answered, score.ComplianceScore, auditStatus)

	c.JSON(http.StatusOK, gin.H{"answer": answer, "progress": score})
}

// GetScore calculates and returns the compliance score.
// GET /sdk/v1/use-case/audits/:id/score
func (h *UseCaseHandler) GetScore(c *gin.Context) {
	auditID, err := uuid.Parse(c.Param("id"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid audit ID"})
		return
	}

	audit, err := h.store.GetAudit(auditID)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "audit not found"})
		return
	}

	answers, err := h.store.ListAnswers(auditID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to load answers"})
		return
	}

	score := usecase.Score(audit, answers)
	c.JSON(http.StatusOK, score)
}

// GetGaps returns missing regulation sources for an audit.
// GET /sdk/v1/use-case/audits/:id/gaps
func (h *UseCaseHandler) GetGaps(c *gin.Context) {
	auditID, err := uuid.Parse(c.Param("id"))
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid audit ID"})
		return
	}

	audit, err := h.store.GetAudit(auditID)
	if err != nil {
		c.JSON(http.StatusNotFound, gin.H{"error": "audit not found"})
		return
	}

	tmpl, ok := usecase.Templates[audit.TemplateID]
	if !ok {
		c.JSON(http.StatusOK, gin.H{"gaps": []interface{}{}, "audit_gaps": []interface{}{}})
		return
	}

	// Missing regulation sources (from MC analysis)
	missingRegs, err := h.gapDetector.DetectMissingRegulations(&tmpl)
	if err != nil {
		missingRegs = nil
	}

	// Audit-specific gaps (from answer analysis)
	answers, _ := h.store.ListAnswers(auditID)
	auditGaps := h.gapDetector.DetectAuditGaps(audit, answers)

	c.JSON(http.StatusOK, gin.H{
		"missing_sources": missingRegs,
		"audit_gaps":      auditGaps,
	})
}