breakpilot-compliance/ai-compliance-sdk/internal/ucca/store.go

package ucca

import (
	"context"
	"encoding/json"
	"strconv"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// Store handles UCCA data persistence
type Store struct {
	pool *pgxpool.Pool
}

// NewStore creates a new UCCA store
func NewStore(pool *pgxpool.Pool) *Store {
	return &Store{pool: pool}
}

// ============================================================================
// Assessment CRUD Operations
// ============================================================================

// CreateAssessment creates a new assessment
func (s *Store) CreateAssessment(ctx context.Context, a *Assessment) error {
	a.ID = uuid.New()
	a.CreatedAt = time.Now().UTC()
	a.UpdatedAt = a.CreatedAt
	if a.PolicyVersion == "" {
		a.PolicyVersion = "1.0.0"
	}
	if a.Status == "" {
		a.Status = "completed"
	}

	// Marshal JSONB fields
	intake, _ := json.Marshal(a.Intake)
	triggeredRules, _ := json.Marshal(a.TriggeredRules)
	requiredControls, _ := json.Marshal(a.RequiredControls)
	recommendedArchitecture, _ := json.Marshal(a.RecommendedArchitecture)
	forbiddenPatterns, _ := json.Marshal(a.ForbiddenPatterns)
	exampleMatches, _ := json.Marshal(a.ExampleMatches)

	_, err := s.pool.Exec(ctx, `
		INSERT INTO ucca_assessments (
			id, tenant_id, namespace_id, title, policy_version, status,
			intake, use_case_text_stored, use_case_text_hash,
			feasibility, risk_level, complexity, risk_score,
			triggered_rules, required_controls, recommended_architecture,
			forbidden_patterns, example_matches,
			dsfa_recommended, art22_risk, training_allowed,
			corpus_version_id, corpus_version,
			explanation_text, explanation_generated_at, explanation_model,
			domain, created_at, updated_at, created_by
		) VALUES (
			$1, $2, $3, $4, $5, $6,
			$7, $8, $9,
			$10, $11, $12, $13,
			$14, $15, $16,
			$17, $18,
			$19, $20, $21,
			$22, $23,
			$24, $25, $26,
			$27, $28, $29, $30
		)
	`,
		a.ID, a.TenantID, a.NamespaceID, a.Title, a.PolicyVersion, a.Status,
		intake, a.UseCaseTextStored, a.UseCaseTextHash,
		string(a.Feasibility), string(a.RiskLevel), string(a.Complexity), a.RiskScore,
		triggeredRules, requiredControls, recommendedArchitecture,
		forbiddenPatterns, exampleMatches,
		a.DSFARecommended, a.Art22Risk, string(a.TrainingAllowed),
		a.CorpusVersionID, a.CorpusVersion,
		a.ExplanationText, a.ExplanationGeneratedAt, a.ExplanationModel,
		string(a.Domain), a.CreatedAt, a.UpdatedAt, a.CreatedBy,
	)

	return err
}

// GetAssessment retrieves an assessment by ID
func (s *Store) GetAssessment(ctx context.Context, id uuid.UUID) (*Assessment, error) {
	var a Assessment
	var intake, triggeredRules, requiredControls, recommendedArchitecture, forbiddenPatterns, exampleMatches []byte
	var feasibility, riskLevel, complexity, trainingAllowed, domain string

	err := s.pool.QueryRow(ctx, `
		SELECT
			id, tenant_id, namespace_id, title, policy_version, status,
			intake, use_case_text_stored, use_case_text_hash,
			feasibility, risk_level, complexity, risk_score,
			triggered_rules, required_controls, recommended_architecture,
			forbidden_patterns, example_matches,
			dsfa_recommended, art22_risk, training_allowed,
			corpus_version_id, corpus_version,
			explanation_text, explanation_generated_at, explanation_model,
			domain, created_at, updated_at, created_by
		FROM ucca_assessments WHERE id = $1
	`, id).Scan(
		&a.ID, &a.TenantID, &a.NamespaceID, &a.Title, &a.PolicyVersion, &a.Status,
		&intake, &a.UseCaseTextStored, &a.UseCaseTextHash,
		&feasibility, &riskLevel, &complexity, &a.RiskScore,
		&triggeredRules, &requiredControls, &recommendedArchitecture,
		&forbiddenPatterns, &exampleMatches,
		&a.DSFARecommended, &a.Art22Risk, &trainingAllowed,
		&a.CorpusVersionID, &a.CorpusVersion,
		&a.ExplanationText, &a.ExplanationGeneratedAt, &a.ExplanationModel,
		&domain, &a.CreatedAt, &a.UpdatedAt, &a.CreatedBy,
	)

	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	// Unmarshal JSONB fields
	json.Unmarshal(intake, &a.Intake)
	json.Unmarshal(triggeredRules, &a.TriggeredRules)
	json.Unmarshal(requiredControls, &a.RequiredControls)
	json.Unmarshal(recommendedArchitecture, &a.RecommendedArchitecture)
	json.Unmarshal(forbiddenPatterns, &a.ForbiddenPatterns)
	json.Unmarshal(exampleMatches, &a.ExampleMatches)

	// Convert string fields to typed constants
	a.Feasibility = Feasibility(feasibility)
	a.RiskLevel = RiskLevel(riskLevel)
	a.Complexity = Complexity(complexity)
	a.TrainingAllowed = TrainingAllowed(trainingAllowed)
	a.Domain = Domain(domain)

	return &a, nil
}

// ListAssessments lists assessments for a tenant with optional filters
func (s *Store) ListAssessments(ctx context.Context, tenantID uuid.UUID, filters *AssessmentFilters) ([]Assessment, int, error) {
	baseWhere := " WHERE tenant_id = $1"
	args := []interface{}{tenantID}
	argIdx := 2

	// Build WHERE clause from filters
	if filters != nil {
		if filters.Feasibility != "" {
			baseWhere += " AND feasibility = $" + itoa(argIdx)
			args = append(args, filters.Feasibility)
			argIdx++
		}
		if filters.Domain != "" {
			baseWhere += " AND domain = $" + itoa(argIdx)
			args = append(args, filters.Domain)
			argIdx++
		}
		if filters.RiskLevel != "" {
			baseWhere += " AND risk_level = $" + itoa(argIdx)
			args = append(args, filters.RiskLevel)
			argIdx++
		}
		if filters.Search != "" {
			baseWhere += " AND title ILIKE $" + strconv.Itoa(argIdx)
			args = append(args, "%"+filters.Search+"%")
			argIdx++
		}
	}

	// COUNT query (same WHERE, no ORDER/LIMIT/OFFSET)
	var total int
	s.pool.QueryRow(ctx, "SELECT COUNT(*) FROM ucca_assessments"+baseWhere, args...).Scan(&total)

	// Data query
	query := `
		SELECT
			id, tenant_id, namespace_id, title, policy_version, status,
			intake, use_case_text_stored, use_case_text_hash,
			feasibility, risk_level, complexity, risk_score,
			triggered_rules, required_controls, recommended_architecture,
			forbidden_patterns, example_matches,
			dsfa_recommended, art22_risk, training_allowed,
			corpus_version_id, corpus_version,
			explanation_text, explanation_generated_at, explanation_model,
			domain, created_at, updated_at, created_by
		FROM ucca_assessments` + baseWhere + " ORDER BY created_at DESC"

	// Apply LIMIT
	if filters != nil && filters.Limit > 0 {
		query += " LIMIT $" + strconv.Itoa(argIdx)
		args = append(args, filters.Limit)
		argIdx++
	}

	// Apply OFFSET
	if filters != nil && filters.Offset > 0 {
		query += " OFFSET $" + strconv.Itoa(argIdx)
		args = append(args, filters.Offset)
		argIdx++
	}

	rows, err := s.pool.Query(ctx, query, args...)
	if err != nil {
		return nil, 0, err
	}
	defer rows.Close()

	var assessments []Assessment
	for rows.Next() {
		var a Assessment
		var intake, triggeredRules, requiredControls, recommendedArchitecture, forbiddenPatterns, exampleMatches []byte
		var feasibility, riskLevel, complexity, trainingAllowed, domain string

		err := rows.Scan(
			&a.ID, &a.TenantID, &a.NamespaceID, &a.Title, &a.PolicyVersion, &a.Status,
			&intake, &a.UseCaseTextStored, &a.UseCaseTextHash,
			&feasibility, &riskLevel, &complexity, &a.RiskScore,
			&triggeredRules, &requiredControls, &recommendedArchitecture,
			&forbiddenPatterns, &exampleMatches,
			&a.DSFARecommended, &a.Art22Risk, &trainingAllowed,
			&a.CorpusVersionID, &a.CorpusVersion,
			&a.ExplanationText, &a.ExplanationGeneratedAt, &a.ExplanationModel,
			&domain, &a.CreatedAt, &a.UpdatedAt, &a.CreatedBy,
		)
		if err != nil {
			return nil, 0, err
		}

		// Unmarshal JSONB fields
		json.Unmarshal(intake, &a.Intake)
		json.Unmarshal(triggeredRules, &a.TriggeredRules)
		json.Unmarshal(requiredControls, &a.RequiredControls)
		json.Unmarshal(recommendedArchitecture, &a.RecommendedArchitecture)
		json.Unmarshal(forbiddenPatterns, &a.ForbiddenPatterns)
		json.Unmarshal(exampleMatches, &a.ExampleMatches)

		// Convert string fields
		a.Feasibility = Feasibility(feasibility)
		a.RiskLevel = RiskLevel(riskLevel)
		a.Complexity = Complexity(complexity)
		a.TrainingAllowed = TrainingAllowed(trainingAllowed)
		a.Domain = Domain(domain)

		assessments = append(assessments, a)
	}

	return assessments, total, nil
}

// DeleteAssessment deletes an assessment by ID
func (s *Store) DeleteAssessment(ctx context.Context, id uuid.UUID) error {
	_, err := s.pool.Exec(ctx, "DELETE FROM ucca_assessments WHERE id = $1", id)
	return err
}

// UpdateAssessment updates an existing assessment's intake and re-evaluated results
func (s *Store) UpdateAssessment(ctx context.Context, id uuid.UUID, a *Assessment) error {
	intake, _ := json.Marshal(a.Intake)
	triggeredRules, _ := json.Marshal(a.TriggeredRules)
	requiredControls, _ := json.Marshal(a.RequiredControls)
	recommendedArchitecture, _ := json.Marshal(a.RecommendedArchitecture)
	forbiddenPatterns, _ := json.Marshal(a.ForbiddenPatterns)
	exampleMatches, _ := json.Marshal(a.ExampleMatches)
	now := time.Now().UTC()

	_, err := s.pool.Exec(ctx, `
		UPDATE ucca_assessments SET
			title = $2, intake = $3, use_case_text_stored = $4, use_case_text_hash = $5,
			feasibility = $6, risk_level = $7, complexity = $8, risk_score = $9,
			triggered_rules = $10, required_controls = $11, recommended_architecture = $12,
			forbidden_patterns = $13, example_matches = $14,
			dsfa_recommended = $15, art22_risk = $16, training_allowed = $17,
			domain = $18, policy_version = $19, updated_at = $20
		WHERE id = $1
	`, id, a.Title, intake, a.UseCaseTextStored, a.UseCaseTextHash,
		string(a.Feasibility), string(a.RiskLevel), string(a.Complexity), a.RiskScore,
		triggeredRules, requiredControls, recommendedArchitecture,
		forbiddenPatterns, exampleMatches,
		a.DSFARecommended, a.Art22Risk, string(a.TrainingAllowed),
		string(a.Domain), a.PolicyVersion, now)
	return err
}

// UpdateExplanation updates the LLM explanation for an assessment
func (s *Store) UpdateExplanation(ctx context.Context, id uuid.UUID, explanation string, model string) error {
	now := time.Now().UTC()
	_, err := s.pool.Exec(ctx, `
		UPDATE ucca_assessments SET
			explanation_text = $2,
			explanation_generated_at = $3,
			explanation_model = $4,
			updated_at = $5
		WHERE id = $1
	`, id, explanation, now, model, now)
	return err
}

// ============================================================================
// Statistics
// ============================================================================

// UCCAStats contains UCCA module statistics
type UCCAStats struct {
	TotalAssessments       int `json:"total_assessments"`
	AssessmentsYES         int `json:"assessments_yes"`
	AssessmentsCONDITIONAL int `json:"assessments_conditional"`
	AssessmentsNO          int `json:"assessments_no"`
	AverageRiskScore       int `json:"average_risk_score"`
	DSFARecommendedCount   int `json:"dsfa_recommended_count"`
}

// GetStats returns UCCA statistics for a tenant
func (s *Store) GetStats(ctx context.Context, tenantID uuid.UUID) (*UCCAStats, error) {
	stats := &UCCAStats{}

	// Total count
	s.pool.QueryRow(ctx,
		"SELECT COUNT(*) FROM ucca_assessments WHERE tenant_id = $1",
		tenantID).Scan(&stats.TotalAssessments)

	// Count by feasibility
	s.pool.QueryRow(ctx,
		"SELECT COUNT(*) FROM ucca_assessments WHERE tenant_id = $1 AND feasibility = 'YES'",
		tenantID).Scan(&stats.AssessmentsYES)

	s.pool.QueryRow(ctx,
		"SELECT COUNT(*) FROM ucca_assessments WHERE tenant_id = $1 AND feasibility = 'CONDITIONAL'",
		tenantID).Scan(&stats.AssessmentsCONDITIONAL)

	s.pool.QueryRow(ctx,
		"SELECT COUNT(*) FROM ucca_assessments WHERE tenant_id = $1 AND feasibility = 'NO'",
		tenantID).Scan(&stats.AssessmentsNO)

	// Average risk score
	s.pool.QueryRow(ctx,
		"SELECT COALESCE(AVG(risk_score)::int, 0) FROM ucca_assessments WHERE tenant_id = $1",
		tenantID).Scan(&stats.AverageRiskScore)

	// DSFA recommended count
	s.pool.QueryRow(ctx,
		"SELECT COUNT(*) FROM ucca_assessments WHERE tenant_id = $1 AND dsfa_recommended = true",
		tenantID).Scan(&stats.DSFARecommendedCount)

	return stats, nil
}

// ============================================================================
// Filter Types
// ============================================================================

// AssessmentFilters defines filters for listing assessments
type AssessmentFilters struct {
	Feasibility string
	Domain      string
	RiskLevel   string
	Search      string // ILIKE on title
	Limit       int
	Offset      int // OFFSET for pagination
}

// ============================================================================
// Decision Tree Result CRUD
// ============================================================================

// CreateDecisionTreeResult stores a new decision tree result
func (s *Store) CreateDecisionTreeResult(ctx context.Context, r *DecisionTreeResult) error {
	r.ID = uuid.New()
	r.CreatedAt = time.Now().UTC()
	r.UpdatedAt = r.CreatedAt

	answers, _ := json.Marshal(r.Answers)
	gpaiResult, _ := json.Marshal(r.GPAIResult)
	obligations, _ := json.Marshal(r.CombinedObligations)
	articles, _ := json.Marshal(r.ApplicableArticles)

	_, err := s.pool.Exec(ctx, `
		INSERT INTO ai_act_decision_tree_results (
			id, tenant_id, project_id, system_name, system_description,
			answers, high_risk_level, gpai_result,
			combined_obligations, applicable_articles,
			created_at, updated_at
		) VALUES (
			$1, $2, $3, $4, $5,
			$6, $7, $8,
			$9, $10,
			$11, $12
		)
	`,
		r.ID, r.TenantID, r.ProjectID, r.SystemName, r.SystemDescription,
		answers, string(r.HighRiskResult), gpaiResult,
		obligations, articles,
		r.CreatedAt, r.UpdatedAt,
	)
	return err
}

// GetDecisionTreeResult retrieves a decision tree result by ID
func (s *Store) GetDecisionTreeResult(ctx context.Context, id uuid.UUID) (*DecisionTreeResult, error) {
	var r DecisionTreeResult
	var answersBytes, gpaiBytes, oblBytes, artBytes []byte
	var highRiskLevel string

	err := s.pool.QueryRow(ctx, `
		SELECT id, tenant_id, project_id, system_name, system_description,
			answers, high_risk_level, gpai_result,
			combined_obligations, applicable_articles,
			created_at, updated_at
		FROM ai_act_decision_tree_results WHERE id = $1
	`, id).Scan(
		&r.ID, &r.TenantID, &r.ProjectID, &r.SystemName, &r.SystemDescription,
		&answersBytes, &highRiskLevel, &gpaiBytes,
		&oblBytes, &artBytes,
		&r.CreatedAt, &r.UpdatedAt,
	)

	if err == pgx.ErrNoRows {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}

	json.Unmarshal(answersBytes, &r.Answers)
	json.Unmarshal(gpaiBytes, &r.GPAIResult)
	json.Unmarshal(oblBytes, &r.CombinedObligations)
	json.Unmarshal(artBytes, &r.ApplicableArticles)
	r.HighRiskResult = AIActRiskLevel(highRiskLevel)

	return &r, nil
}

// ListDecisionTreeResults lists all decision tree results for a tenant
func (s *Store) ListDecisionTreeResults(ctx context.Context, tenantID uuid.UUID) ([]DecisionTreeResult, error) {
	rows, err := s.pool.Query(ctx, `
		SELECT id, tenant_id, project_id, system_name, system_description,
			answers, high_risk_level, gpai_result,
			combined_obligations, applicable_articles,
			created_at, updated_at
		FROM ai_act_decision_tree_results
		WHERE tenant_id = $1
		ORDER BY created_at DESC
		LIMIT 100
	`, tenantID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var results []DecisionTreeResult
	for rows.Next() {
		var r DecisionTreeResult
		var answersBytes, gpaiBytes, oblBytes, artBytes []byte
		var highRiskLevel string

		err := rows.Scan(
			&r.ID, &r.TenantID, &r.ProjectID, &r.SystemName, &r.SystemDescription,
			&answersBytes, &highRiskLevel, &gpaiBytes,
			&oblBytes, &artBytes,
			&r.CreatedAt, &r.UpdatedAt,
		)
		if err != nil {
			return nil, err
		}

		json.Unmarshal(answersBytes, &r.Answers)
		json.Unmarshal(gpaiBytes, &r.GPAIResult)
		json.Unmarshal(oblBytes, &r.CombinedObligations)
		json.Unmarshal(artBytes, &r.ApplicableArticles)
		r.HighRiskResult = AIActRiskLevel(highRiskLevel)

		results = append(results, r)
	}

	return results, nil
}

// DeleteDecisionTreeResult deletes a decision tree result by ID
func (s *Store) DeleteDecisionTreeResult(ctx context.Context, id uuid.UUID) error {
	_, err := s.pool.Exec(ctx, "DELETE FROM ai_act_decision_tree_results WHERE id = $1", id)
	return err
}

// ============================================================================
// Helpers
// ============================================================================

// itoa converts int to string for query building
func itoa(i int) string {
	return string(rune('0' + i))
}