popiart-server/internal/server/sqlite_repo.go

package server

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"

	_ "modernc.org/sqlite"
)

type sqliteRepository struct {
	db        *sql.DB
	secretKey [32]byte
}

func newSQLiteRepository(dbPath, sessionSecret string) (*sqliteRepository, error) {
	if strings.TrimSpace(dbPath) == "" {
		return nil, errors.New("sqlite path is required")
	}
	if err := os.MkdirAll(filepath.Dir(dbPath), 0o755); err != nil {
		return nil, fmt.Errorf("create sqlite directory: %w", err)
	}

	db, err := sql.Open("sqlite", dbPath)
	if err != nil {
		return nil, fmt.Errorf("open sqlite: %w", err)
	}
	db.SetMaxOpenConns(1)

	repo := &sqliteRepository{
		db:        db,
		secretKey: sha256.Sum256([]byte(defaultString(strings.TrimSpace(sessionSecret), "popiart-dev-session-secret"))),
	}
	if err := repo.migrate(); err != nil {
		_ = db.Close()
		return nil, err
	}
	return repo, nil
}

func (r *sqliteRepository) migrate() error {
	statements := []string{
		`PRAGMA foreign_keys = ON;`,
		`CREATE TABLE IF NOT EXISTS sessions (
			session_id TEXT PRIMARY KEY,
			user_id TEXT NOT NULL,
			user_json TEXT NOT NULL,
			token_enc BLOB NOT NULL,
			token_masked TEXT NOT NULL,
			created_at TEXT NOT NULL,
			expires_at TEXT NOT NULL
		);`,
		`CREATE INDEX IF NOT EXISTS idx_sessions_user_id ON sessions(user_id);`,
		`CREATE TABLE IF NOT EXISTS jobs (
			job_id TEXT PRIMARY KEY,
			user_id TEXT NOT NULL,
			session_id TEXT,
			project_id TEXT,
			skill_id TEXT,
			route_key TEXT,
			model_id TEXT,
			exec_mode TEXT NOT NULL,
			newapi_task_id TEXT,
			status TEXT NOT NULL,
			input_json TEXT NOT NULL,
			result_refs_json TEXT,
			usage_json TEXT,
			error_json TEXT,
			idempotency_key TEXT,
			created_at TEXT NOT NULL,
			started_at TEXT,
			finished_at TEXT,
			UNIQUE(user_id, idempotency_key)
		);`,
		`CREATE INDEX IF NOT EXISTS idx_jobs_user_created_at ON jobs(user_id, created_at DESC);`,
		`CREATE INDEX IF NOT EXISTS idx_jobs_status_created_at ON jobs(status, created_at DESC);`,
		`CREATE INDEX IF NOT EXISTS idx_jobs_project_created_at ON jobs(project_id, created_at DESC);`,
		`CREATE INDEX IF NOT EXISTS idx_jobs_skill_created_at ON jobs(skill_id, created_at DESC);`,
		`CREATE TABLE IF NOT EXISTS skill_routes (
			route_key TEXT NOT NULL,
			scope_key TEXT NOT NULL,
			model_id TEXT NOT NULL,
			options_json TEXT,
			updated_at TEXT NOT NULL,
			PRIMARY KEY(route_key, scope_key)
		);`,
	}
	for _, stmt := range statements {
		if _, err := r.db.Exec(stmt); err != nil {
			return fmt.Errorf("apply sqlite migration: %w", err)
		}
	}
	for _, stmt := range []string{
		`ALTER TABLE jobs ADD COLUMN session_id TEXT;`,
		`ALTER TABLE jobs ADD COLUMN route_key TEXT;`,
		`ALTER TABLE jobs ADD COLUMN exec_mode TEXT NOT NULL DEFAULT 'sync_result';`,
		`ALTER TABLE jobs ADD COLUMN newapi_task_id TEXT;`,
		`ALTER TABLE jobs ADD COLUMN result_refs_json TEXT;`,
	} {
		if _, err := r.db.Exec(stmt); err != nil && !strings.Contains(strings.ToLower(err.Error()), "duplicate column name") {
			return fmt.Errorf("upgrade sqlite jobs schema: %w", err)
		}
	}
	return nil
}

func (r *sqliteRepository) CreateSession(upstreamKey string, u user, ttl time.Duration) (session, error) {
	now := time.Now().UTC()
	current := session{
		Token:       "sess_" + randomID(),
		UserID:      u.ID,
		UpstreamKey: upstreamKey,
		User:        u,
		CreatedAt:   now,
		ExpiresAt:   now.Add(ttl),
	}

	enc, err := r.encryptSecret(upstreamKey)
	if err != nil {
		return session{}, err
	}
	userJSON, err := marshalJSON(u)
	if err != nil {
		return session{}, err
	}
	if _, err := r.db.Exec(
		`INSERT INTO sessions (session_id, user_id, user_json, token_enc, token_masked, created_at, expires_at)
		 VALUES (?, ?, ?, ?, ?, ?, ?)`,
		current.Token,
		current.UserID,
		userJSON,
		enc,
		maskSecret(upstreamKey),
		now.Format(time.RFC3339),
		current.ExpiresAt.Format(time.RFC3339),
	); err != nil {
		return session{}, fmt.Errorf("insert session: %w", err)
	}
	return current, nil
}

func (r *sqliteRepository) GetSession(token string) (session, bool, error) {
	row := r.db.QueryRow(
		`SELECT user_id, user_json, token_enc, created_at, expires_at
		 FROM sessions WHERE session_id = ?`,
		token,
	)
	var (
		current              session
		userJSON             string
		tokenEnc             []byte
		createdAt, expiresAt string
	)
	if err := row.Scan(&current.UserID, &userJSON, &tokenEnc, &createdAt, &expiresAt); errors.Is(err, sql.ErrNoRows) {
		return session{}, false, nil
	} else if err != nil {
		return session{}, false, fmt.Errorf("scan session: %w", err)
	}

	if err := unmarshalJSON(userJSON, &current.User); err != nil {
		return session{}, false, fmt.Errorf("decode session user: %w", err)
	}
	upstreamKey, err := r.decryptSecret(tokenEnc)
	if err != nil {
		return session{}, false, err
	}
	current.Token = token
	current.UpstreamKey = upstreamKey
	current.CreatedAt = parseRFC3339(createdAt)
	current.ExpiresAt = parseRFC3339(expiresAt)
	if !current.ExpiresAt.IsZero() && time.Now().UTC().After(current.ExpiresAt) {
		_ = r.DeleteSession(token)
		return session{}, false, nil
	}
	return current, true, nil
}

func (r *sqliteRepository) DeleteSession(token string) error {
	_, err := r.db.Exec(`DELETE FROM sessions WHERE session_id = ?`, token)
	if err != nil {
		return fmt.Errorf("delete session: %w", err)
	}
	return nil
}

func (r *sqliteRepository) RotateSession(oldToken string, ttl time.Duration) (session, bool, error) {
	current, exists, err := r.GetSession(oldToken)
	if err != nil || !exists {
		return session{}, exists, err
	}
	if err := r.DeleteSession(oldToken); err != nil {
		return session{}, false, err
	}
	next, err := r.CreateSession(current.UpstreamKey, current.User, ttl)
	if err != nil {
		return session{}, false, err
	}
	return next, true, nil
}

func (r *sqliteRepository) CreateJob(record *job) (*job, int, error) {
	if record == nil {
		return nil, 0, errors.New("job record is required")
	}
	if record.UserID == "" {
		return nil, 0, errors.New("user_id is required")
	}
	if record.JobID == "" {
		return nil, 0, errors.New("job_id is required")
	}
	if record.CreatedAt == "" {
		record.CreatedAt = time.Now().UTC().Format(time.RFC3339)
	}
	if record.Status == "" {
		record.Status = "pending"
	}

	if record.IdempotencyKey != "" {
		existingID, ok, err := r.findIdempotentJob(record.UserID, record.IdempotencyKey)
		if err != nil {
			return nil, 0, err
		}
		if ok {
			existing, exists, err := r.GetJob(existingID)
			if err != nil {
				return nil, 0, err
			}
			if exists {
				return existing, httpStatusOK, nil
			}
		}
	}

	inputJSON, err := marshalJSON(record.Input)
	if err != nil {
		return nil, 0, err
	}

	if _, err := r.db.Exec(
		`INSERT INTO jobs (
			job_id, user_id, session_id, project_id, skill_id, route_key, model_id,
			exec_mode, newapi_task_id, status, input_json, result_refs_json, usage_json,
			error_json, idempotency_key, created_at, started_at, finished_at
		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, NULL, NULL, NULL, ?, ?, NULL, NULL)`,
		record.JobID,
		record.UserID,
		nullableString(record.SessionID),
		nullableString(record.ProjectID),
		nullableString(record.SkillID),
		nullableString(record.RouteKey),
		nullableString(record.ModelID),
		defaultString(record.ExecMode, "sync_result"),
		nullableString(record.NewAPITaskID),
		record.Status,
		inputJSON,
		nullableString(record.IdempotencyKey),
		record.CreatedAt,
	); err != nil {
		return nil, 0, fmt.Errorf("insert job: %w", err)
	}
	return cloneJob(record), httpStatusAccepted, nil
}

func (r *sqliteRepository) GetJob(jobID string) (*job, bool, error) {
	row := r.db.QueryRow(
		`SELECT job_id, user_id, session_id, project_id, skill_id, route_key, model_id,
		        exec_mode, newapi_task_id, status, input_json, result_refs_json, usage_json,
		        error_json, idempotency_key, created_at, started_at, finished_at
		   FROM jobs WHERE job_id = ?`,
		jobID,
	)
	return scanJob(row)
}

func (r *sqliteRepository) ListJobs(userID, status, skillID, projectID string, limit, offset int) ([]job, int, error) {
	where, args := buildJobFilters(userID, status, skillID, projectID)
	var total int
	if err := r.db.QueryRow(`SELECT COUNT(*) FROM jobs`+where, args...).Scan(&total); err != nil {
		return nil, 0, fmt.Errorf("count jobs: %w", err)
	}

	argsWithPaging := append(append([]any{}, args...), limit, offset)
	rows, err := r.db.Query(
		`SELECT job_id, user_id, session_id, project_id, skill_id, route_key, model_id,
		        exec_mode, newapi_task_id, status, input_json, result_refs_json, usage_json,
		        error_json, idempotency_key, created_at, started_at, finished_at
		   FROM jobs`+where+` ORDER BY datetime(created_at) DESC, job_id DESC LIMIT ? OFFSET ?`,
		argsWithPaging...,
	)
	if err != nil {
		return nil, 0, fmt.Errorf("list jobs: %w", err)
	}
	defer rows.Close()

	items := make([]job, 0, limit)
	for rows.Next() {
		record, _, err := scanJob(rows)
		if err != nil {
			return nil, 0, err
		}
		items = append(items, *record)
	}
	if err := rows.Err(); err != nil {
		return nil, 0, fmt.Errorf("iterate jobs: %w", err)
	}
	return items, total, nil
}

func (r *sqliteRepository) MarkJobRunning(jobID string) (*job, bool, error) {
	record, exists, err := r.GetJob(jobID)
	if err != nil || !exists {
		return nil, exists, err
	}
	if isTerminalStatus(record.Status) {
		return record, true, nil
	}
	now := time.Now().UTC().Format(time.RFC3339)
	if _, err := r.db.Exec(
		`UPDATE jobs SET status = ?, started_at = COALESCE(started_at, ?) WHERE job_id = ?`,
		"running", now, jobID,
	); err != nil {
		return nil, false, fmt.Errorf("mark running job: %w", err)
	}
	return r.GetJob(jobID)
}

func (r *sqliteRepository) LinkUpstreamTask(jobID, upstreamTaskID string) error {
	record, exists, err := r.GetJob(jobID)
	if err != nil {
		return err
	}
	if !exists || isTerminalStatus(record.Status) {
		return nil
	}
	now := time.Now().UTC().Format(time.RFC3339)
	if _, err := r.db.Exec(
		`UPDATE jobs
		    SET status = ?, started_at = COALESCE(started_at, ?), newapi_task_id = ?
		  WHERE job_id = ?`,
		"running", now, nullableString(upstreamTaskID), jobID,
	); err != nil {
		return fmt.Errorf("link upstream task: %w", err)
	}
	return nil
}

func (r *sqliteRepository) FailJob(jobID, code, message string, details map[string]any) error {
	record, exists, err := r.GetJob(jobID)
	if err != nil {
		return err
	}
	if !exists || isTerminalStatus(record.Status) {
		return nil
	}
	now := time.Now().UTC().Format(time.RFC3339)
	errorJSON, err := marshalJSON(&jobError{Code: code, Message: message, Details: details})
	if err != nil {
		return err
	}
	if _, err := r.db.Exec(
		`UPDATE jobs SET status = ?, error_json = ?, finished_at = ? WHERE job_id = ?`,
		"failed", errorJSON, now, jobID,
	); err != nil {
		return fmt.Errorf("update failed job: %w", err)
	}
	return nil
}

func (r *sqliteRepository) CompleteSyncResult(jobID string, refs []resultRef, usage map[string]any) error {
	record, exists, err := r.GetJob(jobID)
	if err != nil {
		return err
	}
	if !exists || isTerminalStatus(record.Status) {
		return nil
	}
	now := time.Now().UTC().Format(time.RFC3339)
	resultJSON, err := marshalJSON(refs)
	if err != nil {
		return err
	}
	usageJSON, err := marshalNullableJSON(usage)
	if err != nil {
		return err
	}
	if _, err := r.db.Exec(
		`UPDATE jobs
		    SET status = ?, result_refs_json = ?, usage_json = ?, finished_at = ?
		  WHERE job_id = ?`,
		"done", resultJSON, usageJSON, now, jobID,
	); err != nil {
		return fmt.Errorf("complete sync result job: %w", err)
	}
	return nil
}

func (r *sqliteRepository) CancelJob(userID, jobID string) (*job, bool, bool, error) {
	record, exists, err := r.GetJob(jobID)
	if err != nil || !exists {
		return nil, exists, false, err
	}
	if record.UserID != userID {
		return nil, false, false, nil
	}
	if isTerminalStatus(record.Status) {
		return nil, true, true, nil
	}
	now := time.Now().UTC().Format(time.RFC3339)
	if _, err := r.db.Exec(
		`UPDATE jobs SET status = ?, finished_at = ? WHERE job_id = ?`,
		"cancelled", now, jobID,
	); err != nil {
		return nil, false, false, fmt.Errorf("cancel job: %w", err)
	}
	record, exists, err = r.GetJob(jobID)
	return record, exists, false, err
}

func (r *sqliteRepository) GetRoutes(projectID string) (map[string]string, error) {
	rows, err := r.db.Query(
		`SELECT route_key, model_id FROM skill_routes WHERE scope_key = ? ORDER BY route_key ASC`,
		scopeKey(projectID),
	)
	if err != nil {
		return nil, fmt.Errorf("query project routes: %w", err)
	}
	defer rows.Close()

	items := map[string]string{}
	for rows.Next() {
		var routeKey, modelID string
		if err := rows.Scan(&routeKey, &modelID); err != nil {
			return nil, fmt.Errorf("scan project route: %w", err)
		}
		items[routeKey] = modelID
	}
	return items, rows.Err()
}

func (r *sqliteRepository) SetRoute(projectID, routeKey, modelID string) error {
	_, err := r.db.Exec(
		`INSERT INTO skill_routes (route_key, scope_key, model_id, options_json, updated_at)
		 VALUES (?, ?, ?, NULL, ?)
		 ON CONFLICT(route_key, scope_key)
		 DO UPDATE SET model_id = excluded.model_id, updated_at = excluded.updated_at`,
		routeKey,
		scopeKey(projectID),
		modelID,
		time.Now().UTC().Format(time.RFC3339),
	)
	if err != nil {
		return fmt.Errorf("upsert route override: %w", err)
	}
	return nil
}

func (r *sqliteRepository) UnsetRoute(projectID, routeKey string) error {
	if _, err := r.db.Exec(
		`DELETE FROM skill_routes WHERE route_key = ? AND scope_key = ?`,
		routeKey,
		scopeKey(projectID),
	); err != nil {
		return fmt.Errorf("delete route override: %w", err)
	}
	return nil
}

func (r *sqliteRepository) findIdempotentJob(userID, idem string) (string, bool, error) {
	var jobID string
	err := r.db.QueryRow(
		`SELECT job_id FROM jobs WHERE user_id = ? AND idempotency_key = ?`,
		userID, idem,
	).Scan(&jobID)
	if errors.Is(err, sql.ErrNoRows) {
		return "", false, nil
	}
	if err != nil {
		return "", false, fmt.Errorf("lookup idempotent job: %w", err)
	}
	return jobID, true, nil
}

func (r *sqliteRepository) encryptSecret(secret string) ([]byte, error) {
	block, err := aes.NewCipher(r.secretKey[:])
	if err != nil {
		return nil, fmt.Errorf("build session cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, fmt.Errorf("build session gcm: %w", err)
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err := rand.Read(nonce); err != nil {
		return nil, fmt.Errorf("generate session nonce: %w", err)
	}
	return gcm.Seal(nonce, nonce, []byte(secret), nil), nil
}

func (r *sqliteRepository) decryptSecret(payload []byte) (string, error) {
	block, err := aes.NewCipher(r.secretKey[:])
	if err != nil {
		return "", fmt.Errorf("build session cipher: %w", err)
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", fmt.Errorf("build session gcm: %w", err)
	}
	if len(payload) < gcm.NonceSize() {
		return "", errors.New("invalid encrypted session token")
	}
	nonce := payload[:gcm.NonceSize()]
	plaintext, err := gcm.Open(nil, nonce, payload[gcm.NonceSize():], nil)
	if err != nil {
		return "", fmt.Errorf("decrypt session token: %w", err)
	}
	return string(plaintext), nil
}

func buildJobFilters(userID, status, skillID, projectID string) (string, []any) {
	filters := []string{"user_id = ?"}
	args := []any{userID}
	if status != "" {
		filters = append(filters, "status = ?")
		args = append(args, status)
	}
	if skillID != "" {
		filters = append(filters, "skill_id = ?")
		args = append(args, skillID)
	}
	if projectID != "" {
		filters = append(filters, "project_id = ?")
		args = append(args, projectID)
	}
	return " WHERE " + strings.Join(filters, " AND "), args
}

func scanJob(row interface{ Scan(dest ...any) error }) (*job, bool, error) {
	var (
		record                                                             job
		sessionID, projectID, skillID, routeKey, modelID                   sql.NullString
		execMode, newapiTaskID, resultRefsJSON, usageJSON, errorJSON, idem sql.NullString
		startedAt, finishedAt                                              sql.NullString
		inputJSON                                                          string
	)
	err := row.Scan(
		&record.JobID,
		&record.UserID,
		&sessionID,
		&projectID,
		&skillID,
		&routeKey,
		&modelID,
		&execMode,
		&newapiTaskID,
		&record.Status,
		&inputJSON,
		&resultRefsJSON,
		&usageJSON,
		&errorJSON,
		&idem,
		&record.CreatedAt,
		&startedAt,
		&finishedAt,
	)
	if errors.Is(err, sql.ErrNoRows) {
		return nil, false, nil
	}
	if err != nil {
		return nil, false, fmt.Errorf("scan job: %w", err)
	}

	record.SessionID = sessionID.String
	record.ProjectID = projectID.String
	record.SkillID = skillID.String
	record.RouteKey = routeKey.String
	record.ModelID = modelID.String
	record.ExecMode = execMode.String
	record.NewAPITaskID = newapiTaskID.String
	record.IdempotencyKey = idem.String
	record.StartedAt = startedAt.String
	record.FinishedAt = finishedAt.String

	if err := unmarshalJSON(inputJSON, &record.Input); err != nil {
		return nil, false, fmt.Errorf("decode job input: %w", err)
	}
	if resultRefsJSON.Valid {
		if err := unmarshalJSON(resultRefsJSON.String, &record.ResultRefs); err != nil {
			return nil, false, fmt.Errorf("decode job result refs: %w", err)
		}
	}
	if usageJSON.Valid {
		if err := unmarshalJSON(usageJSON.String, &record.Usage); err != nil {
			return nil, false, fmt.Errorf("decode job usage: %w", err)
		}
	}
	if errorJSON.Valid {
		var decoded jobError
		if err := unmarshalJSON(errorJSON.String, &decoded); err != nil {
			return nil, false, fmt.Errorf("decode job error: %w", err)
		}
		record.Error = &decoded
	}
	record.ArtifactIDs = buildArtifactIDs(record.JobID, record.ResultRefs)
	return &record, true, nil
}

func marshalJSON(value any) (string, error) {
	data, err := json.Marshal(value)
	if err != nil {
		return "", fmt.Errorf("marshal json: %w", err)
	}
	return string(data), nil
}

func marshalNullableJSON(value any) (any, error) {
	if value == nil {
		return nil, nil
	}
	return marshalJSON(value)
}

func unmarshalJSON(value string, dst any) error {
	if strings.TrimSpace(value) == "" {
		return nil
	}
	return json.Unmarshal([]byte(value), dst)
}

func nullableString(value string) any {
	if strings.TrimSpace(value) == "" {
		return nil
	}
	return value
}

func parseRFC3339(value string) time.Time {
	if strings.TrimSpace(value) == "" {
		return time.Time{}
	}
	ts, err := time.Parse(time.RFC3339, value)
	if err != nil {
		return time.Time{}
	}
	return ts
}

func isTerminalStatus(status string) bool {
	switch status {
	case "done", "failed", "cancelled":
		return true
	default:
		return false
	}
}

func scopeKey(projectID string) string {
	projectID = strings.TrimSpace(projectID)
	if projectID == "" {
		return "__global__"
	}
	return "project:" + projectID
}

const (
	httpStatusOK       = 200
	httpStatusAccepted = 202
)