"""Core data models for Agentbyte workflows."""

from __future__ import annotations

from collections.abc import Iterator, MutableMapping
from datetime import datetime
from enum import Enum
from typing import Any, Callable, NoReturn
from typing import TypeVar
from uuid import uuid4

from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field

InputType = TypeVar("InputType", bound=BaseModel)
OutputType = TypeVar("OutputType", bound=BaseModel)


class WorkflowCompletionMode(str, Enum):
    """Controls when a workflow is considered complete."""

    ANY_END = "any_end"   # Stop as soon as ANY declared end step finishes (default)
    ALL_END = "all_end"   # Wait until ALL declared end steps have finished


class JoinFailureMode(str, Enum):
    """Controls how a fan-in (join) step handles upstream branch failures."""

    FAIL_FAST = "fail_fast"       # Re-raise immediately on any branch failure (default)
    SKIP_FAILED = "skip_failed"   # Tolerate partial failures; pass failed branch IDs in envelope


class StepStatus(str, Enum):
    """Status of a step in workflow execution."""

    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    FAILED_TOLERATED = "failed_tolerated"  # failed but downstream join step has SKIP_FAILED mode
    SKIPPED = "skipped"
    SUSPENDED = "suspended"
    CANCELLED = "cancelled"


class WorkflowStatus(str, Enum):
    """Status of workflow execution."""

    CREATED = "created"
    RUNNING = "running"
    SUSPENDED = "suspended"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class WorkflowStateMode(str, Enum):
    """Controls how workflow state mutations become visible across steps."""

    IMMEDIATE = "immediate"
    STAGED = "staged"


class WorkflowStateConflictPolicy(str, Enum):
    """Controls how staged state writes merge at a wave boundary."""

    FAIL = "fail"
    LAST_WRITE_WINS = "last_write_wins"


class WorkflowEventOrigin(str, Enum):
    """Origin category for workflow events.

    `AGENT` and `SUBWORKFLOW` are reserved for later event-forwarding work so
    the enum does not need a breaking expansion once those surfaces emit native
    workflow events.
    """

    FRAMEWORK = "framework"
    STEP = "step"
    AGENT = "agent"
    SUBWORKFLOW = "subworkflow"


class WorkflowRequestKind(str, Enum):
    """Kinds of external workflow requests."""

    INPUT = "input"
    APPROVAL = "approval"


class EdgeCondition(BaseModel):
    """Defines conditions for workflow edges."""

    type: str = Field(
        default="always",
        description="Type of condition: always, output_based, state_based",
    )
    expression: str | None = Field(
        default=None,
        description="Python expression to evaluate",
    )
    field: str | None = Field(
        default=None,
        description="Field to check in output or state",
    )
    value: Any | None = Field(default=None, description="Expected value")
    operator: str | None = Field(
        default=None,
        description="Comparison operator: ==, !=, >, <, in, etc.",
    )


class Edge(BaseModel):
    """Represents a connection between workflow steps."""

    id: str = Field(default_factory=lambda: str(uuid4()))
    from_step: str = Field(description="Source step ID")
    to_step: str = Field(description="Target step ID")
    condition: EdgeCondition = Field(default_factory=EdgeCondition)

    model_config = ConfigDict(extra="forbid")


class StepExecution(BaseModel):
    """Tracks execution details of a step."""

    step_id: str
    status: StepStatus = StepStatus.PENDING
    start_time: datetime | None = None
    end_time: datetime | None = None
    input_data: dict[str, Any] | None = None
    output_data: dict[str, Any] | None = None
    error: str | None = None
    retry_count: int = 0
    pending_request_id: str | None = None

    model_config = ConfigDict(extra="forbid")


class WorkflowPendingRequest(BaseModel):
    """Persisted metadata for a suspended workflow request."""

    request_id: str = Field(default_factory=lambda: str(uuid4()))
    kind: WorkflowRequestKind
    step_id: str
    key: str
    response_type_name: str
    response_type_module: str | None = None
    response_type_qualname: str | None = None
    prompt: str | None = None
    reason: str | None = None
    payload: Any | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)
    created_at: datetime = Field(default_factory=datetime.now)

    model_config = ConfigDict(extra="forbid")


class WorkflowExecution(BaseModel):
    """Tracks execution of an entire workflow."""

    id: str = Field(default_factory=lambda: str(uuid4()))
    workflow_id: str
    status: WorkflowStatus = WorkflowStatus.CREATED
    start_time: datetime | None = None
    end_time: datetime | None = None
    state: dict[str, Any] = Field(default_factory=dict)
    step_executions: dict[str, StepExecution] = Field(default_factory=dict)
    pending_requests: dict[str, WorkflowPendingRequest] = Field(default_factory=dict)
    error: str | None = None

    model_config = ConfigDict(extra="forbid")


class StepMetadata(BaseModel):
    """Metadata for workflow steps."""

    name: str
    description: str | None = None
    tags: list[str] = Field(default_factory=list)
    max_retries: int = 0
    timeout_seconds: float | None = None
    join_failure_mode: JoinFailureMode = JoinFailureMode.FAIL_FAST

    model_config = ConfigDict(extra="forbid")


class WorkflowMetadata(BaseModel):
    """Metadata for workflows."""

    name: str
    description: str | None = None
    version: str = "1.0.0"
    tags: list[str] = Field(default_factory=list)
    author: str | None = None
    created_at: datetime = Field(default_factory=datetime.now)

    model_config = ConfigDict(extra="forbid")


def _describe_type_name(value: Any) -> str:
    """Return a concise display name for a type-like object."""
    return getattr(value, "__name__", repr(value))


def _describe_type_reference(value: Any) -> tuple[str | None, str | None, str]:
    """Return import metadata plus a human-friendly type name."""
    return (
        getattr(value, "__module__", None),
        getattr(value, "__qualname__", None),
        _describe_type_name(value),
    )


class _WorkflowRequestSuspended(Exception):
    """Internal control-flow signal raised when a step requests external input."""

    def __init__(self, request: WorkflowPendingRequest, response_type: Any) -> None:
        super().__init__(f"Workflow request suspended: {request.request_id}")
        self.request = request
        self.response_type = response_type


_STATE_DELETED = object()


class _WorkflowStateView(MutableMapping[str, Any]):
    """A dict-like workflow state view with optional copy-on-write staging."""

    def __init__(
        self,
        committed_state: dict[str, Any],
        *,
        staged_writes: dict[str, Any] | None = None,
    ) -> None:
        self._committed_state = committed_state
        self._staged_writes = staged_writes

    def __getitem__(self, key: str) -> Any:
        if self._staged_writes is not None and key in self._staged_writes:
            value = self._staged_writes[key]
            if value is _STATE_DELETED:
                raise KeyError(key)
            return value
        return self._committed_state[key]

    def __setitem__(self, key: str, value: Any) -> None:
        if self._staged_writes is None:
            self._committed_state[key] = value
            return
        self._staged_writes[key] = value

    def __delitem__(self, key: str) -> None:
        if key not in self:
            raise KeyError(key)
        if self._staged_writes is None:
            del self._committed_state[key]
            return
        self._staged_writes[key] = _STATE_DELETED

    def __iter__(self) -> Iterator[str]:
        visible_keys = set(self._committed_state)
        if self._staged_writes is not None:
            for key, value in self._staged_writes.items():
                if value is _STATE_DELETED:
                    visible_keys.discard(key)
                else:
                    visible_keys.add(key)
        return iter(visible_keys)

    def __len__(self) -> int:
        return sum(1 for _ in self.__iter__())

    def get(self, key: str, default: Any = None) -> Any:
        try:
            return self[key]
        except KeyError:
            return default

    def update(self, other: Any = None, /, **kwargs: Any) -> None:
        updates: dict[str, Any] = {}
        if other is not None:
            updates.update(dict(other))
        updates.update(kwargs)
        for key, value in updates.items():
            self[key] = value

    def as_dict(self) -> dict[str, Any]:
        return {key: self[key] for key in self}

    def staged_writes(self) -> dict[str, Any]:
        return {} if self._staged_writes is None else self._staged_writes.copy()


class WorkflowContext(BaseModel):
    """Typed execution context for workflow steps."""

    state: dict[str, Any] = Field(
        default_factory=dict,
        description="Shared mutable workflow state",
    )
    _progress_callback: Callable[[dict[str, Any]], None] | None = None
    _step_id: str | None = None

    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)

    @classmethod
    def from_state_ref(
        cls,
        state_dict: dict[str, Any],
        *,
        step_id: str | None = None,
    ) -> "WorkflowContext":
        """Create context with direct reference to state dict (no copy)."""
        instance = cls(state={})
        instance.__dict__["state"] = state_dict
        instance.__dict__["_step_id"] = step_id
        return instance

    @classmethod
    def from_staged_state(
        cls,
        committed_state: dict[str, Any],
        *,
        staged_writes: dict[str, Any],
        step_id: str | None = None,
    ) -> "WorkflowContext":
        """Create context with a committed-state view plus staged writes."""
        instance = cls(state={})
        instance.__dict__["state"] = _WorkflowStateView(
            committed_state,
            staged_writes=staged_writes,
        )
        instance.__dict__["_step_id"] = step_id
        return instance

    def get(self, key: str, default: Any = None) -> Any:
        """Get a value from workflow state."""
        return self.state.get(key, default)

    def set(self, key: str, value: Any) -> None:
        """Set a value in workflow state."""
        self.state[key] = value

    def update(self, updates: dict[str, Any]) -> None:
        """Update multiple values in workflow state."""
        self.state.update(updates)

    def emit_progress(
        self,
        message: str,
        completed: int | None = None,
        total: int | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> None:
        """Emit a progress update from within a workflow step."""
        if self._progress_callback:
            progress_data = {
                "message": message,
                "completed": completed,
                "total": total,
                "metadata": metadata or {},
            }
            self._progress_callback(progress_data)

    async def request_input(
        self,
        *,
        key: str,
        response_type: Any,
        prompt: str,
        metadata: dict[str, Any] | None = None,
    ) -> NoReturn:
        """Suspend the workflow and request typed external input."""
        if self._step_id is None:
            raise RuntimeError("WorkflowContext is missing step binding for request_input")

        response_type_module, response_type_qualname, response_type_name = (
            _describe_type_reference(response_type)
        )
        request = WorkflowPendingRequest(
            kind=WorkflowRequestKind.INPUT,
            step_id=self._step_id,
            key=key,
            prompt=prompt,
            metadata=metadata or {},
            response_type_name=response_type_name,
            response_type_module=response_type_module,
            response_type_qualname=response_type_qualname,
        )
        raise _WorkflowRequestSuspended(request, response_type)

    async def request_approval(
        self,
        *,
        key: str,
        reason: str,
        payload: Any | None = None,
        response_type: Any = bool,
        metadata: dict[str, Any] | None = None,
    ) -> NoReturn:
        """Suspend the workflow and request approval from an external caller."""
        if self._step_id is None:
            raise RuntimeError("WorkflowContext is missing step binding for request_approval")

        response_type_module, response_type_qualname, response_type_name = (
            _describe_type_reference(response_type)
        )
        request = WorkflowPendingRequest(
            kind=WorkflowRequestKind.APPROVAL,
            step_id=self._step_id,
            key=key,
            reason=reason,
            payload=payload,
            metadata=metadata or {},
            response_type_name=response_type_name,
            response_type_module=response_type_module,
            response_type_qualname=response_type_qualname,
        )
        raise _WorkflowRequestSuspended(request, response_type)

    def to_dict(self) -> dict[str, Any]:
        """Convert context to dictionary."""
        return {"workflow_state": dict(self.state)}


class WorkflowValidationResult(BaseModel):
    """Result of workflow validation."""

    is_valid: bool
    errors: list[str] = Field(default_factory=list)
    warnings: list[str] = Field(default_factory=list)
    has_cycles: bool = False
    unreachable_steps: list[str] = Field(default_factory=list)

    model_config = ConfigDict(extra="forbid")


class WorkflowEventType(str, Enum):
    """Types of workflow events."""

    WORKFLOW_STARTED = "workflow_started"
    WORKFLOW_COMPLETED = "workflow_completed"
    WORKFLOW_FAILED = "workflow_failed"
    WORKFLOW_CANCELLED = "workflow_cancelled"
    WORKFLOW_SUSPENDED = "workflow_suspended"
    WORKFLOW_RESUMED = "workflow_resumed"
    CHECKPOINT_SAVED = "checkpoint_saved"
    WORKFLOW_INPUT_REQUESTED = "workflow_input_requested"
    WORKFLOW_APPROVAL_REQUESTED = "workflow_approval_requested"
    STEP_STARTED = "step_started"
    STEP_COMPLETED = "step_completed"
    STEP_FAILED = "step_failed"
    STEP_PROGRESS = "step_progress"
    EDGE_ACTIVATED = "edge_activated"


class WorkflowEvent(BaseModel):
    """Base class for workflow events."""

    event_type: WorkflowEventType
    timestamp: datetime
    workflow_id: str
    origin: WorkflowEventOrigin = WorkflowEventOrigin.FRAMEWORK

    model_config = ConfigDict(extra="forbid")

    def __str__(self) -> str:
        time_str = self.timestamp.strftime("%H:%M:%S")
        return f"[{time_str}] {self.event_type.value}"

    def __repr__(self) -> str:
        class_name = self.__class__.__name__
        return (
            f"{class_name}(event_type='{self.event_type.value}', "
            f"workflow_id='{self.workflow_id[:8]}...', "
            f"timestamp='{self.timestamp}')"
        )


class WorkflowErrorDetails(BaseModel):
    """Structured error payload for workflow failure events."""

    error_type: str
    message: str
    traceback: str | None = None
    workflow_id: str | None = None
    step_id: str | None = None
    execution_id: str | None = None

    model_config = ConfigDict(extra="forbid")


class WorkflowStartedEvent(WorkflowEvent):
    """Workflow execution started."""

    event_type: WorkflowEventType = WorkflowEventType.WORKFLOW_STARTED
    initial_input: dict[str, Any]

    def __str__(self) -> str:
        time_str = self.timestamp.strftime("%H:%M:%S")
        return f"[{time_str}] 🚀 Workflow started with input: {self.initial_input}"


class WorkflowCompletedEvent(WorkflowEvent):
    """Workflow execution completed successfully."""

    event_type: WorkflowEventType = WorkflowEventType.WORKFLOW_COMPLETED
    execution: WorkflowExecution
    final_output: dict[str, Any] | None = None

    def __str__(self) -> str:
        time_str = self.timestamp.strftime("%H:%M:%S")
        duration = None
        if self.execution.start_time and self.execution.end_time:
            duration = (
                self.execution.end_time - self.execution.start_time
            ).total_seconds()
        duration_str = f" in {duration:.2f}s" if duration else ""
        return (
            f"[{time_str}] ✅ Workflow completed{duration_str} "
            f"({len(self.execution.step_executions)} steps)"
        )


class WorkflowFailedEvent(WorkflowEvent):
    """Workflow execution failed."""

    event_type: WorkflowEventType = WorkflowEventType.WORKFLOW_FAILED
    error_details: WorkflowErrorDetails
    execution: WorkflowExecution | None = None

    @property
    def error(self) -> str:
        """Compatibility accessor for the primary error message."""
        return self.error_details.message


class WorkflowCancelledEvent(WorkflowEvent):
    """Workflow execution was cancelled."""

    event_type: WorkflowEventType = WorkflowEventType.WORKFLOW_CANCELLED
    execution: WorkflowExecution
    reason: str


class WorkflowSuspendedEvent(WorkflowEvent):
    """Workflow execution suspended awaiting external input."""

    event_type: WorkflowEventType = WorkflowEventType.WORKFLOW_SUSPENDED
    execution: WorkflowExecution
    reason: str


class WorkflowResumedEvent(WorkflowEvent):
    """Workflow resumed from checkpoint."""

    event_type: WorkflowEventType = WorkflowEventType.WORKFLOW_RESUMED
    checkpoint_id: str
    completed_steps: list[str]
    pending_steps: list[str]

    def __str__(self) -> str:
        time_str = self.timestamp.strftime("%H:%M:%S")
        return (
            f"[{time_str}] 🔄 Resumed from checkpoint "
            f"({len(self.completed_steps)} completed, "
            f"{len(self.pending_steps)} pending)"
        )


class CheckpointSavedEvent(WorkflowEvent):
    """Checkpoint saved during execution."""

    event_type: WorkflowEventType = WorkflowEventType.CHECKPOINT_SAVED
    checkpoint_id: str
    completed_steps: int
    total_steps: int

    def __str__(self) -> str:
        time_str = self.timestamp.strftime("%H:%M:%S")
        progress = (
            (self.completed_steps / self.total_steps * 100)
            if self.total_steps > 0
            else 0
        )
        return (
            f"[{time_str}] 💾 Checkpoint saved "
            f"({self.completed_steps}/{self.total_steps} steps, {progress:.0f}%)"
        )


class WorkflowInputRequestedEvent(WorkflowEvent):
    """A workflow step requested typed external input."""

    event_type: WorkflowEventType = WorkflowEventType.WORKFLOW_INPUT_REQUESTED
    origin: WorkflowEventOrigin = WorkflowEventOrigin.STEP
    step_id: str
    request: WorkflowPendingRequest


class WorkflowApprovalRequestedEvent(WorkflowEvent):
    """A workflow step requested external approval."""

    event_type: WorkflowEventType = WorkflowEventType.WORKFLOW_APPROVAL_REQUESTED
    origin: WorkflowEventOrigin = WorkflowEventOrigin.STEP
    step_id: str
    request: WorkflowPendingRequest


class StepStartedEvent(WorkflowEvent):
    """Step execution started."""

    event_type: WorkflowEventType = WorkflowEventType.STEP_STARTED
    origin: WorkflowEventOrigin = WorkflowEventOrigin.FRAMEWORK
    step_id: str
    input_data: dict[str, Any]

    def __str__(self) -> str:
        time_str = self.timestamp.strftime("%H:%M:%S")
        return f"[{time_str}] ▶️  Step '{self.step_id}' started"


class StepCompletedEvent(WorkflowEvent):
    """Step execution completed successfully."""

    event_type: WorkflowEventType = WorkflowEventType.STEP_COMPLETED
    origin: WorkflowEventOrigin = WorkflowEventOrigin.FRAMEWORK
    step_id: str
    output_data: dict[str, Any]
    duration_seconds: float

    def __str__(self) -> str:
        time_str = self.timestamp.strftime("%H:%M:%S")
        return f"[{time_str}] ✅ Step '{self.step_id}' completed → {self.output_data}"


class StepFailedEvent(WorkflowEvent):
    """Step execution failed."""

    event_type: WorkflowEventType = WorkflowEventType.STEP_FAILED
    origin: WorkflowEventOrigin = WorkflowEventOrigin.FRAMEWORK
    step_id: str
    error_details: WorkflowErrorDetails
    duration_seconds: float

    @property
    def error(self) -> str:
        """Compatibility accessor for the primary error message."""
        return self.error_details.message

    def __str__(self) -> str:
        time_str = self.timestamp.strftime("%H:%M:%S")
        return f"[{time_str}] ❌ Step '{self.step_id}' failed: {self.error_details.message}"


class StepProgressEvent(WorkflowEvent):
    """Progress update from within a step execution."""

    event_type: WorkflowEventType = WorkflowEventType.STEP_PROGRESS
    origin: WorkflowEventOrigin = WorkflowEventOrigin.STEP
    step_id: str
    message: str
    completed: int | None = None
    total: int | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)

    def __str__(self) -> str:
        time_str = self.timestamp.strftime("%H:%M:%S")
        progress_str = ""
        if self.completed is not None and self.total is not None:
            percentage = (self.completed / self.total * 100) if self.total > 0 else 0
            progress_str = f" ({self.completed}/{self.total}, {percentage:.0f}%)"
        return f"[{time_str}] 🔄 Step '{self.step_id}': {self.message}{progress_str}"


class EdgeActivatedEvent(WorkflowEvent):
    """Edge between steps activated (data flowing)."""

    event_type: WorkflowEventType = WorkflowEventType.EDGE_ACTIVATED
    origin: WorkflowEventOrigin = WorkflowEventOrigin.FRAMEWORK
    from_step: str
    to_step: str
    data: dict[str, Any]

    def __str__(self) -> str:
        time_str = self.timestamp.strftime("%H:%M:%S")
        return f"[{time_str}] 🔗 {self.from_step} → {self.to_step}"


__all__ = [
    "InputType",
    "OutputType",
    "WorkflowCompletionMode",
    "JoinFailureMode",
    "StepStatus",
    "WorkflowStatus",
    "WorkflowStateMode",
    "WorkflowStateConflictPolicy",
    "WorkflowEventOrigin",
    "WorkflowRequestKind",
    "EdgeCondition",
    "Edge",
    "StepExecution",
    "WorkflowPendingRequest",
    "WorkflowExecution",
    "StepMetadata",
    "WorkflowMetadata",
    "WorkflowContext",
    "_WorkflowRequestSuspended",
    "WorkflowValidationResult",
    "WorkflowEventType",
    "WorkflowEvent",
    "WorkflowErrorDetails",
    "WorkflowStartedEvent",
    "WorkflowCompletedEvent",
    "WorkflowFailedEvent",
    "WorkflowCancelledEvent",
    "WorkflowSuspendedEvent",
    "WorkflowResumedEvent",
    "CheckpointSavedEvent",
    "WorkflowInputRequestedEvent",
    "WorkflowApprovalRequestedEvent",
    "StepStartedEvent",
    "StepCompletedEvent",
    "StepFailedEvent",
    "StepProgressEvent",
    "EdgeActivatedEvent",
]

