from dataclasses import dataclass
from enum import StrEnum
from typing import Any

ROUTE_CONTEXT_KEY = "_turn_route"


class RouteKind(StrEnum):
    DIRECT_AGENT = "direct_agent"
    TOOLS_ORCHESTRATOR = "tools_orchestrator"
    LEGACY_MULTI_AGENT = "legacy_multi_agent"


@dataclass(frozen=True)
class TurnRoute:
    kind: RouteKind
    model_id: str | None
    agent_ids: tuple[str, ...]
    doc_gen_enabled: bool

    def to_dict(self) -> dict[str, Any]:
        return {
            "kind": self.kind.value,
            "model_id": self.model_id,
            "agent_ids": list(self.agent_ids),
            "doc_gen_enabled": self.doc_gen_enabled,
        }


def route_from_context(context: dict | None) -> TurnRoute | None:
    """Restore a route saved on a pending assistant message."""
    data = (context or {}).get(ROUTE_CONTEXT_KEY)
    if not isinstance(data, dict):
        return None
    try:
        return TurnRoute(
            kind=RouteKind(data["kind"]),
            model_id=data.get("model_id"),
            agent_ids=tuple(data.get("agent_ids") or []),
            doc_gen_enabled=bool(data.get("doc_gen_enabled")),
        )
    except (KeyError, TypeError, ValueError):
        return None


def history_for_route(
    route: TurnRoute,
    history: list[dict],
    history_by_target: dict[str, list[dict]] | None,
    selected_ids: list[str],
) -> list[dict]:
    """Return the attachment-safe history prepared for an orchestrated route."""
    if not history_by_target:
        return history
    target_ids = list(route.agent_ids) or selected_ids or ["__base_llm__"]
    for target_id in target_ids:
        if target_history := history_by_target.get(target_id):
            return target_history
    return history_by_target.get("__base_llm__") or history


def resolve_turn_route(
    *,
    answers_mode: bool,
    draft_mode: bool,
    doc_gen_enabled: bool,
    target_id: str | None,
    target_is_agent: bool,
    agent_ids: list[str],
    agents_as_tools: bool,
    base_model_id: str,
    orchestrator_model_id: str | None,
) -> TurnRoute:
    """Choose one runtime without executing or loading any agents."""
    if answers_mode:
        if target_is_agent:
            if not target_id:
                raise ValueError("An agent target requires an id")
            if not doc_gen_enabled or not orchestrator_model_id:
                return TurnRoute(RouteKind.DIRECT_AGENT, None, (target_id,), False)
            return TurnRoute(
                RouteKind.TOOLS_ORCHESTRATOR,
                orchestrator_model_id,
                (target_id,),
                True,
            )
        return TurnRoute(
            RouteKind.TOOLS_ORCHESTRATOR,
            target_id or base_model_id,
            (),
            doc_gen_enabled,
        )

    selected_ids = tuple(agent_ids)
    if len(selected_ids) == 1 and (draft_mode or not doc_gen_enabled):
        return TurnRoute(RouteKind.DIRECT_AGENT, None, selected_ids, False)
    if len(selected_ids) > 1 and not agents_as_tools:
        return TurnRoute(RouteKind.LEGACY_MULTI_AGENT, base_model_id, selected_ids, doc_gen_enabled)
    return TurnRoute(RouteKind.TOOLS_ORCHESTRATOR, base_model_id, selected_ids, doc_gen_enabled)
