package agentops

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

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

// Direct agent lifecycle actions (start/stop/restart) recorded by the receipt
// ledger. Structural actions (create/update/delete/duplicate) flow through the
// approval operation ledger, not this one.
const (
	LifecycleActionStart   = "start"
	LifecycleActionStop    = "stop"
	LifecycleActionRestart = "restart"

	// LifecycleReceiptStatePrepared marks an intent persisted before dispatch.
	LifecycleReceiptStatePrepared = "prepared"
	// LifecycleReceiptStateDispatching marks a runtime command that has been sent
	// but whose result is not yet recorded.
	LifecycleReceiptStateDispatching = "dispatching"
	// LifecycleReceiptStateSucceeded marks a recorded terminal result.
	LifecycleReceiptStateSucceeded = "succeeded"
	// LifecycleReceiptStateIndeterminate marks a command that may or may not have
	// applied (e.g. a crash after dispatch, before result recording). The daemon
	// runtime command has no dedup commandId today, so this is not exactly-once.
	LifecycleReceiptStateIndeterminate = "indeterminate"
)

// ErrLifecyclePayloadConflict signals a requestId was reused with a different
// payload digest. The caller surfaces this as a conflict error.
var ErrLifecyclePayloadConflict = errors.New("agent lifecycle requestId reused with a different payload")

// LifecycleReceipt is one durable direct-lifecycle operation record.
type LifecycleReceipt struct {
	WorkspaceID      string
	ApplicantAgentID string
	OwnerUserID      string
	Action           string
	TargetAgentID    string
	RequestID        string
	PayloadDigest    string
	State            string
	ResultStatus     string
	CorrelationID    string
	CreatedAt        time.Time
	UpdatedAt        time.Time
}

// LifecycleReceiptIntent is the immutable input reserved before dispatch.
type LifecycleReceiptIntent struct {
	WorkspaceID      string
	ApplicantAgentID string
	OwnerUserID      string
	Action           string
	TargetAgentID    string
	RequestID        string
	PayloadDigest    string
	CorrelationID    string
}

func validateLifecycleReceiptIntent(intent LifecycleReceiptIntent) error {
	if strings.TrimSpace(intent.WorkspaceID) == "" || strings.TrimSpace(intent.ApplicantAgentID) == "" ||
		strings.TrimSpace(intent.OwnerUserID) == "" || strings.TrimSpace(intent.TargetAgentID) == "" ||
		strings.TrimSpace(intent.RequestID) == "" || strings.TrimSpace(intent.PayloadDigest) == "" {
		return fmt.Errorf("agent lifecycle receipt intent is incomplete")
	}
	switch intent.Action {
	case LifecycleActionStart, LifecycleActionStop, LifecycleActionRestart:
	default:
		return fmt.Errorf("unsupported agent lifecycle action %q", intent.Action)
	}
	return nil
}

// legalLifecycleResultStatuses pairs (action, state) -> allowed result_status.
// A succeeded receipt only carries the action's success status; an indeterminate
// receipt only carries runtime_outcome_unknown. The store enforces it so a
// non-handler caller cannot record a mismatched pair (e.g. succeeded + unknown).
var legalLifecycleResultStatuses = map[string]map[string]map[string]bool{
	LifecycleActionStart: {
		LifecycleReceiptStateSucceeded:     {"started": true, "already_started": true},
		LifecycleReceiptStateIndeterminate: {"runtime_outcome_unknown": true},
	},
	LifecycleActionStop: {
		LifecycleReceiptStateSucceeded:     {"stopping": true, "already_stopped": true},
		LifecycleReceiptStateIndeterminate: {"runtime_outcome_unknown": true},
	},
	LifecycleActionRestart: {
		LifecycleReceiptStateSucceeded:     {"restarted": true},
		LifecycleReceiptStateIndeterminate: {"runtime_outcome_unknown": true},
	},
}

// validateLifecycleResultRecord guards the Store API from writing arbitrary or
// backwards state. A terminal result is recorded only from dispatching, lands in
// a terminal state, and carries a result status legal for the (action, state).
func validateLifecycleResultRecord(action, expectedCurrentState, state, resultStatus string) error {
	if expectedCurrentState != LifecycleReceiptStateDispatching {
		return fmt.Errorf("agent lifecycle result can only be recorded from %q, got expected state %q",
			LifecycleReceiptStateDispatching, expectedCurrentState)
	}
	switch state {
	case LifecycleReceiptStateSucceeded, LifecycleReceiptStateIndeterminate:
	default:
		return fmt.Errorf("agent lifecycle result state must be terminal (%q or %q), got %q",
			LifecycleReceiptStateSucceeded, LifecycleReceiptStateIndeterminate, state)
	}
	byState, ok := legalLifecycleResultStatuses[action]
	if !ok {
		return fmt.Errorf("unsupported agent lifecycle action %q", action)
	}
	legal, ok := byState[state]
	if !ok {
		return fmt.Errorf("agent lifecycle result state %q is not legal for action %q", state, action)
	}
	if !legal[resultStatus] {
		return fmt.Errorf("agent lifecycle result status %q is not legal for action %q in state %q", resultStatus, action, state)
	}
	return nil
}

type lifecycleRequestSelector struct {
	ApplicantAgentID string
	RequestID        string
	Action           string
	TargetAgentID    string
	PayloadDigest    string
	ReceiptID        int64
	LegacyAmbiguous  bool
}

// ReserveLifecycleReceipt first reserves the action-independent selector, then
// writes the historical action-scoped receipt. The separate reservation table
// is essential for upgrades: legacy duplicate action rows remain replayable,
// while every new cross-action or changed-payload attempt is rejected before a
// second receipt can commit.
func (txs TxStore) ReserveLifecycleReceipt(ctx context.Context, intent LifecycleReceiptIntent) (LifecycleReceipt, bool, error) {
	if err := validateLifecycleReceiptIntent(intent); err != nil {
		return LifecycleReceipt{}, false, err
	}
	selectorResult, err := txs.tx.Exec(ctx, `
INSERT INTO agent_lifecycle_request_selectors (
    applicant_agent_id, request_id, action, target_agent_id, payload_digest, legacy_ambiguous
) VALUES ($1, $2, $3, $4, $5, 0)
ON DUPLICATE KEY UPDATE id = id`,
		intent.ApplicantAgentID, intent.RequestID, intent.Action, intent.TargetAgentID, intent.PayloadDigest)
	if err != nil {
		return LifecycleReceipt{}, false, fmt.Errorf("reserve agent lifecycle request selector: %w", err)
	}
	selectorInserted := selectorResult.RowsAffected() == 1
	selector, found, err := txs.lockLifecycleSelector(ctx, intent.ApplicantAgentID, intent.RequestID)
	if err != nil {
		return LifecycleReceipt{}, false, err
	}
	if !found {
		return LifecycleReceipt{}, false, fmt.Errorf("reserve agent lifecycle request selector: row absent after upsert")
	}

	if !selectorInserted {
		persisted, found, err := txs.lockLifecycleReceipt(ctx, intent.ApplicantAgentID, intent.Action, intent.RequestID)
		if err != nil {
			return LifecycleReceipt{}, false, err
		}
		if selector.LegacyAmbiguous {
			// The old schema permitted one row per action. Exact historical
			// replay remains valid; no new variant may claim the ambiguous key.
			if !found || !lifecycleReceiptMatchesIntent(persisted, intent) {
				return LifecycleReceipt{}, false, ErrLifecyclePayloadConflict
			}
			return persisted, true, nil
		}
		if selector.Action != intent.Action || selector.TargetAgentID != intent.TargetAgentID || selector.PayloadDigest != intent.PayloadDigest {
			return LifecycleReceipt{}, false, ErrLifecyclePayloadConflict
		}
		if !found || !lifecycleReceiptMatchesIntent(persisted, intent) {
			return LifecycleReceipt{}, false, fmt.Errorf("reserve agent lifecycle receipt: selector points to missing or mismatched receipt")
		}
		return persisted, true, nil
	}

	// A newly reserved selector owns the transaction's receipt insert. The
	// caller's transaction rolls both rows back together if this insert fails.
	_, err = txs.tx.Exec(ctx, `
INSERT INTO agent_lifecycle_receipts (
    workspace_id, applicant_agent_id, owner_user_id, action, target_agent_id,
    request_id, payload_digest, state, result_status, correlation_id
) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, '', $9)`,
		intent.WorkspaceID, intent.ApplicantAgentID, intent.OwnerUserID, intent.Action,
		intent.TargetAgentID, intent.RequestID, intent.PayloadDigest,
		LifecycleReceiptStatePrepared, intent.CorrelationID)
	if err != nil {
		return LifecycleReceipt{}, false, fmt.Errorf("reserve agent lifecycle receipt: %w", err)
	}
	persisted, found, err := txs.lockLifecycleReceipt(ctx, intent.ApplicantAgentID, intent.Action, intent.RequestID)
	if err != nil {
		return LifecycleReceipt{}, false, err
	}
	if !found || !lifecycleReceiptMatchesIntent(persisted, intent) {
		return LifecycleReceipt{}, false, fmt.Errorf("reserve agent lifecycle receipt: row absent or mismatched after insert")
	}
	return persisted, false, nil
}

func lifecycleReceiptMatchesIntent(receipt LifecycleReceipt, intent LifecycleReceiptIntent) bool {
	return receipt.WorkspaceID == intent.WorkspaceID &&
		receipt.ApplicantAgentID == intent.ApplicantAgentID &&
		receipt.OwnerUserID == intent.OwnerUserID &&
		receipt.Action == intent.Action &&
		receipt.TargetAgentID == intent.TargetAgentID &&
		receipt.RequestID == intent.RequestID &&
		receipt.PayloadDigest == intent.PayloadDigest
}

func (txs TxStore) lockLifecycleSelector(ctx context.Context, applicantAgentID, requestID string) (lifecycleRequestSelector, bool, error) {
	row := txs.tx.QueryRow(ctx, `
SELECT applicant_agent_id, request_id, action, target_agent_id, payload_digest,
       COALESCE(receipt_id, 0), legacy_ambiguous
FROM agent_lifecycle_request_selectors
WHERE applicant_agent_id = $1 AND request_id = $2
FOR UPDATE`, applicantAgentID, requestID)
	var selector lifecycleRequestSelector
	var action, target, digest sql.NullString
	var receiptID sql.NullInt64
	var ambiguous int
	err := row.Scan(&selector.ApplicantAgentID, &selector.RequestID, &action, &target, &digest, &receiptID, &ambiguous)
	if errors.Is(err, sql.ErrNoRows) {
		return lifecycleRequestSelector{}, false, nil
	}
	if err != nil {
		return lifecycleRequestSelector{}, false, err
	}
	selector.Action, selector.TargetAgentID, selector.PayloadDigest = action.String, target.String, digest.String
	selector.ReceiptID, selector.LegacyAmbiguous = receiptID.Int64, ambiguous != 0
	return selector, true, nil
}

// MarkLifecycleDispatching CAS transitions prepared -> dispatching. It returns
// false (without error) when the receipt is no longer prepared — already
// dispatching (a crash-retry) or already terminal — so the caller can decide the
// replay/indeterminate path instead of re-dispatching blindly.
func (txs TxStore) MarkLifecycleDispatching(ctx context.Context, applicantAgentID, action, requestID string) (bool, error) {
	result, err := txs.tx.Exec(ctx, `
UPDATE agent_lifecycle_receipts SET state = $4
WHERE applicant_agent_id = $1 AND action = $2 AND request_id = $3 AND state = $5`,
		applicantAgentID, action, requestID, LifecycleReceiptStateDispatching, LifecycleReceiptStatePrepared)
	if err != nil {
		return false, err
	}
	return result.RowsAffected() == 1, nil
}

// RecordLifecycleResult CAS transitions the receipt to a terminal state, but
// only from expectedCurrentState. It returns false (without error) when the
// receipt is no longer in that state — for example another path already
// recorded a terminal result — so a terminal result is never overwritten or
// rolled back. Use succeeded for a recorded outcome and indeterminate when
// dispatch may or may not have applied (the caller must not re-dispatch then).
func (txs TxStore) RecordLifecycleResult(ctx context.Context, applicantAgentID, action, requestID, expectedCurrentState, state, resultStatus string) (bool, error) {
	if err := validateLifecycleResultRecord(action, expectedCurrentState, state, resultStatus); err != nil {
		return false, err
	}
	result, err := txs.tx.Exec(ctx, `
UPDATE agent_lifecycle_receipts SET state = $5, result_status = $6
WHERE applicant_agent_id = $1 AND action = $2 AND request_id = $3 AND state = $4`,
		applicantAgentID, action, requestID, expectedCurrentState, state, resultStatus)
	if err != nil {
		return false, err
	}
	return result.RowsAffected() == 1, nil
}

// LoadLifecycleReceipt returns the receipt for the historical action-scoped
// receipt selector. Cross-action conflict protection is owned by the request
// selector reservation table used during ReserveLifecycleReceipt.
func (txs TxStore) LoadLifecycleReceipt(ctx context.Context, applicantAgentID, action, requestID string) (LifecycleReceipt, bool, error) {
	row := txs.tx.QueryRow(ctx, `
SELECT workspace_id, applicant_agent_id, owner_user_id, action, target_agent_id,
       request_id, payload_digest, state, result_status, correlation_id, created_at, updated_at
FROM agent_lifecycle_receipts
WHERE applicant_agent_id = $1 AND action = $2 AND request_id = $3`,
		applicantAgentID, action, requestID)
	var receipt LifecycleReceipt
	err := row.Scan(
		&receipt.WorkspaceID, &receipt.ApplicantAgentID, &receipt.OwnerUserID, &receipt.Action,
		&receipt.TargetAgentID, &receipt.RequestID, &receipt.PayloadDigest, &receipt.State,
		&receipt.ResultStatus, &receipt.CorrelationID, &receipt.CreatedAt, &receipt.UpdatedAt,
	)
	if errors.Is(err, sql.ErrNoRows) {
		return LifecycleReceipt{}, false, nil
	}
	if err != nil {
		return LifecycleReceipt{}, false, err
	}
	return receipt, true, nil
}

// lockLifecycleReceipt is a current read (SELECT ... FOR UPDATE). It sees the
// latest committed row even under repeatable-read, which is what Reserve needs
// after the request-selector reservation lock arbitrates a concurrent reserve.
func (txs TxStore) lockLifecycleReceipt(ctx context.Context, applicantAgentID, action, requestID string) (LifecycleReceipt, bool, error) {
	row := txs.tx.QueryRow(ctx, `
SELECT workspace_id, applicant_agent_id, owner_user_id, action, target_agent_id,
       request_id, payload_digest, state, result_status, correlation_id, created_at, updated_at
FROM agent_lifecycle_receipts
WHERE applicant_agent_id = $1 AND action = $2 AND request_id = $3
FOR UPDATE`, applicantAgentID, action, requestID)
	var receipt LifecycleReceipt
	err := row.Scan(
		&receipt.WorkspaceID, &receipt.ApplicantAgentID, &receipt.OwnerUserID, &receipt.Action,
		&receipt.TargetAgentID, &receipt.RequestID, &receipt.PayloadDigest, &receipt.State,
		&receipt.ResultStatus, &receipt.CorrelationID, &receipt.CreatedAt, &receipt.UpdatedAt,
	)
	if errors.Is(err, sql.ErrNoRows) {
		return LifecycleReceipt{}, false, nil
	}
	if err != nil {
		return LifecycleReceipt{}, false, err
	}
	return receipt, true, nil
}

// Store-level transactional wrappers mirror the runtime_events.go template.

func (store *Store) ReserveLifecycleReceipt(ctx context.Context, intent LifecycleReceiptIntent) (LifecycleReceipt, bool, error) {
	var receipt LifecycleReceipt
	var replayed bool
	err := coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var err error
		receipt, replayed, err = store.InTx(tx).ReserveLifecycleReceipt(ctx, intent)
		return err
	})
	return receipt, replayed, err
}

func (store *Store) MarkLifecycleDispatching(ctx context.Context, applicantAgentID, action, requestID string) (bool, error) {
	var transitioned bool
	err := coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var err error
		transitioned, err = store.InTx(tx).MarkLifecycleDispatching(ctx, applicantAgentID, action, requestID)
		return err
	})
	return transitioned, err
}

func (store *Store) RecordLifecycleResult(ctx context.Context, applicantAgentID, action, requestID, expectedCurrentState, state, resultStatus string) (bool, error) {
	var recorded bool
	err := coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var err error
		recorded, err = store.InTx(tx).RecordLifecycleResult(ctx, applicantAgentID, action, requestID, expectedCurrentState, state, resultStatus)
		return err
	})
	return recorded, err
}

func (store *Store) LoadLifecycleReceipt(ctx context.Context, applicantAgentID, action, requestID string) (LifecycleReceipt, bool, error) {
	var receipt LifecycleReceipt
	var found bool
	err := coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var err error
		receipt, found, err = store.InTx(tx).LoadLifecycleReceipt(ctx, applicantAgentID, action, requestID)
		return err
	})
	return receipt, found, err
}