"""CheckpointMixin — serialization, persistence, and restoration of paused crew runs."""
from __future__ import annotations

from datetime import UTC, datetime
from typing import Any, cast

from logos_adk.crew import CrewProcessMode, CrewTask, CrewTaskResult
from logos_adk.crew_checkpoint_store import (
    CrewCheckpointRecord,
    JSONObject,
    ensure_json_object,
)

from logos_adk.team._types import (
    _CrewTaskCheckpointPayload,
    _CrewTaskResultCheckpointPayload,
    _PausedCrewRun,
    _PausedCrewRunCheckpointPayload,
)


class CheckpointMixin:
    """Checkpoint serialization and persistence for paused crew runs."""

    def _store_paused_crew_run(
        self,
        *,
        run_id: str,
        process: CrewProcessMode,
        tasks: list[CrewTask],
        kickoff: str,
        context: dict[str, str],
        task_results: list[CrewTaskResult],
        task_results_by_name: dict[str, CrewTaskResult],
        handoffs: list[dict[str, Any]],
        pending: dict[str, CrewTask],
        completed: set[str],
        pending_approval: dict[str, Any],
        max_parallel_tasks: int | None,
        session_id: str | None,
        kwargs: dict[str, Any],
        manager_role: str | None = None,
        carried_role: str | None = None,
        carried_agent: str | None = None,
        last_role: str | None = None,
        last_agent: str | None = None,
    ) -> None:
        self._paused_crew_runs[run_id] = _PausedCrewRun(
            run_id=run_id,
            process=process,
            tasks=list(tasks),
            kickoff=kickoff,
            context=dict(context),
            task_results=list(task_results),
            task_results_by_name=dict(task_results_by_name),
            handoffs=[dict(item) for item in handoffs],
            pending=dict(pending),
            completed=set(completed),
            pending_approval=dict(pending_approval),
            max_parallel_tasks=max_parallel_tasks,
            session_id=session_id,
            kwargs=dict(kwargs),
            manager_role=manager_role,
            carried_role=carried_role,
            carried_agent=carried_agent,
            last_role=last_role,
            last_agent=last_agent,
        )
        if self._checkpoint_store is not None:
            now = datetime.now(UTC).isoformat()
            self._checkpoint_store.save_checkpoint(
                CrewCheckpointRecord(
                    run_id=run_id,
                    team_name=self.name,
                    process=process,
                    status="paused",
                    created_at=now,
                    updated_at=now,
                    payload=self._serialize_paused_crew_run(self._paused_crew_runs[run_id]),
                )
            )

    def _load_paused_crew_run(self, run_id: str) -> _PausedCrewRun | None:
        paused_run = self._paused_crew_runs.get(run_id)
        if paused_run is not None:
            return paused_run
        if self._checkpoint_store is None:
            return None
        record = self._checkpoint_store.load_checkpoint(run_id)
        if record is None:
            return None
        paused_run = self._deserialize_paused_crew_run(record.payload)
        self._paused_crew_runs[run_id] = paused_run
        return paused_run

    def _clear_paused_crew_run(self, run_id: str) -> None:
        self._paused_crew_runs.pop(run_id, None)
        if self._checkpoint_store is not None:
            self._checkpoint_store.delete_checkpoint(run_id)

    def _serialize_paused_crew_run(
        self,
        paused_run: _PausedCrewRun,
    ) -> _PausedCrewRunCheckpointPayload:
        payload: _PausedCrewRunCheckpointPayload = {
            "run_id": paused_run.run_id,
            "process": paused_run.process,
            "tasks": [self._serialize_crew_task(task) for task in paused_run.tasks],
            "kickoff": paused_run.kickoff,
            "context": dict(paused_run.context),
            "task_results": [self._serialize_crew_task_result(item) for item in paused_run.task_results],
            "task_results_by_name": {
                name: self._serialize_crew_task_result(item)
                for name, item in paused_run.task_results_by_name.items()
            },
            "handoffs": list(paused_run.handoffs),
            "pending": {
                name: self._serialize_crew_task(task)
                for name, task in paused_run.pending.items()
            },
            "completed": sorted(paused_run.completed),
            "pending_approval": dict(paused_run.pending_approval),
            "max_parallel_tasks": paused_run.max_parallel_tasks,
            "session_id": paused_run.session_id,
            "kwargs": dict(paused_run.kwargs or {}),
            "manager_role": paused_run.manager_role,
            "carried_role": paused_run.carried_role,
            "carried_agent": paused_run.carried_agent,
            "last_role": paused_run.last_role,
            "last_agent": paused_run.last_agent,
        }
        return cast(
            _PausedCrewRunCheckpointPayload,
            ensure_json_object(payload, path="crew_checkpoint"),
        )

    def _deserialize_paused_crew_run(
        self,
        payload: _PausedCrewRunCheckpointPayload | JSONObject,
    ) -> _PausedCrewRun:
        checkpoint_payload = cast(
            _PausedCrewRunCheckpointPayload,
            ensure_json_object(payload, path="crew_checkpoint"),
        )
        return _PausedCrewRun(
            run_id=checkpoint_payload["run_id"],
            process=checkpoint_payload["process"],
            tasks=[
                self._deserialize_crew_task(item)
                for item in checkpoint_payload.get("tasks", [])
            ],
            kickoff=checkpoint_payload.get("kickoff", ""),
            context=dict(checkpoint_payload.get("context", {})),
            task_results=[
                self._deserialize_crew_task_result(item)
                for item in checkpoint_payload.get("task_results", [])
            ],
            task_results_by_name={
                name: self._deserialize_crew_task_result(item)
                for name, item in checkpoint_payload.get(
                    "task_results_by_name",
                    {},
                ).items()
            },
            handoffs=[dict(item) for item in checkpoint_payload.get("handoffs", [])],
            pending={
                name: self._deserialize_crew_task(item)
                for name, item in checkpoint_payload.get("pending", {}).items()
            },
            completed=set(checkpoint_payload.get("completed", [])),
            pending_approval=dict(checkpoint_payload.get("pending_approval", {})),
            max_parallel_tasks=checkpoint_payload.get("max_parallel_tasks"),
            session_id=checkpoint_payload.get("session_id"),
            kwargs=dict(checkpoint_payload.get("kwargs", {})),
            manager_role=checkpoint_payload.get("manager_role"),
            carried_role=checkpoint_payload.get("carried_role"),
            carried_agent=checkpoint_payload.get("carried_agent"),
            last_role=checkpoint_payload.get("last_role"),
            last_agent=checkpoint_payload.get("last_agent"),
        )

    def _serialize_crew_task(self, task: CrewTask) -> _CrewTaskCheckpointPayload:
        payload: _CrewTaskCheckpointPayload = {
            "name": task.name,
            "description": task.description,
            "task_type": task.task_type,
            "role": task.role,
            "agent_name": task.agent_name,
            "expected_output": task.expected_output,
            "depends_on": list(task.depends_on),
            "context": list(task.context),
            "async_execution": task.async_execution,
            "approval_message": task.approval_message,
            "output_key": task.output_key,
            "handoff_to_role": task.handoff_to_role,
            "handoff_to_agent": task.handoff_to_agent,
            "metadata": dict(task.metadata),
        }
        return cast(
            _CrewTaskCheckpointPayload,
            ensure_json_object(payload, path=f"crew_task[{task.name}]"),
        )

    def _deserialize_crew_task(
        self,
        payload: _CrewTaskCheckpointPayload | JSONObject,
    ) -> CrewTask:
        task_payload = cast(
            _CrewTaskCheckpointPayload,
            ensure_json_object(payload, path="crew_task"),
        )
        return CrewTask(
            name=task_payload["name"],
            description=task_payload["description"],
            task_type=task_payload.get("task_type", "work"),
            role=task_payload.get("role"),
            agent_name=task_payload.get("agent_name"),
            expected_output=task_payload.get("expected_output"),
            depends_on=list(task_payload.get("depends_on", [])),
            context=list(task_payload.get("context", [])),
            async_execution=bool(task_payload.get("async_execution", False)),
            approval_message=task_payload.get("approval_message"),
            output_key=task_payload.get("output_key"),
            handoff_to_role=task_payload.get("handoff_to_role"),
            handoff_to_agent=task_payload.get("handoff_to_agent"),
            metadata=dict(task_payload.get("metadata", {})),
        )

    def _serialize_crew_task_result(
        self,
        result: CrewTaskResult,
    ) -> _CrewTaskResultCheckpointPayload:
        payload: _CrewTaskResultCheckpointPayload = {
            "task_name": result.task_name,
            "task_type": result.task_type,
            "assigned_agent": result.assigned_agent,
            "assigned_role": result.assigned_role,
            "prompt": result.prompt,
            "output": result.output,
            "status": result.status,
            "duration_ms": result.duration_ms,
            "depends_on": list(result.depends_on),
            "async_execution": result.async_execution,
            "handoff_to_role": result.handoff_to_role,
            "handoff_to_agent": result.handoff_to_agent,
            "metadata": dict(result.metadata),
        }
        return cast(
            _CrewTaskResultCheckpointPayload,
            ensure_json_object(
                payload,
                path=f"crew_task_result[{result.task_name}]",
            ),
        )

    def _deserialize_crew_task_result(
        self,
        payload: _CrewTaskResultCheckpointPayload | JSONObject,
    ) -> CrewTaskResult:
        result_payload = cast(
            _CrewTaskResultCheckpointPayload,
            ensure_json_object(payload, path="crew_task_result"),
        )
        return CrewTaskResult(
            task_name=result_payload["task_name"],
            task_type=result_payload.get("task_type", "work"),
            assigned_agent=result_payload["assigned_agent"],
            assigned_role=result_payload.get("assigned_role"),
            prompt=result_payload.get("prompt", ""),
            output=result_payload.get("output", ""),
            status=result_payload.get("status", "completed"),
            duration_ms=float(result_payload.get("duration_ms", 0)),
            depends_on=list(result_payload.get("depends_on", [])),
            async_execution=bool(result_payload.get("async_execution", False)),
            handoff_to_role=result_payload.get("handoff_to_role"),
            handoff_to_agent=result_payload.get("handoff_to_agent"),
            metadata=dict(result_payload.get("metadata", {})),
        )
