from __future__ import annotations

from enum import Enum
from typing import Any

from pydantic import BaseModel, ConfigDict, Field

from nexus.resilience.errors import NexusError


class NodeState(str, Enum):
    PENDING = "pending"
    ACTIVE = "active"
    COMPLETE = "complete"
    FAILED = "failed"
    SKIPPED = "skipped"
    RETRY = "retry"


class NodeConfig(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    node_id: str
    agent_id: str
    task_fn: Any
    task_description: str
    depends_on: list[str] = Field(default_factory=list)
    is_fork_point: bool = False
    merge_strategy: str | None = None
    loop_strategy: str = "max_iterations"
    max_iterations: int = 5
    convergence_threshold: float = 0.95
    enables_graph_mutation: bool = False
    tags: list[str] = Field(default_factory=list)


class AgentNode:
    """Lifecycle wrapper around a configured graph node."""

    def __init__(self, config: NodeConfig) -> None:
        self.config = config
        self._state = NodeState.PENDING
        self._result: Any = None
        self._error: NexusError | None = None
        self._iteration_count = 0

    @property
    def state(self) -> NodeState:
        return self._state

    @property
    def result(self) -> Any:
        return self._result

    @property
    def error(self) -> NexusError | None:
        return self._error

    @property
    def iteration_count(self) -> int:
        return self._iteration_count

    def mark_active(self) -> None:
        """Mark node as active."""
        self._state = NodeState.ACTIVE

    def mark_complete(self, result: Any) -> None:
        """Mark node complete with result."""
        self._state = NodeState.COMPLETE
        self._result = result

    def mark_failed(self, error: NexusError) -> None:
        """Mark node failed with structured error."""
        self._state = NodeState.FAILED
        self._error = error

    def mark_skipped(self) -> None:
        """Mark node skipped."""
        self._state = NodeState.SKIPPED

    def increment_iteration(self) -> int:
        """Increment loop iteration counter and return new value."""
        self._iteration_count += 1
        return self._iteration_count

    def is_ready(self, completed_nodes: set[str]) -> bool:
        """A node is ready when all dependencies are complete."""
        return all(dep in completed_nodes for dep in self.config.depends_on)

    def to_dict(self) -> dict[str, Any]:
        """Serialize node state for checkpointing."""
        return {
            "config": self.config.model_dump(exclude={"task_fn"}),
            "state": self._state.value,
            "result": self._result,
            "error": self._error.to_dict() if self._error else None,
            "iteration_count": self._iteration_count,
        }
