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

package ucca

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

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

// ObligationsStore handles persistence of obligations assessments
type ObligationsStore struct {
	pool *pgxpool.Pool
}

// NewObligationsStore creates a new ObligationsStore
func NewObligationsStore(pool *pgxpool.Pool) *ObligationsStore {
	return &ObligationsStore{pool: pool}
}

// CreateAssessment persists a new obligations assessment
func (s *ObligationsStore) CreateAssessment(ctx context.Context, a *ObligationsAssessment) error {
	if a.ID == uuid.Nil {
		a.ID = uuid.New()
	}
	a.CreatedAt = time.Now().UTC()
	a.UpdatedAt = a.CreatedAt

	// Marshal JSON fields
	factsJSON, err := json.Marshal(a.Facts)
	if err != nil {
		return fmt.Errorf("failed to marshal facts: %w", err)
	}

	overviewJSON, err := json.Marshal(a.Overview)
	if err != nil {
		return fmt.Errorf("failed to marshal overview: %w", err)
	}

	_, err = s.pool.Exec(ctx, `
		INSERT INTO obligations_assessments (
			id, tenant_id, organization_name, facts, overview,
			status, created_at, updated_at, created_by
		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
	`,
		a.ID, a.TenantID, a.OrganizationName, factsJSON, overviewJSON,
		a.Status, a.CreatedAt, a.UpdatedAt, a.CreatedBy,
	)

	if err != nil {
		return fmt.Errorf("failed to insert assessment: %w", err)
	}

	return nil
}

// GetAssessment retrieves an assessment by ID
func (s *ObligationsStore) GetAssessment(ctx context.Context, tenantID, assessmentID uuid.UUID) (*ObligationsAssessment, error) {
	var a ObligationsAssessment
	var factsJSON, overviewJSON []byte

	err := s.pool.QueryRow(ctx, `
		SELECT id, tenant_id, organization_name, facts, overview,
			   status, created_at, updated_at, created_by
		FROM obligations_assessments
		WHERE id = $1 AND tenant_id = $2
	`, assessmentID, tenantID).Scan(
		&a.ID, &a.TenantID, &a.OrganizationName, &factsJSON, &overviewJSON,
		&a.Status, &a.CreatedAt, &a.UpdatedAt, &a.CreatedBy,
	)

	if err != nil {
		return nil, fmt.Errorf("assessment not found: %w", err)
	}

	// Unmarshal JSON fields
	if err := json.Unmarshal(factsJSON, &a.Facts); err != nil {
		return nil, fmt.Errorf("failed to unmarshal facts: %w", err)
	}
	if err := json.Unmarshal(overviewJSON, &a.Overview); err != nil {
		return nil, fmt.Errorf("failed to unmarshal overview: %w", err)
	}

	return &a, nil
}

// ListAssessments returns all assessments for a tenant
func (s *ObligationsStore) ListAssessments(ctx context.Context, tenantID uuid.UUID, limit, offset int) ([]*ObligationsAssessment, error) {
	if limit <= 0 {
		limit = 20
	}
	if limit > 100 {
		limit = 100
	}

	rows, err := s.pool.Query(ctx, `
		SELECT id, tenant_id, organization_name, facts, overview,
			   status, created_at, updated_at, created_by
		FROM obligations_assessments
		WHERE tenant_id = $1
		ORDER BY created_at DESC
		LIMIT $2 OFFSET $3
	`, tenantID, limit, offset)

	if err != nil {
		return nil, fmt.Errorf("failed to query assessments: %w", err)
	}
	defer rows.Close()

	var assessments []*ObligationsAssessment
	for rows.Next() {
		var a ObligationsAssessment
		var factsJSON, overviewJSON []byte

		if err := rows.Scan(
			&a.ID, &a.TenantID, &a.OrganizationName, &factsJSON, &overviewJSON,
			&a.Status, &a.CreatedAt, &a.UpdatedAt, &a.CreatedBy,
		); err != nil {
			return nil, fmt.Errorf("failed to scan row: %w", err)
		}

		if err := json.Unmarshal(factsJSON, &a.Facts); err != nil {
			return nil, fmt.Errorf("failed to unmarshal facts: %w", err)
		}
		if err := json.Unmarshal(overviewJSON, &a.Overview); err != nil {
			return nil, fmt.Errorf("failed to unmarshal overview: %w", err)
		}

		assessments = append(assessments, &a)
	}

	return assessments, nil
}

// UpdateAssessment updates an existing assessment
func (s *ObligationsStore) UpdateAssessment(ctx context.Context, a *ObligationsAssessment) error {
	a.UpdatedAt = time.Now().UTC()

	factsJSON, err := json.Marshal(a.Facts)
	if err != nil {
		return fmt.Errorf("failed to marshal facts: %w", err)
	}

	overviewJSON, err := json.Marshal(a.Overview)
	if err != nil {
		return fmt.Errorf("failed to marshal overview: %w", err)
	}

	result, err := s.pool.Exec(ctx, `
		UPDATE obligations_assessments
		SET organization_name = $3, facts = $4, overview = $5,
			status = $6, updated_at = $7
		WHERE id = $1 AND tenant_id = $2
	`,
		a.ID, a.TenantID, a.OrganizationName, factsJSON, overviewJSON,
		a.Status, a.UpdatedAt,
	)

	if err != nil {
		return fmt.Errorf("failed to update assessment: %w", err)
	}

	if result.RowsAffected() == 0 {
		return fmt.Errorf("assessment not found")
	}

	return nil
}

// DeleteAssessment deletes an assessment
func (s *ObligationsStore) DeleteAssessment(ctx context.Context, tenantID, assessmentID uuid.UUID) error {
	result, err := s.pool.Exec(ctx, `
		DELETE FROM obligations_assessments
		WHERE id = $1 AND tenant_id = $2
	`, assessmentID, tenantID)

	if err != nil {
		return fmt.Errorf("failed to delete assessment: %w", err)
	}

	if result.RowsAffected() == 0 {
		return fmt.Errorf("assessment not found")
	}

	return nil
}

// GetMigrationSQL returns the SQL to create the required table
func (s *ObligationsStore) GetMigrationSQL() string {
	return `
-- Obligations Assessments Table
CREATE TABLE IF NOT EXISTS obligations_assessments (
    id UUID PRIMARY KEY,
    tenant_id UUID NOT NULL,
    organization_name VARCHAR(255),
    facts JSONB NOT NULL,
    overview JSONB NOT NULL,
    status VARCHAR(50) NOT NULL DEFAULT 'completed',
    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    created_by UUID,

    CONSTRAINT fk_tenant FOREIGN KEY (tenant_id) REFERENCES tenants(id) ON DELETE CASCADE
);

-- Indexes
CREATE INDEX IF NOT EXISTS idx_obligations_assessments_tenant ON obligations_assessments(tenant_id);
CREATE INDEX IF NOT EXISTS idx_obligations_assessments_created ON obligations_assessments(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_obligations_assessments_status ON obligations_assessments(status);

-- GIN index for JSON queries
CREATE INDEX IF NOT EXISTS idx_obligations_assessments_overview ON obligations_assessments USING GIN (overview);
`
}