"""Core types for workflow pipelines."""
from __future__ import annotations

from dataclasses import asdict, dataclass, field
from datetime import datetime, UTC
from typing import Any, Iterator, MutableMapping

from logos_adk.state import AgentState


class BlackboardConflictError(RuntimeError):
    """Raised when concurrent blackboard writes conflict under optimistic merge."""


class WorkflowBlackboard(MutableMapping[str, Any]):
    """Versioned blackboard with optimistic concurrency metadata."""

    def __init__(
        self,
        initial: dict[str, Any] | None = None,
        *,
        versions: dict[str, int] | None = None,
    ) -> None:
        self._data: dict[str, Any] = dict(initial or {})
        if versions is None:
            self._versions = {key: 1 for key in self._data}
        else:
            self._versions = {
                str(key): max(0, int(value))
                for key, value in dict(versions).items()
            }
            for key in self._data:
                self._versions.setdefault(key, 1)
        self._dirty_keys: set[str] = set()
        self._deleted_keys: set[str] = set()

    def __getitem__(self, key: str) -> Any:
        return self._data[key]

    def __setitem__(self, key: str, value: Any) -> None:
        self._data[key] = value
        self._versions[key] = self.version(key) + 1
        self._dirty_keys.add(key)
        self._deleted_keys.discard(key)

    def __delitem__(self, key: str) -> None:
        if key not in self._data:
            raise KeyError(key)
        del self._data[key]
        self._versions[key] = self.version(key) + 1
        self._dirty_keys.add(key)
        self._deleted_keys.add(key)

    def __iter__(self) -> Iterator[str]:
        return iter(self._data)

    def __len__(self) -> int:
        return len(self._data)

    def clone(self) -> WorkflowBlackboard:
        return WorkflowBlackboard(self._data, versions=self._versions)

    def snapshot(self) -> dict[str, Any]:
        return dict(self._data)

    def versions_snapshot(self) -> dict[str, int]:
        return dict(self._versions)

    def version(self, key: str) -> int:
        return int(self._versions.get(key, 0))

    def compare_and_set(self, key: str, value: Any, *, expected_version: int) -> int:
        current_version = self.version(key)
        if current_version != int(expected_version):
            raise BlackboardConflictError(
                f"Blackboard key '{key}' version conflict: expected {expected_version}, got {current_version}"
            )
        self[key] = value
        return self.version(key)

    def delete_if_version(self, key: str, *, expected_version: int) -> int:
        current_version = self.version(key)
        if current_version != int(expected_version):
            raise BlackboardConflictError(
                f"Blackboard key '{key}' version conflict: expected {expected_version}, got {current_version}"
            )
        del self[key]
        return self.version(key)

    def dirty_writes(self) -> dict[str, Any]:
        return {
            key: self._data[key]
            for key in self._dirty_keys
            if key not in self._deleted_keys and key in self._data
        }

    def dirty_deletes(self) -> set[str]:
        return set(self._deleted_keys)

    def reset_tracking(self) -> None:
        self._dirty_keys.clear()
        self._deleted_keys.clear()


@dataclass
class PipelineContext:
    """Carries data between pipeline steps."""

    pipeline_id: str
    data: dict[str, Any] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)
    blackboard: WorkflowBlackboard | dict[str, Any] = field(default_factory=WorkflowBlackboard)

    def __post_init__(self) -> None:
        if isinstance(self.blackboard, WorkflowBlackboard):
            return
        versions = self.metadata.pop("_blackboard_versions", None)
        self.blackboard = WorkflowBlackboard(
            self.blackboard,
            versions=versions if isinstance(versions, dict) else None,
        )

    def __getitem__(self, key: str) -> Any:
        return self.data[key]

    def __setitem__(self, key: str, value: Any) -> None:
        self.data[key] = value

    def get(self, key: str, default: Any = None) -> Any:
        return self.data.get(key, default)


@dataclass
class StepResult:
    """Result of executing a single step."""

    output: Any
    status: str  # "completed", "failed", "skipped", "paused"
    duration_ms: float
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class PipelineResult:
    """Result of executing an entire pipeline."""

    status: str  # "completed", "failed", "paused"
    outputs: dict[str, StepResult]
    context: PipelineContext
    duration_ms: float


@dataclass
class PipelineState:
    """Serializable snapshot of pipeline execution progress."""

    pipeline_id: str
    pipeline_name: str
    current_step_index: int
    status: str  # "running", "paused", "completed", "failed"
    context_data: dict[str, Any] = field(default_factory=dict)
    metadata_data: dict[str, Any] = field(default_factory=dict)
    blackboard_data: dict[str, Any] = field(default_factory=dict)
    blackboard_versions_data: dict[str, int] = field(default_factory=dict)
    step_results: dict[str, dict] = field(default_factory=dict)
    created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    updated_at: datetime = field(default_factory=lambda: datetime.now(UTC))

    def to_agent_state(self) -> AgentState:
        """Wrap into AgentState for StateStore compatibility."""
        return AgentState(
            session_id=self.pipeline_id,
            metadata={"pipeline_state": asdict(self)},
        )

    @classmethod
    def from_agent_state(cls, state: AgentState) -> PipelineState:
        """Restore from an AgentState loaded from StateStore."""
        data = state.metadata["pipeline_state"]

        def _parse_dt(val: Any) -> datetime:
            if isinstance(val, datetime):
                return val
            if isinstance(val, str):
                return datetime.fromisoformat(val)
            return datetime.now(UTC)

        return cls(
            pipeline_id=data["pipeline_id"],
            pipeline_name=data["pipeline_name"],
            current_step_index=data["current_step_index"],
            status=data["status"],
            context_data=data.get("context_data", {}),
            metadata_data=data.get("metadata_data", {}),
            blackboard_data=data.get("blackboard_data", {}),
            blackboard_versions_data=data.get("blackboard_versions_data", {}),
            step_results=data.get("step_results", {}),
            created_at=_parse_dt(data.get("created_at")),
            updated_at=_parse_dt(data.get("updated_at")),
        )
