package taskrun

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

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"

	"chat-agent-system/internal/stream"
)

var ErrStoreNotAvailable = errors.New("taskrun: store not available")
var ErrIdempotencyConflict = errors.New("taskrun: idempotency key already exists")

// Store 持久化 task run 与 SSE 事件，支持重启后查询与重放。
type Store interface {
	Available() bool
	InsertRun(ctx context.Context, rec *RunRecord) error
	UpdateRun(ctx context.Context, rec *RunRecord) error
	GetByIdempotencyKey(ctx context.Context, key string) (*RunRecord, error)
	AppendEvent(ctx context.Context, runID string, seq int64, ev stream.StreamEvent) error
	GetRun(ctx context.Context, runID string) (*RunRecord, error)
	ListActiveBySession(ctx context.Context, sessionID string) ([]*RunRecord, error)
	LatestBySession(ctx context.Context, sessionID string) (*RunRecord, error)
	ListEvents(ctx context.Context, runID string, afterSeq int64) ([]StoredEvent, error)
	AcquireLease(ctx context.Context, runID, workerID string, ttl time.Duration) error
	HeartbeatLease(ctx context.Context, runID, workerID string, ttl time.Duration) error
	ReleaseLease(ctx context.Context, runID, workerID string) error
	ResetForResume(ctx context.Context, runID string, attempt int) error
	MarkOrphanedRunning(ctx context.Context) (int, error)
}

// PGStore PostgreSQL 实现。
type PGStore struct {
	pool *pgxpool.Pool
}

func NewPGStore(pool *pgxpool.Pool) *PGStore {
	return &PGStore{pool: pool}
}

func (s *PGStore) Available() bool {
	if s == nil || s.pool == nil {
		return false
	}
	var ok bool
	_ = s.pool.QueryRow(context.Background(), `
		SELECT EXISTS (
			SELECT 1 FROM information_schema.tables
			WHERE table_schema = 'public' AND table_name = 'task_runs'
		)
	`).Scan(&ok)
	return ok
}

func (s *PGStore) hasLeaseColumns(ctx context.Context) bool {
	var ok bool
	_ = s.pool.QueryRow(ctx, `
		SELECT EXISTS (
			SELECT 1 FROM information_schema.columns
			WHERE table_schema = 'public' AND table_name = 'task_runs' AND column_name = 'worker_id'
		)
	`).Scan(&ok)
	return ok
}

func (s *PGStore) InsertRun(ctx context.Context, rec *RunRecord) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	if rec == nil || rec.RunID == "" {
		return fmt.Errorf("taskrun: run_id required")
	}
	attempt := rec.Attempt
	if attempt <= 0 {
		attempt = 1
	}
	idemKey := nullScopeText(rec.IdempotencyKey)
	if s.hasLeaseColumns(ctx) {
		var leaseHeartbeat *time.Time
		if rec.LeaseExpiresAt != nil {
			now := time.Now()
			leaseHeartbeat = &now
		}
		_, err := s.pool.Exec(ctx, `
			INSERT INTO task_runs (
				run_id, session_id, project_id, user_id, client_turn_id, idempotency_key,
				task_key, stage, status, workspace_id, assistant_message_id,
				error_message, worker_id, lease_expires_at, lease_heartbeat_at, attempt,
				workflow_instance_id, stage_instance_id, workflow_stage_key,
				created_at, updated_at
			) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21)
			ON CONFLICT (run_id) DO NOTHING
		`,
			rec.RunID, rec.SessionID, nullInt64(rec.ProjectID), rec.UserID, rec.ClientTurnID, idemKey,
			rec.TaskKey, rec.Stage, rec.Status, rec.WorkspaceID, nullInt64(rec.AssistantMsgID),
			rec.ErrorMessage, rec.WorkerID, rec.LeaseExpiresAt, leaseHeartbeat, attempt,
			nullInt64(rec.WorkflowInstanceID), nullInt64(rec.StageInstanceID), nullScopeText(rec.WorkflowStageKey),
			rec.CreatedAt, rec.UpdatedAt,
		)
		return mapIdempotencyConflict(err)
	}
	_, err := s.pool.Exec(ctx, `
		INSERT INTO task_runs (
			run_id, session_id, project_id, user_id, client_turn_id, idempotency_key,
			task_key, stage, status, workspace_id, assistant_message_id,
			error_message, workflow_instance_id, stage_instance_id, workflow_stage_key,
			created_at, updated_at
		) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17)
		ON CONFLICT (run_id) DO NOTHING
	`,
		rec.RunID, rec.SessionID, nullInt64(rec.ProjectID), rec.UserID, rec.ClientTurnID, idemKey,
		rec.TaskKey, rec.Stage, rec.Status, rec.WorkspaceID, nullInt64(rec.AssistantMsgID),
		rec.ErrorMessage, nullInt64(rec.WorkflowInstanceID), nullInt64(rec.StageInstanceID), nullScopeText(rec.WorkflowStageKey),
		rec.CreatedAt, rec.UpdatedAt,
	)
	return mapIdempotencyConflict(err)
}

func mapIdempotencyConflict(err error) error {
	if err == nil {
		return nil
	}
	var pgErr *pgconn.PgError
	if errors.As(err, &pgErr) && pgErr.Code == "23505" {
		return fmt.Errorf("%w: %v", ErrIdempotencyConflict, err)
	}
	msg := strings.ToLower(err.Error())
	if strings.Contains(msg, "idempotency_key") && (strings.Contains(msg, "duplicate") || strings.Contains(msg, "unique") || strings.Contains(msg, "23505")) {
		return fmt.Errorf("%w: %v", ErrIdempotencyConflict, err)
	}
	return err
}

func (s *PGStore) UpdateRun(ctx context.Context, rec *RunRecord) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	if rec == nil || rec.RunID == "" {
		return fmt.Errorf("taskrun: run_id required")
	}
	if s.hasLeaseColumns(ctx) {
		_, err := s.pool.Exec(ctx, `
			UPDATE task_runs SET
				status = CASE
					WHEN status IN ('completed','cancelled','failed','quality_failed','pending_author_confirmation','finalizing_failed') THEN status
					ELSE $2
				END,
				workspace_id = $3,
				error_message = $4,
				worker_id = $5,
				lease_expires_at = $6,
				assistant_message_id = $7,
				updated_at = $8
			WHERE run_id = $1
		`, rec.RunID, rec.Status, rec.WorkspaceID, rec.ErrorMessage, rec.WorkerID, rec.LeaseExpiresAt, nullInt64(rec.AssistantMsgID), rec.UpdatedAt)
		return err
	}
	_, err := s.pool.Exec(ctx, `
		UPDATE task_runs SET
			status = CASE
				WHEN status IN ('completed','cancelled','failed','quality_failed','pending_author_confirmation','finalizing_failed') THEN status
				ELSE $2
			END,
			workspace_id = $3,
			error_message = $4,
			assistant_message_id = $5,
			updated_at = $6
		WHERE run_id = $1
	`, rec.RunID, rec.Status, rec.WorkspaceID, rec.ErrorMessage, nullInt64(rec.AssistantMsgID), rec.UpdatedAt)
	return err
}

func (s *PGStore) AppendEvent(ctx context.Context, runID string, seq int64, ev stream.StreamEvent) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	payload, err := json.Marshal(ev)
	if err != nil {
		return err
	}
	_, err = s.pool.Exec(ctx, `
		INSERT INTO task_run_events (run_id, seq, event_type, payload)
		VALUES ($1, $2, $3, $4::jsonb)
		ON CONFLICT (run_id, seq) DO NOTHING
	`, runID, seq, string(ev.Type), payload)
	return err
}

func (s *PGStore) GetRun(ctx context.Context, runID string) (*RunRecord, error) {
	if !s.Available() {
		return nil, ErrStoreNotAvailable
	}
	var rec RunRecord
	var projectID, assistantMsgID *int64
	var idempotencyKey *string
	if s.hasLeaseColumns(ctx) {
		err := s.pool.QueryRow(ctx, `
			SELECT run_id, session_id, project_id, user_id, client_turn_id, idempotency_key,
				task_key, stage, status, workspace_id, assistant_message_id,
				error_message, worker_id, lease_expires_at, attempt, created_at, updated_at
			FROM task_runs
			WHERE run_id = $1
		`, runID).Scan(
			&rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID, &idempotencyKey,
			&rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID,
			&rec.ErrorMessage, &rec.WorkerID, &rec.LeaseExpiresAt, &rec.Attempt,
			&rec.CreatedAt, &rec.UpdatedAt,
		)
		if err != nil {
			if errors.Is(err, pgx.ErrNoRows) {
				return nil, fmt.Errorf("run not found")
			}
			return nil, err
		}
	} else {
		err := s.pool.QueryRow(ctx, `
			SELECT run_id, session_id, project_id, user_id, client_turn_id, idempotency_key,
				task_key, stage, status, workspace_id, assistant_message_id,
				error_message, created_at, updated_at
			FROM task_runs
			WHERE run_id = $1
		`, runID).Scan(
			&rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID, &idempotencyKey,
			&rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID,
			&rec.ErrorMessage, &rec.CreatedAt, &rec.UpdatedAt,
		)
		if err != nil {
			if errors.Is(err, pgx.ErrNoRows) {
				return nil, fmt.Errorf("run not found")
			}
			return nil, err
		}
	}
	if projectID != nil {
		rec.ProjectID = *projectID
	}
	if assistantMsgID != nil {
		rec.AssistantMsgID = *assistantMsgID
	}
	if idempotencyKey != nil {
		rec.IdempotencyKey = *idempotencyKey
	}
	return &rec, nil
}

func (s *PGStore) ListActiveBySession(ctx context.Context, sessionID string) ([]*RunRecord, error) {
	if !s.Available() {
		return nil, ErrStoreNotAvailable
	}
	if sessionID == "" {
		return nil, nil
	}
	where := `WHERE session_id = $1 AND status IN ('pending', 'running')`
	if s.hasLeaseColumns(ctx) {
		where += ` AND (lease_expires_at IS NULL OR lease_expires_at > now())`
	}
	where += ` ORDER BY updated_at DESC`
	return s.queryRuns(ctx, where, sessionID)
}

func (s *PGStore) LatestBySession(ctx context.Context, sessionID string) (*RunRecord, error) {
	if !s.Available() {
		return nil, ErrStoreNotAvailable
	}
	if sessionID == "" {
		return nil, nil
	}
	rows, err := s.queryRuns(ctx, `WHERE session_id = $1 ORDER BY updated_at DESC LIMIT 1`, sessionID)
	if err != nil {
		return nil, err
	}
	if len(rows) == 0 {
		return nil, nil
	}
	return rows[0], nil
}

func (s *PGStore) queryRuns(ctx context.Context, where string, args ...any) ([]*RunRecord, error) {
	withLease := s.hasLeaseColumns(ctx)
	q := runSelectSQL(withLease) + " " + where
	rows, err := s.pool.Query(ctx, q, args...)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := make([]*RunRecord, 0)
	for rows.Next() {
		rec, err := scanRunRecord(rows, withLease)
		if err != nil {
			return nil, err
		}
		out = append(out, rec)
	}
	return out, rows.Err()
}

func runSelectSQL(withLease bool) string {
	if withLease {
		return `SELECT run_id, session_id, project_id, user_id, client_turn_id, idempotency_key,
			task_key, stage, status, workspace_id, assistant_message_id,
			error_message, worker_id, lease_expires_at, attempt, created_at, updated_at
			FROM task_runs`
	}
	return `SELECT run_id, session_id, project_id, user_id, client_turn_id, idempotency_key,
		task_key, stage, status, workspace_id, assistant_message_id,
		error_message, created_at, updated_at
		FROM task_runs`
}

type runScanner interface {
	Scan(dest ...any) error
}

func scanRunRecord(row runScanner, withLease bool) (*RunRecord, error) {
	var rec RunRecord
	var projectID, assistantMsgID *int64
	var idempotencyKey *string
	var err error
	if withLease {
		err = row.Scan(
			&rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID, &idempotencyKey,
			&rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID,
			&rec.ErrorMessage, &rec.WorkerID, &rec.LeaseExpiresAt, &rec.Attempt,
			&rec.CreatedAt, &rec.UpdatedAt,
		)
	} else {
		err = row.Scan(
			&rec.RunID, &rec.SessionID, &projectID, &rec.UserID, &rec.ClientTurnID, &idempotencyKey,
			&rec.TaskKey, &rec.Stage, &rec.Status, &rec.WorkspaceID, &assistantMsgID,
			&rec.ErrorMessage, &rec.CreatedAt, &rec.UpdatedAt,
		)
	}
	if err != nil {
		return nil, err
	}
	if projectID != nil {
		rec.ProjectID = *projectID
	}
	if assistantMsgID != nil {
		rec.AssistantMsgID = *assistantMsgID
	}
	if idempotencyKey != nil {
		rec.IdempotencyKey = *idempotencyKey
	}
	return &rec, nil
}

// GetByIdempotencyKey 按稳定幂等键查找既有 run；无键或未命中返回 (nil, nil)。
func (s *PGStore) GetByIdempotencyKey(ctx context.Context, key string) (*RunRecord, error) {
	if !s.Available() {
		return nil, ErrStoreNotAvailable
	}
	if key == "" {
		return nil, nil
	}
	var runID string
	err := s.pool.QueryRow(ctx, `
		SELECT run_id FROM task_runs
		WHERE idempotency_key = $1
		LIMIT 1
	`, key).Scan(&runID)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	return s.GetRun(ctx, runID)
}

func (s *PGStore) ListEvents(ctx context.Context, runID string, afterSeq int64) ([]StoredEvent, error) {
	if !s.Available() {
		return nil, ErrStoreNotAvailable
	}
	rows, err := s.pool.Query(ctx, `
		SELECT seq, payload
		FROM task_run_events
		WHERE run_id = $1 AND seq > $2
		ORDER BY seq ASC
	`, runID, afterSeq)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	out := make([]StoredEvent, 0)
	for rows.Next() {
		var seq int64
		var payload []byte
		if err := rows.Scan(&seq, &payload); err != nil {
			return nil, err
		}
		var ev stream.StreamEvent
		if err := json.Unmarshal(payload, &ev); err != nil {
			return nil, err
		}
		out = append(out, StoredEvent{Seq: seq, Event: ev})
	}
	return out, rows.Err()
}

// CancelOpenByProject cancels pending/running runs for a project inside an existing transaction.
// Table ownership stays with taskrun; workflow terminate must call this instead of writing task_runs directly.
func CancelOpenByProject(ctx context.Context, tx pgx.Tx, projectID int64, at time.Time, errMsg string) error {
	if tx == nil {
		return fmt.Errorf("taskrun: CancelOpenByProject requires tx")
	}
	if strings.TrimSpace(errMsg) == "" {
		errMsg = "workflow terminated"
	}
	_, err := tx.Exec(ctx, `
		UPDATE task_runs SET status = 'cancelled', updated_at = $2, error_message = $3
		WHERE project_id = $1 AND status IN ('pending', 'running')
	`, projectID, at, errMsg)
	return err
}

func (s *PGStore) AcquireLease(ctx context.Context, runID, workerID string, ttl time.Duration) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	if !s.hasLeaseColumns(ctx) {
		return nil
	}
	expires := time.Now().Add(ttl)
	tag, err := s.pool.Exec(ctx, `
		UPDATE task_runs
		SET worker_id = $2,
		    lease_expires_at = $3,
		    lease_heartbeat_at = now(),
		    updated_at = now()
		WHERE run_id = $1 AND status = 'running'
	`, runID, workerID, expires)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("taskrun: lease acquire failed for %s", runID)
	}
	return nil
}

func (s *PGStore) HeartbeatLease(ctx context.Context, runID, workerID string, ttl time.Duration) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	if !s.hasLeaseColumns(ctx) {
		return nil
	}
	expires := time.Now().Add(ttl)
	tag, err := s.pool.Exec(ctx, `
		UPDATE task_runs
		SET lease_expires_at = $3,
		    lease_heartbeat_at = now(),
		    updated_at = now()
		WHERE run_id = $1 AND worker_id = $2 AND status = 'running'
	`, runID, workerID, expires)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("taskrun: heartbeat rejected for %s", runID)
	}
	return nil
}

func (s *PGStore) ReleaseLease(ctx context.Context, runID, workerID string) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	if !s.hasLeaseColumns(ctx) {
		return nil
	}
	_, err := s.pool.Exec(ctx, `
		UPDATE task_runs
		SET worker_id = '',
		    lease_expires_at = NULL,
		    lease_heartbeat_at = NULL,
		    updated_at = now()
		WHERE run_id = $1 AND (worker_id = $2 OR worker_id = '')
	`, runID, workerID)
	return err
}

func (s *PGStore) ResetForResume(ctx context.Context, runID string, attempt int) error {
	if !s.Available() {
		return ErrStoreNotAvailable
	}
	tag, err := s.pool.Exec(ctx, `
		UPDATE task_runs
		SET status = 'running',
		    error_message = '',
		    attempt = $2,
		    updated_at = now()
		WHERE run_id = $1 AND status = 'interrupted'
	`, runID, attempt)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("taskrun: run %s not resumable", runID)
	}
	return nil
}

func (s *PGStore) MarkOrphanedRunning(ctx context.Context) (int, error) {
	if !s.Available() {
		return 0, ErrStoreNotAvailable
	}
	if s.hasLeaseColumns(ctx) {
		tag, err := s.pool.Exec(ctx, `
			UPDATE task_runs
			SET status = 'interrupted',
			    error_message = 'worker lease 过期，任务已中断（进程重启或 worker 失联）',
			    worker_id = '',
			    lease_expires_at = NULL,
			    lease_heartbeat_at = NULL,
			    updated_at = now()
			WHERE status = 'running'
			  AND lease_expires_at IS NOT NULL
			  AND lease_expires_at < now() - interval '2 seconds'
		`)
		if err != nil {
			return 0, err
		}
		return int(tag.RowsAffected()), nil
	}
	tag, err := s.pool.Exec(ctx, `
		UPDATE task_runs
		SET status = 'interrupted',
		    error_message = '进程重启导致任务中断，请重新提交',
		    updated_at = now()
		WHERE status = 'running'
	`)
	if err != nil {
		return 0, err
	}
	return int(tag.RowsAffected()), nil
}

func nullInt64(v int64) *int64 {
	if v <= 0 {
		return nil
	}
	return &v
}

func nullScopeText(v string) *string {
	if v == "" {
		return nil
	}
	return &v
}

// TablesReady 检查 task_runs 表是否已迁移。
func TablesReady(ctx context.Context, pool *pgxpool.Pool) bool {
	if pool == nil {
		return false
	}
	var ok bool
	_ = pool.QueryRow(ctx, `
		SELECT EXISTS (
			SELECT 1 FROM information_schema.tables
			WHERE table_schema = 'public' AND table_name = 'task_runs'
		)
	`).Scan(&ok)
	return ok
}

// LeaseColumnsReady 检查 worker lease 列是否已迁移。
func LeaseColumnsReady(ctx context.Context, pool *pgxpool.Pool) bool {
	if pool == nil {
		return false
	}
	var ok bool
	_ = pool.QueryRow(ctx, `
		SELECT EXISTS (
			SELECT 1 FROM information_schema.columns
			WHERE table_schema = 'public' AND table_name = 'task_runs' AND column_name = 'worker_id'
		)
	`).Scan(&ok)
	return ok
}

// noopStore 无持久化时的空实现。
type noopStore struct{}

func (noopStore) Available() bool { return false }
func (noopStore) InsertRun(context.Context, *RunRecord) error {
	return ErrStoreNotAvailable
}
func (noopStore) UpdateRun(context.Context, *RunRecord) error { return ErrStoreNotAvailable }
func (noopStore) AppendEvent(context.Context, string, int64, stream.StreamEvent) error {
	return ErrStoreNotAvailable
}
func (noopStore) GetRun(context.Context, string) (*RunRecord, error) {
	return nil, ErrStoreNotAvailable
}
func (noopStore) ListActiveBySession(context.Context, string) ([]*RunRecord, error) {
	return nil, ErrStoreNotAvailable
}
func (noopStore) LatestBySession(context.Context, string) (*RunRecord, error) {
	return nil, ErrStoreNotAvailable
}
func (noopStore) GetByIdempotencyKey(context.Context, string) (*RunRecord, error) {
	return nil, ErrStoreNotAvailable
}
func (noopStore) ListEvents(context.Context, string, int64) ([]StoredEvent, error) {
	return nil, ErrStoreNotAvailable
}
func (noopStore) AcquireLease(context.Context, string, string, time.Duration) error {
	return ErrStoreNotAvailable
}
func (noopStore) HeartbeatLease(context.Context, string, string, time.Duration) error {
	return ErrStoreNotAvailable
}
func (noopStore) ReleaseLease(context.Context, string, string) error { return ErrStoreNotAvailable }
func (noopStore) ResetForResume(context.Context, string, int) error {
	return ErrStoreNotAvailable
}
func (noopStore) MarkOrphanedRunning(context.Context) (int, error) { return 0, ErrStoreNotAvailable }

// EnsureUpdatedAt 若 UpdatedAt 为零则设为当前时间。
func EnsureUpdatedAt(rec *RunRecord) {
	if rec != nil && rec.UpdatedAt.IsZero() {
		rec.UpdatedAt = time.Now()
	}
}