package extension

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

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

// CatalogUpdate is one durable catalog-change event awaiting delivery to the
// Agent's machine (Extension Platform stage 6). It carries only THAT the Agent's
// catalog changed at a version — never Action/Schema/credential/endpoint.
type CatalogUpdate struct {
	WorkspaceID    string
	AgentID        string
	Version        int64
	Status         string
	Attempts       int
	LastError      string
	LeaseOwner     string
	LeaseExpiresAt *time.Time
	NextRetryAt    *time.Time
	DeliveredAt    *time.Time
	CreatedAt      time.Time
	UpdatedAt      time.Time
}

// validateCatalogUpdateState enforces the ledger state-machine invariants in the
// app layer (TiDB silently ignores CHECK constraints). Called after scanning a
// row. leased IFF owner non-empty + expiry non-null; delivered IFF delivered_at
// non-null; version > 0; status in the closed set.
func validateCatalogUpdateState(u CatalogUpdate) error {
	if u.Version <= 0 {
		return fmt.Errorf("catalog update version must be > 0, got %d", u.Version)
	}
	// attempts is a monotonic, non-negative retry counter; a negative value is
	// corrupt ledger state (stage-6 final-review fail-closed).
	if u.Attempts < 0 {
		return fmt.Errorf("catalog update attempts must be >= 0, got %d", u.Attempts)
	}
	switch u.Status {
	case "pending", "leased", "delivered", "failed":
	default:
		return fmt.Errorf("catalog update has invalid status %q", u.Status)
	}
	leased := u.Status == "leased"
	hasLease := u.LeaseOwner != "" && u.LeaseExpiresAt != nil
	if leased != hasLease {
		return fmt.Errorf("catalog update %s/%s/%d lease/state inconsistency", u.WorkspaceID, u.AgentID, u.Version)
	}
	delivered := u.Status == "delivered"
	hasDelivered := u.DeliveredAt != nil
	if delivered != hasDelivered {
		return fmt.Errorf("catalog update %s/%s/%d delivered/state inconsistency", u.WorkspaceID, u.AgentID, u.Version)
	}
	return nil
}

// CatalogUpdateClaim claims the next claimable catalog update for delivery. It
// generates a UNIQUE lease token per claim (so an old worker whose lease was
// reclaimed cannot settle the new worker's lease). It claims rows that are
// 'pending' and due (next_retry_at is NULL or in the past) OR 'leased' with an
// expired lease. Returns (update, found).
func (tx *TxStore) ClaimNextCatalogUpdate(ctx context.Context, workerID string, leaseDuration time.Duration, now time.Time) (CatalogUpdate, bool, error) {
	if tx == nil || tx.db == nil {
		return CatalogUpdate{}, false, fmt.Errorf("extension store transaction is not configured")
	}
	if workerID == "" {
		return CatalogUpdate{}, false, fmt.Errorf("worker ID is required")
	}
	if leaseDuration <= 0 {
		return CatalogUpdate{}, false, fmt.Errorf("lease duration must be positive")
	}
	if now.IsZero() {
		now = time.Now().UTC()
	}
	leaseToken, err := identifiers.NewEntityID("extlease")
	if err != nil {
		return CatalogUpdate{}, false, fmt.Errorf("generate lease token: %w", err)
	}
	// Prefix with the worker ID for diagnostics; the full string is the unique
	// token that CAS-guards the settle methods.
	uniqueLeaseOwner := workerID + "/" + leaseToken
	leaseExpires := now.Add(leaseDuration)

	// Lock the next claimable row: pending+due OR expired-lease.
	row := tx.db.QueryRow(ctx, `
SELECT workspace_id, agent_id, version, status, attempts, last_error,
       lease_owner, lease_expires_at, next_retry_at, delivered_at,
       created_at, updated_at
FROM agent_extension_catalog_updates
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 created_at ASC
LIMIT 1
FOR UPDATE`, now)
	update, err := scanCatalogUpdate(row)
	if err != nil {
		if coredb.IsNotFound(err) {
			return CatalogUpdate{}, false, nil
		}
		return CatalogUpdate{}, false, fmt.Errorf("claim catalog update: %w", err)
	}
	if err := validateCatalogUpdateState(update); err != nil {
		return CatalogUpdate{}, false, err
	}

	result, err := tx.db.Exec(ctx, `
UPDATE agent_extension_catalog_updates
SET status = 'leased', lease_owner = $4, lease_expires_at = $5,
    attempts = attempts + 1, next_retry_at = NULL, updated_at = $6
WHERE workspace_id = $1 AND agent_id = $2 AND version = $3`,
		update.WorkspaceID, update.AgentID, update.Version, uniqueLeaseOwner, leaseExpires, now)
	if err != nil {
		return CatalogUpdate{}, false, fmt.Errorf("lease catalog update: %w", err)
	}
	if result.RowsAffected() == 0 {
		return CatalogUpdate{}, false, fmt.Errorf("catalog update %s/%s/%d was concurrently claimed", update.WorkspaceID, update.AgentID, update.Version)
	}

	update.Status = "leased"
	update.LeaseOwner = uniqueLeaseOwner
	update.LeaseExpiresAt = &leaseExpires
	update.Attempts++
	return update, true, nil
}

// MarkCatalogUpdateDelivered marks a leased update as delivered (terminal). Only
// the lease HOLDER (matching leaseOwner) can settle — CAS prevents an old worker
// from settling a reclaimed lease.
func (tx *TxStore) MarkCatalogUpdateDelivered(ctx context.Context, workspaceID, agentID string, version int64, leaseOwner string, now time.Time) error {
	if tx == nil || tx.db == nil {
		return fmt.Errorf("extension store transaction is not configured")
	}
	if leaseOwner == "" {
		return fmt.Errorf("lease owner is required")
	}
	if now.IsZero() {
		now = time.Now().UTC()
	}
	result, err := tx.db.Exec(ctx, `
UPDATE agent_extension_catalog_updates
SET status = 'delivered', delivered_at = $5,
    lease_owner = '', lease_expires_at = NULL, updated_at = $5
WHERE workspace_id = $1 AND agent_id = $2 AND version = $3
  AND status = 'leased' AND lease_owner = $4 AND lease_expires_at > $5`,
		workspaceID, agentID, version, leaseOwner, now)
	if err != nil {
		return fmt.Errorf("mark catalog update delivered: %w", err)
	}
	if result.RowsAffected() == 0 {
		return fmt.Errorf("catalog update %s/%s/%d lease %q not found or not held", workspaceID, agentID, version, leaseOwner)
	}
	return nil
}

// RequeueCatalogUpdate returns a leased update to 'pending' for retry. Only the
// lease HOLDER can requeue.
func (tx *TxStore) RequeueCatalogUpdate(ctx context.Context, workspaceID, agentID string, version int64, leaseOwner string, lastError string, nextRetryAt time.Time, now time.Time) error {
	if tx == nil || tx.db == nil {
		return fmt.Errorf("extension store transaction is not configured")
	}
	if leaseOwner == "" {
		return fmt.Errorf("lease owner is required")
	}
	if now.IsZero() {
		now = time.Now().UTC()
	}
	boundedError := truncateLastError(lastError)
	result, err := tx.db.Exec(ctx, `
UPDATE agent_extension_catalog_updates
SET status = 'pending', lease_owner = '', lease_expires_at = NULL,
    last_error = $5, next_retry_at = $6, delivered_at = NULL, updated_at = $7
WHERE workspace_id = $1 AND agent_id = $2 AND version = $3
  AND status = 'leased' AND lease_owner = $4 AND lease_expires_at > $7`,
		workspaceID, agentID, version, leaseOwner, boundedError, nextRetryAt, now)
	if err != nil {
		return fmt.Errorf("requeue catalog update: %w", err)
	}
	if result.RowsAffected() == 0 {
		return fmt.Errorf("catalog update %s/%s/%d lease %q not found or not held", workspaceID, agentID, version, leaseOwner)
	}
	return nil
}

// FailCatalogUpdate marks a leased update as 'failed' (terminal). Only the lease
// HOLDER can fail.
func (tx *TxStore) FailCatalogUpdate(ctx context.Context, workspaceID, agentID string, version int64, leaseOwner string, lastError string, now time.Time) error {
	if tx == nil || tx.db == nil {
		return fmt.Errorf("extension store transaction is not configured")
	}
	if leaseOwner == "" {
		return fmt.Errorf("lease owner is required")
	}
	if now.IsZero() {
		now = time.Now().UTC()
	}
	boundedError := truncateLastError(lastError)
	result, err := tx.db.Exec(ctx, `
UPDATE agent_extension_catalog_updates
SET status = 'failed', lease_owner = '', lease_expires_at = NULL,
    last_error = $5, next_retry_at = NULL, updated_at = $6
WHERE workspace_id = $1 AND agent_id = $2 AND version = $3
  AND status = 'leased' AND lease_owner = $4 AND lease_expires_at > $6`,
		workspaceID, agentID, version, leaseOwner, boundedError, now)
	if err != nil {
		return fmt.Errorf("fail catalog update: %w", err)
	}
	if result.RowsAffected() == 0 {
		return fmt.Errorf("catalog update %s/%s/%d lease %q not found or not held", workspaceID, agentID, version, leaseOwner)
	}
	return nil
}

func truncateLastError(s string) string {
	s = strings.TrimSpace(s)
	if len(s) > lastErrorMaxLen {
		return s[:lastErrorMaxLen]
	}
	return s
}

func scanCatalogUpdate(row coredb.Row) (CatalogUpdate, error) {
	var u CatalogUpdate
	var leaseExpires, nextRetry, delivered sql.NullTime
	if err := row.Scan(
		&u.WorkspaceID, &u.AgentID, &u.Version, &u.Status, &u.Attempts,
		&u.LastError, &u.LeaseOwner, &leaseExpires, &nextRetry, &delivered,
		&u.CreatedAt, &u.UpdatedAt,
	); err != nil {
		return CatalogUpdate{}, err
	}
	if leaseExpires.Valid {
		t := leaseExpires.Time
		u.LeaseExpiresAt = &t
	}
	if nextRetry.Valid {
		t := nextRetry.Time
		u.NextRetryAt = &t
	}
	if delivered.Valid {
		t := delivered.Time
		u.DeliveredAt = &t
	}
	return u, nil
}