"""Correction domain models for decision tree editing.

Supports revert and append correction modes with impact analysis,
dry-run reporting, and execution flow tracking.
"""

from __future__ import annotations

from datetime import UTC, datetime
from enum import StrEnum
from typing import Any

from pydantic import BaseModel, ConfigDict, Field, field_validator
from ulid import ULID


class CorrectionMode(StrEnum):
    """Available correction strategies."""

    REVERT = "revert"
    APPEND = "append"


class CorrectionStatus(StrEnum):
    """Lifecycle states for a correction request."""

    PENDING = "pending"
    ANALYZING = "analyzing"
    EXECUTING = "executing"
    APPLIED = "applied"
    FAILED = "failed"
    CANCELLED = "cancelled"
    REJECTED = "rejected"


class ChildPlanState(StrEnum):
    """Observed state of a child plan for cross-plan correction cascading.

    Determined by inspecting the child plan's phase and processing state
    to classify it into one of the four spec-defined categories.
    """

    NOT_STARTED = "not_started"
    IN_PROGRESS = "in_progress"
    COMPLETED_UNAPPLIED = "completed_unapplied"
    APPLIED = "applied"


class CorrectionRequest(BaseModel):
    """A request to correct a decision within a plan's decision tree.

    Each request targets a specific decision node and applies either a
    *revert* (invalidate subtree) or *append* (add new child plan) strategy.
    """

    model_config = ConfigDict(frozen=False, populate_by_name=True)

    correction_id: str = Field(
        default_factory=lambda: str(ULID()),
        description="Unique correction identifier (ULID).",
    )
    plan_id: str = Field(
        ...,
        description="Plan that owns the targeted decision tree.",
    )
    target_decision_id: str = Field(
        ...,
        description="Decision node to apply correction at.",
    )
    mode: CorrectionMode = Field(
        ...,
        description="Correction strategy: revert or append.",
    )
    guidance: str = Field(
        default="",
        max_length=10_000,
        description="Human-supplied guidance for the correction.",
    )
    dry_run: bool = Field(
        default=False,
        description="When True, only compute impact without executing.",
    )
    created_at: datetime = Field(
        default_factory=lambda: datetime.now(UTC),
        description="Timestamp when the request was created.",
    )
    status: CorrectionStatus = Field(
        default=CorrectionStatus.PENDING,
        description="Current lifecycle status.",
    )

    @field_validator("plan_id")
    @classmethod
    def _plan_id_not_empty(cls, v: str) -> str:
        if not v or not v.strip():
            raise ValueError("plan_id must not be empty")
        return v

    @field_validator("target_decision_id")
    @classmethod
    def _target_decision_id_not_empty(cls, v: str) -> str:
        if not v or not v.strip():
            raise ValueError("target_decision_id must not be empty")
        return v

    @field_validator("guidance")
    @classmethod
    def _guidance_stripped(cls, v: str) -> str:
        return v.strip()


class CorrectionImpact(BaseModel):
    """Impact analysis result for a proposed correction.

    Computed via BFS traversal of the decision tree rooted at the
    targeted decision node.
    """

    model_config = ConfigDict(frozen=True)

    affected_decisions: list[str] = Field(
        default_factory=list,
        description="Decision IDs in the affected subtree (BFS order).",
    )
    excluded_decisions: list[str] = Field(
        default_factory=list,
        description="Decision IDs NOT in the affected subtree (preserved).",
    )
    affected_files: list[str] = Field(
        default_factory=list,
        description="Files touched by affected decisions.",
    )
    affected_child_plans: list[str] = Field(
        default_factory=list,
        description="Child plan IDs spawned by affected decisions.",
    )
    estimated_cost: float | None = Field(
        default=None,
        description="Estimated recompute cost (arbitrary units).",
    )
    risk_level: str = Field(
        default="low",
        description="Risk classification: low, medium, or high.",
    )
    rollback_tier: str = Field(
        default="full",
        description="Rollback granularity: full, phase, or append_only.",
    )
    rollback_tier_depth: int = Field(
        default=0,
        description=(
            "Depth of the targeted decision from the tree root. "
            "Tier 0 means the root itself is targeted; tier N means "
            "the target is N hops below the root."
        ),
    )
    artifacts_to_archive: list[str] = Field(
        default_factory=list,
        description="Artifact paths that should be archived on revert.",
    )

    @field_validator("risk_level")
    @classmethod
    def _valid_risk_level(cls, v: str) -> str:
        allowed = {"low", "medium", "high"}
        if v not in allowed:
            raise ValueError(f"risk_level must be one of {allowed}, got '{v}'")
        return v

    @field_validator("rollback_tier")
    @classmethod
    def _valid_rollback_tier(cls, v: str) -> str:
        allowed = {"full", "phase", "append_only"}
        if v not in allowed:
            raise ValueError(f"rollback_tier must be one of {allowed}, got '{v}'")
        return v


class CorrectionResult(BaseModel):
    """Outcome of executing a correction.

    For revert corrections, the result includes fields for checkpoint
    restoration, actor state recovery, phase transition signalling,
    and user-intervention decision creation — enabling full re-execution
    from the decision point per the specification (§ Correction Flow,
    Revert Mode).
    """

    model_config = ConfigDict(frozen=True)

    correction_id: str = Field(
        ...,
        description="Correction request this result belongs to.",
    )
    status: CorrectionStatus = Field(
        ...,
        description="Terminal status after execution.",
    )
    new_decisions: list[str] = Field(
        default_factory=list,
        description="Decision IDs created by an append correction.",
    )
    reverted_decisions: list[str] = Field(
        default_factory=list,
        description="Decision IDs invalidated by a revert correction.",
    )
    spawned_child_plan_id: str | None = Field(
        default=None,
        description="Child plan spawned for append corrections.",
    )
    archived_artifacts: list[str] = Field(
        default_factory=list,
        description="Artifacts archived during revert.",
    )
    error_message: str | None = Field(
        default=None,
        description="Error details when status is FAILED.",
    )
    completed_at: datetime = Field(
        default_factory=lambda: datetime.now(UTC),
        description="Timestamp when execution finished.",
    )

    # --- Revert re-execution fields (spec § Correction Flow) ---

    checkpoint_restored: bool = Field(
        default=False,
        description=(
            "Whether a sandbox checkpoint was restored during revert. "
            "True when the CheckpointService successfully rolled back "
            "resources to the target decision's checkpoint."
        ),
    )
    actor_state_ref: str = Field(
        default="",
        description=(
            "LangGraph actor checkpoint reference restored from the "
            "target decision's context_snapshot.actor_state_ref. "
            "Downstream consumers use this to restore the actor's "
            "reasoning state for re-execution."
        ),
    )
    user_intervention_decision_id: str = Field(
        default="",
        description=(
            "ULID of the user_intervention decision created to inject "
            "the correction guidance into the decision tree."
        ),
    )
    phase_transition_target: str = Field(
        default="",
        description=(
            "Target plan phase after revert (e.g. 'strategize'). "
            "Empty when no phase transition is needed."
        ),
    )


class CorrectionAttempt(BaseModel):
    """Record of a single correction execution attempt.

    Multiple attempts may occur if a correction is retried after failure.
    """

    model_config = ConfigDict(frozen=False)

    attempt_id: str = Field(
        default_factory=lambda: str(ULID()),
        description="Unique attempt identifier (ULID).",
    )
    correction_id: str = Field(
        ...,
        description="Parent correction request ID.",
    )
    started_at: datetime = Field(
        default_factory=lambda: datetime.now(UTC),
        description="When the attempt started.",
    )
    completed_at: datetime | None = Field(
        default=None,
        description="When the attempt finished (None if still running).",
    )
    success: bool = Field(
        default=False,
        description="Whether the attempt succeeded.",
    )
    details: dict[str, Any] = Field(
        default_factory=dict,
        description="Arbitrary execution details.",
    )


class CorrectionDryRunReport(BaseModel):
    """Detailed dry-run report showing what a correction *would* do.

    Generated by ``CorrectionService.generate_dry_run_report`` without
    actually modifying any state.

    Design rationale (Option B — nest-only):
    The embedded ``impact`` object is the single authoritative source for
    impact data.  Previously, ``excluded_decisions``, ``rollback_tier_depth``,
    and ``child_plans_to_rollback`` were duplicated at the top level, creating
    a divergence risk: both models are ``frozen=True``, so post-construction
    mutation is impossible, but nothing prevented constructing an instance
    where the top-level copies disagreed with ``impact``.  Removing the
    redundant fields eliminates that ambiguity.  Consumers should access
    these values via ``report.impact.excluded_decisions``,
    ``report.impact.rollback_tier_depth``, and
    ``report.impact.affected_child_plans`` respectively.
    """

    model_config = ConfigDict(frozen=True)

    correction_id: str = Field(
        ...,
        description="Correction request this report describes.",
    )
    mode: CorrectionMode = Field(
        ...,
        description="Correction strategy that would be applied.",
    )
    impact: CorrectionImpact = Field(
        ...,
        description=(
            "Full impact analysis.  Use ``impact.excluded_decisions``, "
            "``impact.rollback_tier_depth``, and ``impact.affected_child_plans`` "
            "to access data that was previously duplicated at the top level."
        ),
    )
    decisions_to_invalidate: list[str] = Field(
        default_factory=list,
        description="Decisions that would be marked invalid.",
    )
    estimated_recompute_time_seconds: float | None = Field(
        default=None,
        description="Estimated wall-clock seconds to recompute.",
    )
    warnings: list[str] = Field(
        default_factory=list,
        description="Human-readable warnings about the correction.",
    )


class CascadeAction(BaseModel):
    """Describes the action taken on a single child plan during cascade.

    Each affected child plan is inspected and a cascade action is
    determined based on its current state.
    """

    model_config = ConfigDict(frozen=True)

    child_plan_id: str = Field(
        ...,
        description="ID of the affected child plan.",
    )
    child_plan_state: ChildPlanState = Field(
        ...,
        description="Observed state of the child plan.",
    )
    action: str = Field(
        ...,
        description="Action taken: 'cancel', 'cancel_and_rollback', or 'reject'.",
    )
    sandbox_rolled_back: bool = Field(
        default=False,
        description="Whether the child plan's sandbox was rolled back.",
    )

    @field_validator("action")
    @classmethod
    def _valid_action(cls, v: str) -> str:
        allowed = {"cancel", "cancel_and_rollback", "reject"}
        if v not in allowed:
            raise ValueError(f"action must be one of {allowed}, got '{v}'")
        return v


class CorrectionAttemptState(StrEnum):
    """Lifecycle states for a correction attempt record (spec DDL).

    Spec lifecycle: ``pending → executing → complete|failed``.
    Terminal states (``complete``, ``failed``) allow no further transitions.
    """

    PENDING = "pending"
    EXECUTING = "executing"
    COMPLETE = "complete"
    FAILED = "failed"


#: Spec-defined lifecycle transitions for correction attempts.
#: ``pending`` may only move to ``executing``;
#: ``executing`` may only move to ``complete`` or ``failed``.
#: Terminal states have no outgoing transitions.
CORRECTION_ATTEMPT_VALID_TRANSITIONS: dict[
    CorrectionAttemptState, frozenset[CorrectionAttemptState]
] = {
    CorrectionAttemptState.PENDING: frozenset({CorrectionAttemptState.EXECUTING}),
    CorrectionAttemptState.EXECUTING: frozenset(
        {CorrectionAttemptState.COMPLETE, CorrectionAttemptState.FAILED}
    ),
}

#: Terminal states from which no further transitions are allowed.
CORRECTION_ATTEMPT_TERMINAL_STATES: frozenset[CorrectionAttemptState] = frozenset(
    {CorrectionAttemptState.COMPLETE, CorrectionAttemptState.FAILED}
)


def validate_correction_state_transition(
    current: CorrectionAttemptState,
    target: CorrectionAttemptState,
) -> None:
    """Validate a correction attempt state transition per the spec lifecycle.

    Args:
        current: The current state.
        target: The desired target state.

    Raises:
        ValueError: If the transition is not allowed by the spec
            lifecycle (``pending → executing → complete|failed``).
    """
    allowed = CORRECTION_ATTEMPT_VALID_TRANSITIONS.get(
        current,
        frozenset(),
    )
    if target not in allowed:
        msg = f"Invalid state transition from '{current.value}' to '{target.value}'"
        raise ValueError(msg)


class CorrectionAttemptRecord(BaseModel):
    """Spec-aligned record of a correction attempt workflow.

    Maps to the ``correction_attempts`` table defined in the specification
    DDL.  Tracks the full lifecycle of a decision correction including the
    plan, original and new decisions, mode (revert/append), guidance text,
    archived artifacts path, state, and timestamps.
    """

    model_config = ConfigDict(frozen=False, populate_by_name=True)

    correction_attempt_id: str = Field(
        default_factory=lambda: str(ULID()),
        description="Unique identifier (ULID) for the correction attempt.",
    )
    plan_id: str = Field(
        ...,
        description="Plan that owns the targeted decision tree.",
    )
    original_decision_id: str = Field(
        ...,
        description="Decision node being corrected.",
    )
    new_decision_id: str | None = Field(
        default=None,
        description="New decision created by the correction (if any).",
    )
    mode: CorrectionMode = Field(
        ...,
        description="Correction strategy: revert or append.",
    )
    guidance: str = Field(
        ...,
        max_length=10_000,
        description="Human-supplied guidance for the correction.",
    )
    archived_artifacts_path: str | None = Field(
        default=None,
        description="Path to archived artifacts (revert mode).",
    )
    state: CorrectionAttemptState = Field(
        default=CorrectionAttemptState.PENDING,
        description="Current lifecycle state.",
    )
    created_at: datetime = Field(
        default_factory=lambda: datetime.now(UTC),
        description="Timestamp when the record was created.",
    )
    completed_at: datetime | None = Field(
        default=None,
        description="Timestamp when the attempt completed (None if still running).",
    )

    @field_validator("plan_id")
    @classmethod
    def _plan_id_not_empty(cls, v: str) -> str:
        stripped = v.strip()
        if not stripped:
            raise ValueError("plan_id must not be empty")
        return stripped

    @field_validator("original_decision_id")
    @classmethod
    def _original_decision_id_not_empty(cls, v: str) -> str:
        stripped = v.strip()
        if not stripped:
            raise ValueError("original_decision_id must not be empty")
        return stripped

    @field_validator("new_decision_id")
    @classmethod
    def _new_decision_id_not_empty(cls, v: str | None) -> str | None:
        if v is None:
            return v
        stripped = v.strip()
        if not stripped:
            raise ValueError("new_decision_id must not be empty when set")
        return stripped

    @field_validator("guidance")
    @classmethod
    def _guidance_not_empty(cls, v: str) -> str:
        stripped = v.strip()
        if not stripped:
            raise ValueError("guidance must not be empty")
        return stripped

    @field_validator("created_at")
    @classmethod
    def _created_at_tz_aware(cls, v: datetime) -> datetime:
        """Ensure ``created_at`` is timezone-aware (default to UTC if naive)."""
        if v.tzinfo is None:
            return v.replace(tzinfo=UTC)
        return v

    @field_validator("completed_at")
    @classmethod
    def _completed_at_tz_aware(cls, v: datetime | None) -> datetime | None:
        """Ensure ``completed_at`` is timezone-aware when set."""
        if v is not None and v.tzinfo is None:
            return v.replace(tzinfo=UTC)
        return v


class CorrectionRejection(BaseModel):
    """Result when a correction is rejected due to already-applied child plans.

    Returned instead of a ``CorrectionResult`` when the affected subtree
    includes child plans whose changes have already been applied and
    cannot be unilaterally reverted.
    """

    model_config = ConfigDict(frozen=True)

    correction_id: str = Field(
        ...,
        description="Correction request that was rejected.",
    )
    reason: str = Field(
        ...,
        description="Human-readable explanation of why the correction was rejected.",
    )
    affected_applied_child_plan_ids: list[str] = Field(
        default_factory=list,
        description="Child plan IDs that are already applied and block the correction.",
    )
    rejected_at: datetime = Field(
        default_factory=lambda: datetime.now(UTC),
        description="When the rejection occurred.",
    )


class CascadeResult(BaseModel):
    """Outcome of a cross-plan correction cascade operation.

    Captures all actions taken on child plans and whether the cascade
    succeeded or was rejected.
    """

    model_config = ConfigDict(frozen=True)

    correction_id: str = Field(
        ...,
        description="Parent correction request ID.",
    )
    cascade_actions: list[CascadeAction] = Field(
        default_factory=list,
        description="Actions taken on each affected child plan.",
    )
    rejected: bool = Field(
        default=False,
        description="Whether the cascade was rejected.",
    )
    rejection: CorrectionRejection | None = Field(
        default=None,
        description="Rejection details when rejected is True.",
    )
    all_cancelled_plan_ids: list[str] = Field(
        default_factory=list,
        description="Child plan IDs successfully cancelled during cascade.",
    )
    all_rolled_back_plan_ids: list[str] = Field(
        default_factory=list,
        description="Child plan IDs whose sandboxes were rolled back.",
    )


__all__ = [
    "CORRECTION_ATTEMPT_TERMINAL_STATES",
    "CORRECTION_ATTEMPT_VALID_TRANSITIONS",
    "CascadeAction",
    "CascadeResult",
    "ChildPlanState",
    "CorrectionAttempt",
    "CorrectionAttemptRecord",
    "CorrectionAttemptState",
    "CorrectionDryRunReport",
    "CorrectionImpact",
    "CorrectionMode",
    "CorrectionRejection",
    "CorrectionRequest",
    "CorrectionResult",
    "CorrectionStatus",
    "validate_correction_state_transition",
]
