"""Crew-style process orchestration layered on top of Team."""
from __future__ import annotations

import asyncio
from dataclasses import asdict, dataclass, field
from typing import Any, AsyncIterator, Literal
from uuid import uuid4

from pydantic import BaseModel, Field

from logos_adk.runnable import Runnable
from logos_adk.types import Response, StreamChunk, StreamDone, TextChunk, Usage
from logos_adk.ui_events import make_ui_event


CrewProcessMode = Literal["sequential", "hierarchical"]
CrewTaskType = Literal["work", "approval"]


class CrewRoutingDecisionSchema(BaseModel):
    """Structured manager routing decision for hierarchical crews."""

    role: str | None = Field(default=None, min_length=1)
    agent_name: str | None = Field(default=None, min_length=1)
    instruction: str | None = None


@dataclass
class CrewTask:
    """Explicit task description for a crew process."""

    name: str
    description: str
    task_type: CrewTaskType = "work"
    role: str | None = None
    agent_name: str | None = None
    expected_output: str | None = None
    depends_on: list[str] = field(default_factory=list)
    context: list[str] = field(default_factory=list)
    async_execution: bool = False
    approval_message: str | None = None
    output_key: str | None = None
    handoff_to_role: str | None = None
    handoff_to_agent: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class CrewTaskResult:
    """Execution result of a single crew task."""

    task_name: str
    task_type: CrewTaskType
    assigned_agent: str
    assigned_role: str | None
    prompt: str
    output: str
    status: str
    duration_ms: float
    depends_on: list[str] = field(default_factory=list)
    async_execution: bool = False
    handoff_to_role: str | None = None
    handoff_to_agent: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)


@dataclass
class CrewResult:
    """Serializable result of a crew process run."""

    run_id: str
    crew_name: str
    process: CrewProcessMode
    status: str
    final_output: str
    task_results: list[CrewTaskResult]
    context: dict[str, str]
    handoffs: list[dict[str, Any]]
    manager_role: str | None = None
    manager_agent: str | None = None
    pending_approval: dict[str, Any] | None = None


class CrewProcess(Runnable):
    """Runnable wrapper that executes a crew task list over a Team."""

    def __init__(
        self,
        name: str,
        *,
        team: Any,
        tasks: list[CrewTask],
        process: CrewProcessMode = "sequential",
        manager_role: str | None = None,
        max_parallel_tasks: int | None = None,
        approval_responses: dict[str, bool] | None = None,
        approval_resolver: Any | None = None,
    ) -> None:
        self._name = name
        self._team = team
        self._tasks = tasks
        self._process = process
        self._manager_role = manager_role
        self._max_parallel_tasks = max_parallel_tasks
        self._approval_responses = approval_responses
        self._approval_resolver = approval_resolver
        self._last_paused_run_id: str | None = None

    @property
    def name(self) -> str:
        return self._name

    async def run(self, prompt: str, *, session_id: str | None = None, **kwargs: Any) -> Response:
        result = await self._team.run_tasks(
            self._tasks,
            kickoff=prompt,
            process=self._process,
            manager_role=self._manager_role,
            max_parallel_tasks=self._max_parallel_tasks,
            approval_responses=self._approval_responses,
            approval_resolver=self._approval_resolver,
            session_id=session_id,
            **kwargs,
        )
        self._last_paused_run_id = result.run_id if result.status == "paused" else None
        return Response(
            content=result.final_output,
            tool_calls=[],
            usage=Usage(input_tokens=0, output_tokens=0),
            raw={"crew_result": serialize_crew_result(result)},
        )

    async def resume_approval(
        self,
        task_name: str,
        approved: bool,
        *,
        run_id: str | None = None,
        session_id: str | None = None,
        **kwargs: Any,
    ) -> Response:
        target_run_id = run_id or self._last_paused_run_id
        if target_run_id is None:
            raise ValueError("No paused crew run is available to resume")
        result = await self._team.resume_approval(
            target_run_id,
            task_name=task_name,
            approved=approved,
            session_id=session_id,
            **kwargs,
        )
        self._last_paused_run_id = result.run_id if result.status == "paused" else None
        return Response(
            content=result.final_output,
            tool_calls=[],
            usage=Usage(input_tokens=0, output_tokens=0),
            raw={"crew_result": serialize_crew_result(result)},
        )

    async def stream(self, prompt: str, **kwargs: Any) -> AsyncIterator[StreamChunk]:
        response = await self.run(prompt, session_id=kwargs.pop("session_id", None), **kwargs)
        if response.content:
            yield TextChunk(text=response.content)
        yield StreamDone(response=response)

    async def stream_events(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        correlation_id: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[dict[str, Any]]:
        run_id = uuid4().hex
        sequence = 0
        trace_id = kwargs.get("_trace_id")
        effective_correlation_id = correlation_id or session_id or run_id
        yield make_ui_event(
            event="graph_started",
            event_type="crew",
            graph_name=self.name,
            graph_run_id=run_id,
            sequence=sequence,
            session_id=session_id,
            correlation_id=effective_correlation_id,
            trace_id=trace_id,
            node_name=self.name,
            status="running",
        )
        sequence += 1
        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue()

        async def callback(event: dict[str, Any]) -> None:
            await queue.put(event)

        response_task = asyncio.create_task(
            self._team.run_tasks(
                self._tasks,
                kickoff=prompt,
                process=self._process,
                manager_role=self._manager_role,
                max_parallel_tasks=self._max_parallel_tasks,
                approval_responses=self._approval_responses,
                approval_resolver=self._approval_resolver,
                session_id=session_id,
                _workflow_event_callback=callback,
                correlation_id=effective_correlation_id,
                **kwargs,
            )
        )
        while True:
            if response_task.done() and queue.empty():
                break
            try:
                event = await asyncio.wait_for(queue.get(), timeout=0.05)
            except asyncio.TimeoutError:
                continue
            yield event
        response = await response_task
        self._last_paused_run_id = response.run_id if response.status == "paused" else None
        yield make_ui_event(
            event="graph_completed",
            event_type="crew",
            graph_name=self.name,
            graph_run_id=run_id,
            sequence=sequence,
            session_id=session_id,
            correlation_id=effective_correlation_id,
            trace_id=trace_id,
            node_name=self.name,
            status=response.status,
            output=response.final_output,
        )


def serialize_crew_result(result: CrewResult) -> dict[str, Any]:
    return {
        "run_id": result.run_id,
        "crew_name": result.crew_name,
        "process": result.process,
        "status": result.status,
        "final_output": result.final_output,
        "task_results": [asdict(task_result) for task_result in result.task_results],
        "context": dict(result.context),
        "handoffs": list(result.handoffs),
        "manager_role": result.manager_role,
        "manager_agent": result.manager_agent,
        "pending_approval": result.pending_approval,
    }


__all__ = [
    "CrewProcess",
    "CrewProcessMode",
    "CrewResult",
    "CrewRoutingDecisionSchema",
    "CrewTask",
    "CrewTaskResult",
    "CrewTaskType",
    "serialize_crew_result",
]
