package run

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"time"

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

// CancelRun atomically terminals the aggregate, cancels every uncommitted
// activation/effect, releases its Run lease, and appends exactly one event. A
// dispatching Effect is already a commitment to perform its external write.
// The same is true after a retryable writer failure returns that Effect to
// pending with attempt_count>0: the external write may have succeeded while
// its response was lost, so its stable idempotent retry must converge first.
// Cancellation therefore conflicts until every crossed-dispatch attempt is
// durably resolved. A repeated cancel is idempotent; successful/failed Runs
// cannot be rewritten.
func (store *Store) CancelRun(ctx context.Context, params CancelRunParams) (Run, *Event, error) {
	return store.cancelRun(ctx, params, nil)
}

// CancelRunForAgent rechecks the current owner policy inside the same
// row-locked transaction that terminals the Run. The preflight Agent read is
// intentionally insufficient: a machine/member/owner lifecycle change between
// show and cancel must make the mutation fail closed.
func (store *Store) CancelRunForAgent(ctx context.Context, params CancelRunParams, principal AgentPrincipal) (Run, *Event, error) {
	return store.cancelRun(ctx, params, func(ctx context.Context, tx coredb.Tx, run Run) error {
		return store.authorizeAgentRunTx(ctx, tx, principal, run.WorkspaceID, run.ID)
	})
}

type cancelRunAuthorization func(context.Context, coredb.Tx, Run) error

func (store *Store) cancelRun(ctx context.Context, params CancelRunParams, authorize cancelRunAuthorization) (Run, *Event, error) {
	if err := store.Validate(); err != nil {
		return Run{}, nil, err
	}
	params.WorkspaceID, params.RunID = strings.TrimSpace(params.WorkspaceID), strings.TrimSpace(params.RunID)
	if params.WorkspaceID == "" || params.RunID == "" {
		return Run{}, nil, storeerr.BadRequest("workspaceId and runId are required")
	}
	if params.Now.IsZero() {
		params.Now = time.Now()
	}
	params.Now = utc(params.Now)
	if len(params.ReasonJSON) == 0 {
		params.ReasonJSON = json.RawMessage(`{"reason":"canceled"}`)
	}
	reason, err := deterministicObject("workflow Run cancellation reason", params.ReasonJSON, true)
	if err != nil {
		return Run{}, nil, storeerr.BadRequest(err.Error())
	}
	if err := store.requireRunExistsForCancellation(ctx, params.WorkspaceID, params.RunID); err != nil {
		return Run{}, nil, err
	}
	store.runRowLockAttempting("cancel") // test-only pre-conflict signal (inert in production)
	var canceled Run
	var appended *Event
	err = coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		// Match Effect claims' lock order (placement -> Run -> Effect) so
		// cancellation serializes without a Run/Effect inversion.
		var lockedWorkspaceID string
		if err := tx.QueryRow(ctx, `
SELECT workspace_id
FROM workspace_placements
WHERE workspace_id=$1
FOR UPDATE`, params.WorkspaceID).Scan(&lockedWorkspaceID); err != nil {
			if coredb.IsNotFound(err) {
				return storeerr.Conflict(fmt.Sprintf("workspace placement for workflow Run %s is missing", params.WorkspaceID))
			}
			return mapReadError("lock workflow Run workspace placement for cancellation", err)
		}
		run, err := scanRun(tx.QueryRow(ctx, runSelectSQL()+` WHERE r.workspace_id=$1 AND r.id=$2 FOR UPDATE`, params.WorkspaceID, params.RunID))
		if err != nil {
			return mapReadError("lock workflow Run for cancellation", err)
		}
		if run.Status == RunStatusCanceled {
			canceled = run
			return nil
		}
		if run.Terminal() {
			return storeerr.Conflict("completed workflow Run cannot be canceled")
		}
		if err := store.runRowLockBarrier("cancel"); err != nil { // test-only fence
			return err
		}
		if authorize != nil {
			if err := authorize(ctx, tx, run); err != nil {
				return err
			}
		}
		state, stateStatus, err := cancelStateSnapshot(run.StateJSON)
		if err != nil {
			return fmt.Errorf("cancel workflow Run state snapshot: %w", err)
		}
		if stateStatus == RunStatusSucceeded || stateStatus == RunStatusFailed {
			// The reducer may have reached a logical terminal state while the
			// aggregate stays waiting only to drain a dispatching Effect. That is
			// completed work, not a still-cancelable graph.
			return storeerr.Conflict("completed workflow Run cannot be canceled")
		}
		var committedEffectID string
		err = tx.QueryRow(ctx, `
SELECT id
FROM workflow_run_effects
WHERE workspace_id=$1 AND run_id=$2
  AND (status='dispatching' OR (status='pending' AND attempt_count>0))
ORDER BY id
FOR UPDATE`, params.WorkspaceID, params.RunID).Scan(&committedEffectID)
		switch {
		case err == nil:
			return storeerr.Conflict("workflow Run has an effect committed for dispatch; retry cancellation after it resolves")
		case !coredb.IsNotFound(err):
			return fmt.Errorf("lock crossed-dispatch workflow Run effects for cancellation: %w", err)
		}
		if _, err := tx.Exec(ctx, `
UPDATE workflow_run_activations
SET status='canceled',completed_at=$3,retry_at=NULL,wake_at=NULL,error_code='run.canceled',error_json=$4,updated_at=$3
WHERE workspace_id=$1 AND run_id=$2 AND status IN ('ready','waiting')`, params.WorkspaceID, params.RunID, params.Now, reason); err != nil {
			return fmt.Errorf("cancel workflow Run activations: %w", err)
		}
		// queued_for_agent callbacks are cancelable and awaiting callbacks already
		// hold an execution lease; terminal them and release every lease here.
		if err := cancelCallbackExecutionLeases(ctx, tx, params.WorkspaceID, params.RunID, "run.canceled", reason, params.Now); err != nil {
			return err
		}
		txStore := store.InTx(tx)
		events, lastSeq, err := txStore.insertEvents(ctx, run, []EventDraft{{Type: "run_canceled", PayloadJSON: reason}}, params.Now)
		if err != nil {
			return err
		}
		result, err := tx.Exec(ctx, `
UPDATE workflow_run_instances
SET status='canceled',state_json=$3,error_code='run.canceled',error_json=$4,
    completed_at=$5,canceled_at=$5,runnable_at=$5,
    lease_owner=NULL,lease_token=NULL,lease_expires_at=NULL,
    last_event_seq=$6,revision=revision+1,updated_at=$5
WHERE workspace_id=$1 AND id=$2 AND revision=$7`, params.WorkspaceID, params.RunID, state, reason, params.Now, lastSeq, run.Revision)
		if err != nil {
			return fmt.Errorf("cancel workflow Run: %w", err)
		}
		if result.RowsAffected() == 0 {
			return storeerr.Conflict("workflow Run changed during cancellation")
		}
		canceled, err = scanRun(tx.QueryRow(ctx, runSelectSQL()+` WHERE r.workspace_id=$1 AND r.id=$2`, params.WorkspaceID, params.RunID))
		if err != nil {
			return fmt.Errorf("load canceled workflow Run: %w", err)
		}
		appended = &events[0]
		return nil
	})
	return canceled, appended, err
}

// requireRunExistsForCancellation distinguishes a missing Run from the
// integrity error caused by its missing placement. The placement must remain
// the first locked row in the mutation transaction, so this proof deliberately
// uses a non-locking read before cancellation starts that transaction.
func (store *Store) requireRunExistsForCancellation(ctx context.Context, workspaceID, runID string) error {
	var existingRunID string
	if err := store.db.QueryRow(ctx, `
SELECT id
FROM workflow_run_instances
WHERE workspace_id=$1 AND id=$2`, workspaceID, runID).Scan(&existingRunID); err != nil {
		return mapReadError("load workflow Run before cancellation", err)
	}
	return nil
}

// cancelStateSnapshot applies the same terminal projection as the pure graph
// reducer while preserving every other reducer field verbatim at the JSON
// value level. It also returns the prior logical status so a reducer-terminal
// Run that is only waiting for Effect drain cannot be rewritten as canceled.
// An empty object is the durable pre-initialization sentinel and is valid for
// a queued Run that was canceled before its first reducer turn.
func cancelStateSnapshot(payload json.RawMessage) (json.RawMessage, string, error) {
	decoder := json.NewDecoder(bytes.NewReader(payload))
	decoder.UseNumber()
	var state map[string]any
	if err := decoder.Decode(&state); err != nil {
		return nil, "", fmt.Errorf("decode state: %w", err)
	}
	if err := requireJSONEOF(decoder); err != nil {
		return nil, "", fmt.Errorf("decode state: %w", err)
	}
	if len(state) == 0 {
		return json.RawMessage(`{}`), "", nil
	}
	status, ok := state["status"].(string)
	if !ok {
		return nil, "", fmt.Errorf("initialized state status is missing or invalid")
	}
	if status == RunStatusSucceeded || status == RunStatusFailed {
		return nil, status, nil
	}
	activations, ok := state["activations"].([]any)
	if !ok {
		return nil, "", fmt.Errorf("initialized state activations are missing or invalid")
	}
	scopes, ok := state["scopes"].([]any)
	if !ok {
		return nil, "", fmt.Errorf("initialized state scopes are missing or invalid")
	}
	state["status"] = RunStatusCanceled
	for index, raw := range activations {
		activation, ok := raw.(map[string]any)
		if !ok {
			return nil, "", fmt.Errorf("initialized state activation %d is invalid", index)
		}
		status, ok := activation["status"].(string)
		if !ok {
			return nil, "", fmt.Errorf("initialized state activation %d status is missing or invalid", index)
		}
		if status == ActivationStatusReady || status == ActivationStatusWaiting {
			activation["status"] = ActivationStatusCanceled
		}
	}
	for index, raw := range scopes {
		scope, ok := raw.(map[string]any)
		if !ok {
			return nil, "", fmt.Errorf("initialized state scope %d is invalid", index)
		}
		status, ok := scope["status"].(string)
		if !ok {
			return nil, "", fmt.Errorf("initialized state scope %d status is missing or invalid", index)
		}
		if status == "active" {
			scope["status"] = "canceled"
		}
	}
	encoded, err := json.Marshal(state)
	if err != nil {
		return nil, "", fmt.Errorf("encode canceled state: %w", err)
	}
	canonical, err := deterministicRunState(encoded)
	if err != nil {
		// A legal snapshot may sit exactly at the aggregate limit; replacing a
		// short live status with "canceled" can add a handful of bytes. Preserve
		// the immutable input but fall back to a bounded terminal snapshot rather
		// than making that Run impossible to cancel.
		input, ok := state["input"].(map[string]any)
		if !ok {
			input = map[string]any{}
		}
		minimal, marshalErr := json.Marshal(map[string]any{
			"status": RunStatusCanceled, "input": input, "variables": map[string]any{},
			"activations": []any{}, "scopes": []any{}, "nextActivation": 0, "nextScope": 0,
		})
		if marshalErr != nil {
			return nil, "", fmt.Errorf("encode bounded canceled workflow Run state: %w", marshalErr)
		}
		canonical, err = deterministicRunState(minimal)
		if err != nil {
			return nil, "", err
		}
	}
	return canonical, status, nil
}

func requireJSONEOF(decoder *json.Decoder) error {
	var extra any
	if err := decoder.Decode(&extra); err != io.EOF {
		if err == nil {
			return fmt.Errorf("must contain exactly one JSON value")
		}
		return err
	}
	return nil
}