from __future__ import annotations

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, ClassVar, Literal
from uuid import uuid4

from core.models import utc_now

WorkflowRunStatus = Literal["pending", "running", "waiting", "completed", "failed", "cancelled"]


@dataclass(frozen=True)
class ConditionSpec:
    kind: str
    params: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        return {"kind": self.kind, "params": dict(self.params)}

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "ConditionSpec":
        return cls(kind=str(data["kind"]), params=dict(data.get("params") or {}))


@dataclass(frozen=True)
class PlannedGoalStep:
    step_id: str
    goal: str
    next_step_id: str | None = None
    kind: ClassVar[str] = "planned_goal"

    def to_dict(self) -> dict[str, Any]:
        return {
            "kind": self.kind,
            "step_id": self.step_id,
            "goal": self.goal,
            "next_step_id": self.next_step_id,
        }


@dataclass(frozen=True)
class SkillInvocationStep:
    step_id: str
    skill_id: str
    args: dict[str, Any] = field(default_factory=dict)
    next_step_id: str | None = None
    kind: ClassVar[str] = "skill_invocation"

    def to_dict(self) -> dict[str, Any]:
        return {
            "kind": self.kind,
            "step_id": self.step_id,
            "skill_id": self.skill_id,
            "args": dict(self.args),
            "next_step_id": self.next_step_id,
        }


@dataclass(frozen=True)
class WaitForConditionStep:
    step_id: str
    condition: ConditionSpec
    timeout_seconds: float
    poll_interval_seconds: float
    next_step_id: str | None = None
    kind: ClassVar[str] = "wait_for_condition"

    def to_dict(self) -> dict[str, Any]:
        return {
            "kind": self.kind,
            "step_id": self.step_id,
            "condition": self.condition.to_dict(),
            "timeout_seconds": self.timeout_seconds,
            "poll_interval_seconds": self.poll_interval_seconds,
            "next_step_id": self.next_step_id,
        }


@dataclass(frozen=True)
class BranchStep:
    step_id: str
    condition: ConditionSpec
    on_true: str
    on_false: str
    kind: ClassVar[str] = "branch"

    def to_dict(self) -> dict[str, Any]:
        return {
            "kind": self.kind,
            "step_id": self.step_id,
            "condition": self.condition.to_dict(),
            "on_true": self.on_true,
            "on_false": self.on_false,
        }


WorkflowStep = PlannedGoalStep | SkillInvocationStep | WaitForConditionStep | BranchStep


@dataclass(frozen=True)
class WorkflowDefinition:
    name: str
    steps: list[WorkflowStep]
    entry_step_id: str
    id: str = field(default_factory=lambda: uuid4().hex)

    def __post_init__(self) -> None:
        if not self.steps:
            raise ValueError("workflow definition requires at least one step")
        step_ids = [step.step_id for step in self.steps]
        if len(step_ids) != len(set(step_ids)):
            raise ValueError("workflow definition contains duplicate step_id values")
        known = set(step_ids)
        if self.entry_step_id not in known:
            raise ValueError(f"entry_step_id {self.entry_step_id} is not a workflow step")
        for step in self.steps:
            for target in _step_targets(step):
                if target is not None and target not in known:
                    raise ValueError(
                        f"step {step.step_id} references unknown step {target}"
                    )

    def step_by_id(self, step_id: str) -> WorkflowStep:
        for step in self.steps:
            if step.step_id == step_id:
                return step
        raise KeyError(step_id)

    def next_step_id_after(self, step_id: str) -> str | None:
        for index, step in enumerate(self.steps):
            if step.step_id == step_id:
                if index + 1 >= len(self.steps):
                    return None
                return self.steps[index + 1].step_id
        raise KeyError(step_id)

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "name": self.name,
            "entry_step_id": self.entry_step_id,
            "steps": [step.to_dict() for step in self.steps],
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "WorkflowDefinition":
        return cls(
            id=str(data["id"]),
            name=str(data["name"]),
            entry_step_id=str(data["entry_step_id"]),
            steps=[workflow_step_from_dict(step) for step in data.get("steps", [])],
        )


@dataclass(frozen=True)
class WorkflowStepResult:
    step_id: str
    kind: str
    success: bool
    detail: dict[str, Any] = field(default_factory=dict)
    task_id: str | None = None
    timestamp: datetime = field(default_factory=utc_now)

    def to_dict(self) -> dict[str, Any]:
        return {
            "step_id": self.step_id,
            "kind": self.kind,
            "success": self.success,
            "detail": dict(self.detail),
            "task_id": self.task_id,
            "timestamp": self.timestamp.isoformat(),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "WorkflowStepResult":
        return cls(
            step_id=str(data["step_id"]),
            kind=str(data["kind"]),
            success=bool(data["success"]),
            detail=dict(data.get("detail") or {}),
            task_id=data.get("task_id"),
            timestamp=_parse_datetime(data.get("timestamp")),
        )


@dataclass(frozen=True)
class WorkflowRun:
    definition_id: str
    status: WorkflowRunStatus
    current_step_id: str | None
    variables: dict[str, Any] = field(default_factory=dict)
    step_results: list[WorkflowStepResult] = field(default_factory=list)
    id: str = field(default_factory=lambda: uuid4().hex)
    device_id: str | None = None
    created_at: datetime = field(default_factory=utc_now)
    updated_at: datetime = field(default_factory=utc_now)

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "definition_id": self.definition_id,
            "status": self.status,
            "current_step_id": self.current_step_id,
            "device_id": self.device_id,
            "variables": dict(self.variables),
            "step_results": [result.to_dict() for result in self.step_results],
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
        }

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "WorkflowRun":
        return cls(
            id=str(data["id"]),
            definition_id=str(data["definition_id"]),
            status=data["status"],
            current_step_id=data.get("current_step_id"),
            device_id=data.get("device_id"),
            variables=dict(data.get("variables") or {}),
            step_results=[
                WorkflowStepResult.from_dict(result)
                for result in data.get("step_results", [])
            ],
            created_at=_parse_datetime(data.get("created_at")),
            updated_at=_parse_datetime(data.get("updated_at")),
        )


def workflow_step_from_dict(data: dict[str, Any]) -> WorkflowStep:
    kind = data.get("kind")
    if kind == PlannedGoalStep.kind:
        return PlannedGoalStep(
            step_id=str(data["step_id"]),
            goal=str(data["goal"]),
            next_step_id=data.get("next_step_id"),
        )
    if kind == SkillInvocationStep.kind:
        return SkillInvocationStep(
            step_id=str(data["step_id"]),
            skill_id=str(data["skill_id"]),
            args=dict(data.get("args") or {}),
            next_step_id=data.get("next_step_id"),
        )
    if kind == WaitForConditionStep.kind:
        return WaitForConditionStep(
            step_id=str(data["step_id"]),
            condition=ConditionSpec.from_dict(data["condition"]),
            timeout_seconds=float(data["timeout_seconds"]),
            poll_interval_seconds=float(data["poll_interval_seconds"]),
            next_step_id=data.get("next_step_id"),
        )
    if kind == BranchStep.kind:
        return BranchStep(
            step_id=str(data["step_id"]),
            condition=ConditionSpec.from_dict(data["condition"]),
            on_true=str(data["on_true"]),
            on_false=str(data["on_false"]),
        )
    raise ValueError(f"unknown workflow step kind {kind}")


def _step_targets(step: WorkflowStep) -> list[str | None]:
    if isinstance(step, BranchStep):
        return [step.on_true, step.on_false]
    return [step.next_step_id]


def _parse_datetime(value: Any) -> datetime:
    if isinstance(value, datetime):
        return value
    if isinstance(value, str):
        try:
            return datetime.fromisoformat(value)
        except ValueError:
            pass
    return utc_now()