package extension

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
)

const (
	CatalogReconcileJobStatusPending = "pending"
	CatalogReconcileJobStatusLeased  = "leased"
	CatalogReconcileJobStatusBlocked = "blocked"
)

// WorkspaceCatalogReconcileJob is one coalesced dirty marker for a Workspace.
// generation is advanced by source-state mutations; processing_generation is
// the generation a worker leased. If generation advances while a worker is
// processing, the release path resets the cursor so the next pass replays the
// whole Workspace under the newer source state.
type WorkspaceCatalogReconcileJob struct {
	WorkspaceID          string
	Generation           int64
	ProcessingGeneration int64
	CursorAgentID        string
	Status               string
	Attempts             int
	LastErrorCode        string
	LastError            string
	LeaseOwner           string
	LeaseExpiresAt       *time.Time
	NextRetryAt          *time.Time
	CreatedAt            time.Time
	UpdatedAt            time.Time
}

// EnqueueWorkspaceCatalogReconcile records that an App/Connection/selection
// mutation may have changed the Workspace's effective Agent Extension Catalog.
// Multiple mutations coalesce into one row by incrementing generation; blocked
// or pending jobs are made pending again, while an active lease is preserved and
// fenced by generation for the worker release path.
func (tx *TxStore) EnqueueWorkspaceCatalogReconcile(ctx context.Context, workspaceID string, now time.Time) (int64, error) {
	if tx == nil || tx.db == nil {
		return 0, fmt.Errorf("extension store transaction is not configured")
	}
	workspaceID = strings.TrimSpace(workspaceID)
	if workspaceID == "" {
		return 0, fmt.Errorf("workspaceId is required")
	}
	now = now.UTC()
	if now.IsZero() {
		now = time.Now().UTC()
	}
	if _, err := tx.db.Exec(ctx, `
INSERT INTO extension_catalog_reconcile_jobs (
    workspace_id, generation, processing_generation, cursor_agent_id,
    status, attempts, last_error_code, last_error, lease_owner,
    lease_expires_at, next_retry_at, created_at, updated_at
) VALUES (
    $1, 1, 0, '',
    'pending', 0, '', '', '',
    NULL, NULL, $2, $2
)
ON DUPLICATE KEY UPDATE
    generation = generation + 1,
    processing_generation = CASE
      WHEN status = 'leased' THEN processing_generation
      ELSE 0
    END,
    cursor_agent_id = CASE
      WHEN status = 'leased' THEN cursor_agent_id
      ELSE ''
    END,
    status = CASE
      WHEN status = 'leased' THEN status
      ELSE 'pending'
    END,
    attempts = CASE
      WHEN status = 'leased' THEN attempts
      ELSE 0
    END,
    last_error_code = '',
    last_error = '',
    next_retry_at = CASE
      WHEN status = 'leased' THEN next_retry_at
      ELSE NULL
    END,
    updated_at = VALUES(updated_at)
`, workspaceID, now); err != nil {
		return 0, fmt.Errorf("enqueue Workspace Extension Catalog reconcile: %w", err)
	}
	var generation int64
	if err := tx.db.QueryRow(ctx, `
SELECT generation
FROM extension_catalog_reconcile_jobs
WHERE workspace_id = $1
`, workspaceID).Scan(&generation); err != nil {
		return 0, fmt.Errorf("read Workspace Extension Catalog reconcile generation: %w", err)
	}
	return generation, nil
}

// ClaimNextWorkspaceCatalogReconcileJob leases one due Workspace dirty job. A
// stale lease whose generation changed while the previous worker was processing
// restarts from the beginning by clearing cursor_agent_id.
func (tx *TxStore) ClaimNextWorkspaceCatalogReconcileJob(ctx context.Context, leaseOwner string, leaseDuration time.Duration, now time.Time) (WorkspaceCatalogReconcileJob, bool, error) {
	if tx == nil || tx.db == nil {
		return WorkspaceCatalogReconcileJob{}, false, fmt.Errorf("extension store transaction is not configured")
	}
	leaseOwner = strings.TrimSpace(leaseOwner)
	if leaseOwner == "" {
		return WorkspaceCatalogReconcileJob{}, false, fmt.Errorf("lease owner is required")
	}
	if leaseDuration <= 0 {
		return WorkspaceCatalogReconcileJob{}, false, fmt.Errorf("lease duration must be positive")
	}
	now = now.UTC()
	if now.IsZero() {
		now = time.Now().UTC()
	}
	var job WorkspaceCatalogReconcileJob
	var leaseExpiresAt sql.NullTime
	var nextRetryAt sql.NullTime
	err := tx.db.QueryRow(ctx, `
SELECT workspace_id, generation, processing_generation, cursor_agent_id,
       status, attempts, last_error_code, last_error,
       lease_owner, lease_expires_at, next_retry_at, created_at, updated_at
FROM extension_catalog_reconcile_jobs
WHERE (
    status = 'pending'
    AND (next_retry_at IS NULL OR next_retry_at <= $1)
  )
  OR (
    status = 'leased'
    AND lease_expires_at IS NOT NULL
    AND lease_expires_at <= $1
  )
ORDER BY updated_at ASC, workspace_id ASC
LIMIT 1
FOR UPDATE
`, now).Scan(
		&job.WorkspaceID,
		&job.Generation,
		&job.ProcessingGeneration,
		&job.CursorAgentID,
		&job.Status,
		&job.Attempts,
		&job.LastErrorCode,
		&job.LastError,
		&job.LeaseOwner,
		&leaseExpiresAt,
		&nextRetryAt,
		&job.CreatedAt,
		&job.UpdatedAt,
	)
	if err != nil {
		if coredb.IsNotFound(err) {
			return WorkspaceCatalogReconcileJob{}, false, nil
		}
		return WorkspaceCatalogReconcileJob{}, false, fmt.Errorf("claim Workspace Extension Catalog reconcile job: %w", err)
	}
	if leaseExpiresAt.Valid {
		value := leaseExpiresAt.Time
		job.LeaseExpiresAt = &value
	}
	if nextRetryAt.Valid {
		value := nextRetryAt.Time
		job.NextRetryAt = &value
	}
	cursor := job.CursorAgentID
	if job.Generation != job.ProcessingGeneration {
		cursor = ""
	}
	attempts := job.Attempts + 1
	leaseExpiresAtTime := now.Add(leaseDuration).UTC()
	result, err := tx.db.Exec(ctx, `
UPDATE extension_catalog_reconcile_jobs
SET status = 'leased',
    processing_generation = generation,
    cursor_agent_id = $2,
    attempts = $3,
    lease_owner = $4,
    lease_expires_at = $5,
    next_retry_at = NULL,
    updated_at = $6
WHERE workspace_id = $1
`, job.WorkspaceID, cursor, attempts, leaseOwner, leaseExpiresAtTime, now)
	if err != nil {
		return WorkspaceCatalogReconcileJob{}, false, fmt.Errorf("lease Workspace Extension Catalog reconcile job: %w", err)
	}
	if result.RowsAffected() == 0 {
		return WorkspaceCatalogReconcileJob{}, false, nil
	}
	job.ProcessingGeneration = job.Generation
	job.CursorAgentID = cursor
	job.Status = CatalogReconcileJobStatusLeased
	job.Attempts = attempts
	job.LeaseOwner = leaseOwner
	job.LeaseExpiresAt = &leaseExpiresAtTime
	job.NextRetryAt = nil
	job.UpdatedAt = now
	return job, true, nil
}

// ReleaseWorkspaceCatalogReconcileBatch releases the lease after one bounded
// batch. If there is more work and generation did not change, it keeps the next
// cursor; if a concurrent source mutation advanced generation, it resets the
// cursor so the newer generation is processed from the beginning. If the batch
// exhausted the Workspace and no newer generation exists, the job row is deleted
// so idle workers have no Workspace state to revisit.
func (tx *TxStore) ReleaseWorkspaceCatalogReconcileBatch(
	ctx context.Context,
	workspaceID string,
	leaseOwner string,
	processingGeneration int64,
	nextCursorAgentID string,
	hasMore bool,
	now time.Time,
) error {
	if tx == nil || tx.db == nil {
		return fmt.Errorf("extension store transaction is not configured")
	}
	workspaceID = strings.TrimSpace(workspaceID)
	leaseOwner = strings.TrimSpace(leaseOwner)
	if workspaceID == "" || leaseOwner == "" || processingGeneration <= 0 {
		return fmt.Errorf("workspaceId, leaseOwner, and processingGeneration are required")
	}
	now = now.UTC()
	if now.IsZero() {
		now = time.Now().UTC()
	}
	if !hasMore {
		result, err := tx.db.Exec(ctx, `
DELETE FROM extension_catalog_reconcile_jobs
WHERE workspace_id = $1
  AND status = 'leased'
  AND lease_owner = $2
  AND processing_generation = $3
  AND generation = $3
`, workspaceID, leaseOwner, processingGeneration)
		if err != nil {
			return fmt.Errorf("complete Workspace Extension Catalog reconcile job: %w", err)
		}
		if result.RowsAffected() != 0 {
			return nil
		}
	}
	if _, err := tx.db.Exec(ctx, `
UPDATE extension_catalog_reconcile_jobs
SET status = 'pending',
    cursor_agent_id = CASE
      WHEN generation = $3 THEN $4
      ELSE ''
    END,
    processing_generation = CASE
      WHEN generation = $3 THEN processing_generation
      ELSE 0
    END,
    lease_owner = '',
    lease_expires_at = NULL,
    next_retry_at = NULL,
    attempts = 0,
    updated_at = $5
WHERE workspace_id = $1
  AND status = 'leased'
  AND lease_owner = $2
  AND processing_generation = $3
`, workspaceID, leaseOwner, processingGeneration, strings.TrimSpace(nextCursorAgentID), now); err != nil {
		return fmt.Errorf("release Workspace Extension Catalog reconcile job: %w", err)
	}
	return nil
}

// BlockWorkspaceCatalogReconcileJob records one terminal/configuration error
// for the Workspace generation. A later source mutation increments generation
// and moves the job back to pending.
func (tx *TxStore) BlockWorkspaceCatalogReconcileJob(ctx context.Context, workspaceID string, leaseOwner string, processingGeneration int64, code string, message string, now time.Time) error {
	if tx == nil || tx.db == nil {
		return fmt.Errorf("extension store transaction is not configured")
	}
	workspaceID = strings.TrimSpace(workspaceID)
	leaseOwner = strings.TrimSpace(leaseOwner)
	if workspaceID == "" || leaseOwner == "" || processingGeneration <= 0 {
		return fmt.Errorf("workspaceId, leaseOwner, and processingGeneration are required")
	}
	now = now.UTC()
	if now.IsZero() {
		now = time.Now().UTC()
	}
	if _, err := tx.db.Exec(ctx, `
UPDATE extension_catalog_reconcile_jobs
SET status = CASE
      WHEN generation = $3 THEN 'blocked'
      ELSE 'pending'
    END,
    cursor_agent_id = CASE
      WHEN generation = $3 THEN cursor_agent_id
      ELSE ''
    END,
    processing_generation = CASE
      WHEN generation = $3 THEN processing_generation
      ELSE 0
    END,
    lease_owner = '',
    lease_expires_at = NULL,
    next_retry_at = NULL,
    last_error_code = CASE
      WHEN generation = $3 THEN $4
      ELSE ''
    END,
    last_error = CASE
      WHEN generation = $3 THEN $5
      ELSE ''
    END,
    updated_at = $6
WHERE workspace_id = $1
  AND status = 'leased'
  AND lease_owner = $2
  AND processing_generation = $3
`, workspaceID, leaseOwner, processingGeneration, truncateCatalogError(code, 255), truncateCatalogError(message, lastErrorMaxLen), now); err != nil {
		return fmt.Errorf("block Workspace Extension Catalog reconcile job: %w", err)
	}
	return nil
}

// RequeueWorkspaceCatalogReconcileJob releases a transiently failed job with a
// bounded retry delay. Retrying starts at the beginning of the Workspace because
// the failed batch may have committed no partial cursor progress.
func (tx *TxStore) RequeueWorkspaceCatalogReconcileJob(ctx context.Context, workspaceID string, leaseOwner string, processingGeneration int64, nextRetryAt time.Time, code string, message string, now time.Time) error {
	if tx == nil || tx.db == nil {
		return fmt.Errorf("extension store transaction is not configured")
	}
	workspaceID = strings.TrimSpace(workspaceID)
	leaseOwner = strings.TrimSpace(leaseOwner)
	if workspaceID == "" || leaseOwner == "" || processingGeneration <= 0 {
		return fmt.Errorf("workspaceId, leaseOwner, and processingGeneration are required")
	}
	now = now.UTC()
	if now.IsZero() {
		now = time.Now().UTC()
	}
	nextRetryAt = nextRetryAt.UTC()
	if _, err := tx.db.Exec(ctx, `
UPDATE extension_catalog_reconcile_jobs
SET status = 'pending',
    cursor_agent_id = '',
    processing_generation = 0,
    lease_owner = '',
    lease_expires_at = NULL,
    next_retry_at = CASE
      WHEN generation = $3 THEN $4
      ELSE NULL
    END,
    last_error_code = CASE
      WHEN generation = $3 THEN $5
      ELSE ''
    END,
    last_error = CASE
      WHEN generation = $3 THEN $6
      ELSE ''
    END,
    updated_at = $7
WHERE workspace_id = $1
  AND status = 'leased'
  AND lease_owner = $2
  AND processing_generation = $3
`, workspaceID, leaseOwner, processingGeneration, nextRetryAt, truncateCatalogError(code, 255), truncateCatalogError(message, lastErrorMaxLen), now); err != nil {
		return fmt.Errorf("requeue Workspace Extension Catalog reconcile job: %w", err)
	}
	return nil
}

func truncateCatalogError(value string, maxLen int) string {
	value = strings.TrimSpace(value)
	if maxLen <= 0 || len(value) <= maxLen {
		return value
	}
	return value[:maxLen]
}