from __future__ import annotations

import asyncio
from dataclasses import dataclass
from typing import Any, Awaitable, Callable

from core.agent.cost_tracker import CostTracker
from core.agent import events

StageCallable = Callable[..., Awaitable[dict[str, Any]]]


@dataclass(frozen=True)
class OrchestratorResult:
    output: str
    usage: dict[str, Any]
    events: list[dict[str, Any]]
    context: dict[str, Any]
    failed: bool
    error: str | None


class AgentChatOrchestrator:
    def __init__(
        self,
        *,
        intent_stage: StageCallable,
        execution_stage: StageCallable,
        organization_stage: StageCallable,
    ) -> None:
        self._intent_stage = intent_stage
        self._execution_stage = execution_stage
        self._organization_stage = organization_stage

    def run_sync(self, *, run_id: str, user_message: str) -> OrchestratorResult:
        return asyncio.run(self.run(run_id=run_id, user_message=user_message))

    async def run(self, *, run_id: str, user_message: str) -> OrchestratorResult:
        tracker = CostTracker()
        emitted_events: list[dict[str, Any]] = [events.run_started(run_id=run_id)]
        context: dict[str, Any] = {}

        stage_pipeline: list[tuple[str, StageCallable]] = [
            ("intent", self._intent_stage),
            ("execution", self._execution_stage),
            ("organization", self._organization_stage),
        ]

        stage_output = user_message
        try:
            for stage_name, stage_callable in stage_pipeline:
                stage_result = await stage_callable(
                    message=stage_output, context=context
                )
                stage_output = str(stage_result.get("content", stage_output))
                usage = stage_result.get("usage", {})
                if isinstance(usage, dict):
                    tracker.add_usage(usage)
                emitted_events.append(
                    events.stage_completed(
                        run_id=run_id,
                        stage=stage_name,
                        usage=tracker.snapshot(),
                    )
                )
        except Exception as exc:  # noqa: BLE001
            emitted_events.append(events.run_failed(run_id=run_id, error=str(exc)))
            return OrchestratorResult(
                output="",
                usage=tracker.snapshot(),
                events=emitted_events,
                context=context,
                failed=True,
                error=str(exc),
            )

        summary = tracker.snapshot()
        emitted_events.append(
            events.run_completed(run_id=run_id, output=stage_output, usage=summary)
        )
        return OrchestratorResult(
            output=stage_output,
            usage=summary,
            events=emitted_events,
            context=context,
            failed=False,
            error=None,
        )
