"""Pydantic models for the Areteos Workflow Runtime API."""

from __future__ import annotations

from datetime import datetime
from enum import StrEnum
from typing import Any
from uuid import UUID

from pydantic import BaseModel, ConfigDict, Field

# --- Enums ---


class StepType(StrEnum):
    AGENT = "agent"
    TOOL = "tool"
    HITL = "hitl"
    MAP = "map"


class RunStatus(StrEnum):
    PENDING = "pending"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class StepRunStatus(StrEnum):
    PENDING = "pending"
    RUNNING = "running"
    PAUSED = "paused"
    COMPLETED = "completed"
    FAILED = "failed"
    SKIPPED = "skipped"


class ApprovalStatus(StrEnum):
    PENDING = "pending"
    APPROVED = "approved"
    REJECTED = "rejected"
    RESPONDED = "responded"
    EXPIRED = "expired"
    REASSIGNED = "reassigned"


class ApprovalRequestType(StrEnum):
    APPROVAL = "approval"
    INPUT = "input"
    REVIEW = "review"


class AgentStrategy(StrEnum):
    REACT = "react"
    ADAPTIVE = "adaptive"
    COT = "cot"
    COD = "cod"
    TOT = "tot"
    GOT = "got"
    TRM = "trm"


class ModelAlias(StrEnum):
    FAST = "fast"
    CAPABLE = "capable"
    THINKING = "thinking"


# --- Base ---


class _ApiModel(BaseModel):
    model_config = ConfigDict(
        populate_by_name=True,
        use_enum_values=True,
        strict=False,
    )


# --- Step definition models (for building workflows) ---


class InputMapping(_ApiModel):
    """Maps output from one step (or run input) to another step's input."""

    source_step_id: str | None = None
    """Step ID to read output from. None = use run input."""
    source_field: str
    """Dot-separated path into the source output (e.g., 'result.summary')."""
    target_field: str
    """Key name in this step's input."""


class BranchCondition(_ApiModel):
    """Conditional edge — only followed if the condition is met."""

    field: str
    operator: str  # equals, not_equals, greater_than, less_than, contains, exists
    value: str | None = None


class Edge(_ApiModel):
    """Connects two steps in the workflow DAG."""

    from_step_id: str
    to_step_id: str
    condition: BranchCondition | None = None


class Step(_ApiModel):
    """A single step in a workflow."""

    id: str
    type: StepType
    name: str
    config: dict[str, Any] = Field(default_factory=dict)
    timeout_ms: int = 60_000
    max_retries: int = 3
    mcp_server_ids: list[str] = Field(default_factory=list)
    input_mappings: list[InputMapping] = Field(default_factory=list)

    # --- Convenience constructors ---

    @classmethod
    def agent(
        cls,
        id: str,
        name: str,
        *,
        model: ModelAlias | str = ModelAlias.CAPABLE,
        strategy: AgentStrategy | str = AgentStrategy.REACT,
        system_prompt: str = "",
        max_turns: int = 10,
        max_tokens: int = 4096,
        input_mappings: list[InputMapping] | None = None,
        timeout_ms: int = 300_000,
        max_retries: int = 2,
        mcp_server_ids: list[str] | None = None,
    ) -> Step:
        """Create an LLM agent step."""
        return cls(
            id=id,
            type=StepType.AGENT,
            name=name,
            config={
                "model": str(model),
                "strategy": str(strategy),
                "system_prompt": system_prompt,
                "max_turns": max_turns,
                "max_tokens": max_tokens,
            },
            timeout_ms=timeout_ms,
            max_retries=max_retries,
            mcp_server_ids=mcp_server_ids or [],
            input_mappings=input_mappings or [],
        )

    @classmethod
    def python(
        cls,
        id: str,
        name: str,
        *,
        code: str,
        input_mappings: list[InputMapping] | None = None,
        timeout_ms: int = 30_000,
        max_retries: int = 1,
    ) -> Step:
        """Create a Python inline code step."""
        return cls(
            id=id,
            type=StepType.TOOL,
            name=name,
            config={"tool": "python", "script_content": code},
            timeout_ms=timeout_ms,
            max_retries=max_retries,
            input_mappings=input_mappings or [],
        )

    @classmethod
    def hitl(
        cls,
        id: str,
        name: str,
        *,
        request_type: ApprovalRequestType | str = ApprovalRequestType.APPROVAL,
        timeout_minutes: int | None = None,
        input_mappings: list[InputMapping] | None = None,
    ) -> Step:
        """Create a human-in-the-loop approval step."""
        config: dict[str, Any] = {"request_type": str(request_type)}
        if timeout_minutes is not None:
            config["timeout_minutes"] = timeout_minutes
        return cls(
            id=id,
            type=StepType.HITL,
            name=name,
            config=config,
            input_mappings=input_mappings or [],
        )


# --- Response models (returned from the API) ---


class Workflow(_ApiModel):
    """A registered workflow."""

    id: UUID
    name: str
    description: str | None = None
    scope_type: str | None = None
    current_version_id: UUID | None = None
    owner_id: UUID | None = None
    engagement_id: UUID | None = None
    inserted_at: datetime | None = None
    updated_at: datetime | None = None


class Run(_ApiModel):
    """A workflow execution."""

    id: UUID
    status: RunStatus
    input: dict[str, Any] = Field(default_factory=dict)
    output: dict[str, Any] | None = None
    error_message: str | None = None
    started_at: datetime | None = None
    completed_at: datetime | None = None
    pause_requested: bool = False
    workflow_version_id: UUID | None = None
    triggered_by_id: UUID | None = None
    engagement_id: UUID | None = None
    inserted_at: datetime | None = None
    updated_at: datetime | None = None

    @property
    def is_terminal(self) -> bool:
        return self.status in (RunStatus.COMPLETED, RunStatus.FAILED, RunStatus.CANCELLED)


class StepRun(_ApiModel):
    """Execution record for a single step within a run."""

    id: UUID
    step_id: str
    status: StepRunStatus
    input: dict[str, Any] = Field(default_factory=dict)
    output: dict[str, Any] | None = None
    error_message: str | None = None
    attempt_number: int = 0
    started_at: datetime | None = None
    completed_at: datetime | None = None
    run_id: UUID | None = None
    approval_request_id: UUID | None = None
    inserted_at: datetime | None = None
    updated_at: datetime | None = None


class ApprovalRequest(_ApiModel):
    """A human-in-the-loop approval request."""

    id: UUID
    status: ApprovalStatus
    request_type: ApprovalRequestType | None = None
    title: str | None = None
    instructions: str | None = None
    context: dict[str, Any] = Field(default_factory=dict)
    response: dict[str, Any] | None = None
    response_text: str | None = None
    edited_output: dict[str, Any] | None = None
    responded_at: datetime | None = None
    expires_at: datetime | None = None
    run_id: UUID | None = None
    step_run_id: UUID | None = None
    responded_by_id: UUID | None = None
    assigned_to_id: UUID | None = None
    inserted_at: datetime | None = None
    updated_at: datetime | None = None
