"""Dynamic department-level orchestration across teams, crews, and swarms."""
from __future__ import annotations

import importlib
import json
import re
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Literal
from uuid import uuid4

from pydantic import BaseModel, Field

from logos_adk.crew import CrewProcess
from logos_adk.dynamic_orchestration_config import DepartmentSupervisorConfig
from logos_adk.errors import ConfigError
from logos_adk.event_bus import InMemoryEventBus
from logos_adk.runnable import Runnable
from logos_adk.subscription_registry import subscription_handler_registry
from logos_adk.swarm import Swarm
from logos_adk.team import Team
from logos_adk.types import Response, StreamChunk, StreamDone, StructuredResponse, TextChunk
from logos_adk.ui_events import make_ui_event


DepartmentKind = Literal["team", "crew", "swarm", "runnable"]


class DepartmentRoutingDecisionSchema(BaseModel):
    """Structured supervisor decision for department selection."""

    department_name: str = Field(min_length=1)
    delegated_prompt: str | None = None
    reason: str | None = None


@dataclass
class DepartmentSpec:
    """Descriptor for one executable department under the supervisor."""

    name: str
    target: Runnable | Swarm
    description: str = ""
    kind: DepartmentKind = "runnable"
    capability_tags: list[str] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)
    request_topic: str | None = None
    response_topic: str | None = None


class DepartmentSupervisor(Runnable):
    """LLM-guided supervisor that routes goals to the most suitable department."""

    def __init__(
        self,
        name: str,
        *,
        router: Any | None = None,
        departments: list[DepartmentSpec] | None = None,
        fallback_department: str | None = None,
    ) -> None:
        self._name = name
        self._router = router
        self._fallback_department = fallback_department
        self._departments: dict[str, DepartmentSpec] = {}
        for department in departments or []:
            self.add_department_spec(department)

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

    def department(
        self,
        name: str,
        target: Runnable | Swarm,
        *,
        description: str = "",
        kind: DepartmentKind | None = None,
        capability_tags: list[str] | None = None,
        metadata: dict[str, Any] | None = None,
        request_topic: str | None = None,
        response_topic: str | None = None,
    ) -> DepartmentSupervisor:
        return self.add_department_spec(
            DepartmentSpec(
                name=name,
                target=target,
                description=description,
                kind=kind or _infer_department_kind(target),
                capability_tags=list(capability_tags or []),
                metadata=dict(metadata or {}),
                request_topic=request_topic,
                response_topic=response_topic,
            )
        )

    def add_department_spec(self, department: DepartmentSpec) -> DepartmentSupervisor:
        if department.name in self._departments:
            raise ValueError(f"Department '{department.name}' is already registered")
        self._departments[department.name] = department
        if self._fallback_department is None:
            self._fallback_department = department.name
        return self

    def list_departments(self) -> list[DepartmentSpec]:
        return list(self._departments.values())

    @classmethod
    def from_definition(cls, config: DepartmentSupervisorConfig) -> DepartmentSupervisor:
        """Build a DepartmentSupervisor from a typed config."""
        from logos_adk.agent import Agent
        from logos_adk.config_loader import build_agent_from_definition

        router = (
            build_agent_from_definition(config.router, default_cls=Agent)
            if config.router is not None
            else None
        )
        supervisor = cls(
            config.name,
            router=router,
            fallback_department=config.fallback_department,
        )
        for department in config.departments:
            if department.kind == "team":
                if department.team is None:
                    raise ConfigError(f"Department '{department.name}' requires a team config")
                target = Team.from_definition(department.team)
            elif department.kind == "crew":
                if department.team is None:
                    raise ConfigError(f"Department '{department.name}' requires a team config")
                team = Team.from_definition(department.team)
                target = CrewProcess(
                    department.name,
                    team=team,
                    tasks=department.tasks,
                    process=department.process,
                    manager_role=department.manager_role,
                    max_parallel_tasks=department.max_parallel_tasks,
                )
            elif department.kind == "swarm":
                target = Swarm(
                    department.name,
                    event_bus=Team._build_event_bus_from_config(department.event_bus) or InMemoryEventBus(),
                    retry_delay_seconds=department.retry_delay_seconds,
                    lease_seconds=department.lease_seconds,
                )
                for subscription in department.subscriptions:
                    target.subscribe(
                        subscription.topic,
                        _resolve_swarm_handler(subscription.handler),
                        name=subscription.name,
                        publish_to=subscription.publish_to,
                    )
            else:
                raise ConfigError(f"Unsupported department kind '{department.kind}'")
            supervisor.department(
                department.name,
                target,
                description=department.description,
                kind=department.kind,
                capability_tags=department.capability_tags,
                metadata=department.metadata,
                request_topic=department.request_topic,
                response_topic=department.response_topic,
            )
        return supervisor

    @classmethod
    def from_config(cls, path: str) -> DepartmentSupervisor:
        """Load a DepartmentSupervisor from a YAML config file."""
        from logos_adk.dynamic_orchestration_config_parser import load_department_supervisor_config

        return cls.from_definition(load_department_supervisor_config(path))

    async def run(
        self,
        prompt: str,
        *,
        session_id: str | None = None,
        **kwargs: Any,
    ) -> Response:
        department, decision = await self._select_department(prompt, session_id=session_id)
        delegated_prompt = (decision.delegated_prompt or prompt).strip() or prompt
        response = await self._run_department(
            department,
            delegated_prompt,
            session_id=session_id,
            **kwargs,
        )
        raw = dict(response.raw or {})
        raw["dynamic_orchestration"] = {
            "supervisor": self._name,
            "department_name": department.name,
            "department_kind": department.kind,
            "delegated_prompt": delegated_prompt,
            "reason": decision.reason,
        }
        return Response(
            content=response.content,
            content_parts=response.content_parts,
            tool_calls=response.tool_calls,
            usage=response.usage,
            raw=raw,
            reasoning=response.reasoning,
            status=response.status,
        )

    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, reasoning=response.reasoning)
        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="supervisor",
            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
        yield make_ui_event(
            event="node_started",
            event_type="routing",
            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="router",
            status="running",
            metadata={
                "user_facing": {
                    "kind": "step",
                    "title": "Выбираем маршрут",
                    "detail": "Определяем подходящий отдел для выполнения задачи",
                    "step_key": "route:select",
                    "status": "running",
                    "visible": True,
                }
            },
        )
        sequence += 1
        department, decision = await self._select_department(prompt, session_id=session_id)
        delegated_prompt = (decision.delegated_prompt or prompt).strip() or prompt
        yield make_ui_event(
            event="node_completed",
            event_type="routing",
            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="router",
            status="completed",
            output={
                "department_name": department.name,
                "department_kind": department.kind,
                "delegated_prompt": delegated_prompt,
                "reason": decision.reason,
            },
            metadata={
                "user_facing": {
                    "kind": "step",
                    "title": "Маршрут выбран",
                    "detail": f"Передаём задачу в {department.name}",
                    "step_key": "route:select",
                    "status": "completed",
                    "visible": True,
                }
            },
        )
        sequence += 1
        target = department.target
        if isinstance(target, Swarm):
            final_event: dict[str, Any] | None = None
            async for event in target.dispatch_events(
                delegated_prompt,
                request_topic=department.request_topic or department.name,
                response_topic=department.response_topic,
                session_id=session_id,
                correlation_id=effective_correlation_id,
                trace_id=trace_id,
            ):
                payload = _decorate_forwarded_event(event, department_name=department.name)
                if payload is None:
                    continue
                if payload.get("event") == "graph_completed":
                    final_event = payload
                    continue
                yield payload
            yield _build_supervisor_terminal_event(
                supervisor_name=self._name,
                run_id=run_id,
                sequence=sequence,
                session_id=session_id,
                correlation_id=effective_correlation_id,
                trace_id=(final_event or {}).get("trace_id") or trace_id,
                department=department,
                decision=decision,
                delegated_prompt=delegated_prompt,
                final_event=final_event,
            )
            return
        if hasattr(target, "stream_ui_events"):
            final_event = None
            async for event in target.stream_ui_events(
                delegated_prompt,
                session_id=session_id,
                correlation_id=effective_correlation_id,
                **kwargs,
            ):
                payload = _decorate_forwarded_event(event, department_name=department.name)
                if payload is None:
                    continue
                if payload.get("event") == "graph_completed":
                    final_event = payload
                    continue
                yield payload
            yield _build_supervisor_terminal_event(
                supervisor_name=self._name,
                run_id=run_id,
                sequence=sequence,
                session_id=session_id,
                correlation_id=effective_correlation_id,
                trace_id=(final_event or {}).get("trace_id") or trace_id,
                department=department,
                decision=decision,
                delegated_prompt=delegated_prompt,
                final_event=final_event,
            )
            return
        if hasattr(target, "stream_events"):
            final_event = None
            async for event in target.stream_events(
                delegated_prompt,
                session_id=session_id,
                correlation_id=effective_correlation_id,
                **kwargs,
            ):
                payload = _decorate_forwarded_event(event, department_name=department.name)
                if payload is None:
                    continue
                if payload.get("event") == "graph_completed":
                    final_event = payload
                    continue
                yield payload
            yield _build_supervisor_terminal_event(
                supervisor_name=self._name,
                run_id=run_id,
                sequence=sequence,
                session_id=session_id,
                correlation_id=effective_correlation_id,
                trace_id=(final_event or {}).get("trace_id") or trace_id,
                department=department,
                decision=decision,
                delegated_prompt=delegated_prompt,
                final_event=final_event,
            )
            return
        response = await self._run_department(
            department,
            delegated_prompt,
            session_id=session_id,
            **kwargs,
        )
        yield make_ui_event(
            event="graph_completed",
            event_type="supervisor",
            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=department.name,
            status=response.status,
            output=response,
        )

    async def _select_department(
        self,
        prompt: str,
        *,
        session_id: str | None,
    ) -> tuple[DepartmentSpec, DepartmentRoutingDecisionSchema]:
        if not self._departments:
            raise ValueError("DepartmentSupervisor requires at least one department")
        if len(self._departments) == 1:
            only = next(iter(self._departments.values()))
            return only, DepartmentRoutingDecisionSchema(
                department_name=only.name,
                delegated_prompt=prompt,
                reason="Only one department is registered.",
            )
        routed = await self._route_with_router(prompt, session_id=session_id)
        if routed is not None:
            department = self._departments.get(routed.department_name)
            if department is not None:
                return department, routed
        fallback = self._heuristic_department(prompt)
        return fallback, DepartmentRoutingDecisionSchema(
            department_name=fallback.name,
            delegated_prompt=prompt,
            reason="Heuristic routing fallback.",
        )

    async def _route_with_router(
        self,
        prompt: str,
        *,
        session_id: str | None,
    ) -> DepartmentRoutingDecisionSchema | None:
        if self._router is None or not hasattr(self._router, "run"):
            return None
        routing_prompt = self._build_routing_prompt(prompt)
        try:
            response = await self._router.run(
                routing_prompt,
                session_id=session_id,
                output_schema=DepartmentRoutingDecisionSchema,
            )
        except TypeError:
            response = await self._router.run(routing_prompt, session_id=session_id)
        if isinstance(response, StructuredResponse) and response.parsed is not None:
            return response.parsed
        try:
            return DepartmentRoutingDecisionSchema.model_validate_json(response.content)
        except Exception:
            return None

    def _build_routing_prompt(self, prompt: str) -> str:
        lines = [
            "You are routing a business goal to the best department.",
            "Return strict JSON with keys: department_name, delegated_prompt, reason.",
            "Choose exactly one department from the list below.",
            "Departments:",
        ]
        for department in self._departments.values():
            tags = ", ".join(department.capability_tags) if department.capability_tags else "-"
            lines.append(
                f"- {department.name} [{department.kind}] tags={tags} description={department.description or '-'}"
            )
        lines.extend(
            [
                "",
                "Business goal:",
                prompt,
            ]
        )
        return "\n".join(lines)

    def _heuristic_department(self, prompt: str) -> DepartmentSpec:
        prompt_tokens = _score_tokens(prompt)
        scored = [
            (
                _department_score(prompt_tokens, department, prompt),
                index,
                department,
            )
            for index, department in enumerate(self._departments.values())
        ]
        scored.sort(key=lambda item: (item[0], -item[1]), reverse=True)
        top_score = scored[0][0]
        if top_score <= 0 and self._fallback_department is not None:
            return self._departments[self._fallback_department]
        return scored[0][2]

    async def _run_department(
        self,
        department: DepartmentSpec,
        prompt: str,
        *,
        session_id: str | None,
        **kwargs: Any,
    ) -> Response:
        if isinstance(department.target, Swarm):
            request_topic = department.request_topic or department.name
            return await department.target.dispatch(
                prompt,
                request_topic=request_topic,
                response_topic=department.response_topic,
                session_id=session_id,
                campaign_id=kwargs.pop("campaign_id", None),
                run_id=kwargs.pop("run_id", None),
            )
        return await department.target.run(prompt, session_id=session_id, **kwargs)


def _infer_department_kind(target: Runnable | Swarm) -> DepartmentKind:
    if isinstance(target, Swarm):
        return "swarm"
    target_name = type(target).__name__
    if target_name == "Team":
        return "team"
    if target_name == "CrewProcess":
        return "crew"
    return "runnable"


def _decorate_forwarded_event(
    event: Any,
    *,
    department_name: str | None = None,
    team_name: str | None = None,
) -> dict[str, Any] | None:
    payload = _normalize_routed_event(event)
    payload.setdefault("metadata", {})
    if department_name is not None:
        payload["metadata"]["department_name"] = department_name
    if team_name is not None:
        payload["metadata"]["team_name"] = team_name
    if payload.get("event") == "graph_started":
        return None
    return payload


def _build_supervisor_terminal_event(
    *,
    supervisor_name: str,
    run_id: str,
    sequence: int,
    session_id: str | None,
    correlation_id: str | None,
    trace_id: str | None,
    department: DepartmentSpec,
    decision: DepartmentRoutingDecisionSchema,
    delegated_prompt: str,
    final_event: dict[str, Any] | None,
) -> dict[str, Any]:
    metadata = dict((final_event or {}).get("metadata") or {})
    metadata["department_name"] = department.name
    metadata["department_kind"] = department.kind
    metadata["delegated_prompt"] = delegated_prompt
    metadata["reason"] = decision.reason
    if final_event is None:
        metadata["user_facing"] = {
            "kind": "error",
            "title": "Не удалось завершить задачу",
            "detail": f"Отдел '{department.name}' завершил поток без итогового terminal event.",
            "step_key": "run:department_missing_terminal",
            "status": "failed",
            "visible": True,
        }
        output = {
            "error": "department_missing_terminal",
            "department_name": department.name,
            "department_kind": department.kind,
            "reason": decision.reason,
        }
        status = "failed"
    else:
        output = final_event.get("output")
        status = final_event.get("status") or "completed"
    return make_ui_event(
        event="graph_completed",
        event_type="supervisor",
        graph_name=supervisor_name,
        graph_run_id=run_id,
        sequence=sequence,
        session_id=session_id,
        correlation_id=correlation_id,
        trace_id=trace_id,
        node_name=department.name,
        status=status,
        output=output,
        metadata=metadata,
    )


def _normalize_routed_event(event: Any) -> dict[str, Any]:
    if isinstance(event, dict):
        return dict(event)
    from logos_adk.workflow.graph import GraphNodeEvent, _serialize_graph_ui_event

    if isinstance(event, GraphNodeEvent):
        return _serialize_graph_ui_event(event)
    raise TypeError(f"Unsupported routed event type: {type(event)!r}")


def _score_tokens(value: str) -> set[str]:
    return {
        token
        for token in re.split(r"[^a-zA-Z0-9_]+", value.lower())
        if len(token) >= 3
    }


def _department_score(prompt_tokens: set[str], department: DepartmentSpec, prompt_text: str) -> int:
    haystacks = [
        department.name,
        department.description,
        " ".join(department.capability_tags),
        json.dumps(department.metadata, ensure_ascii=False, sort_keys=True),
    ]
    haystack_tokens: set[str] = set()
    for haystack in haystacks:
        haystack_tokens.update(_score_tokens(haystack))
    score = len(prompt_tokens & haystack_tokens)
    if department.name.lower() in prompt_text.lower():
        score += 2
    return score


def _resolve_swarm_handler(handler: str) -> Any:
    registered = subscription_handler_registry.get(handler)
    if registered is not None:
        return registered
    imported = _load_handler_from_path(handler)
    if imported is not None:
        return imported
    raise ConfigError(
        f"Unknown swarm handler '{handler}'. Register it in subscription_handler_registry "
        "or provide an import path like 'package.module:handler'."
    )


def _load_handler_from_path(path: str) -> Any | None:
    if ":" in path:
        module_name, attr_name = path.split(":", 1)
    elif "." in path:
        module_name, attr_name = path.rsplit(".", 1)
    else:
        return None
    if not module_name or not attr_name:
        return None
    try:
        module = importlib.import_module(module_name)
    except Exception:
        return None
    candidate = getattr(module, attr_name, None)
    return candidate if callable(candidate) else None


__all__ = [
    "DepartmentKind",
    "DepartmentRoutingDecisionSchema",
    "DepartmentSpec",
    "DepartmentSupervisor",
]
