from __future__ import annotations

from pathlib import Path

import torch

from .config import WorkerPoolConfig, load_worker_pool
from .model import RouterModel
from .providers import WorkerCallError, WorkerRunner
from .schemas import TaskRecord, route_text
from .training_common import resolve_device


class FuguLiteOrchestrator:
    def __init__(self, checkpoint: str | Path, worker_config: str | Path):
        self.pool: WorkerPoolConfig = load_worker_pool(worker_config)
        self.model = RouterModel.from_checkpoint(checkpoint, device=resolve_device())
        if self.model.worker_ids != self.pool.worker_ids:
            raise ValueError(
                f"Checkpoint workers {self.model.worker_ids} do not match "
                f"configured workers {self.pool.worker_ids}"
            )
        self.model.eval()
        self.runner = WorkerRunner(self.pool)

    @torch.no_grad()
    def route(self, prompt: str, domain: str = "general", tags: list[str] | None = None) -> dict:
        text = route_text(prompt, domain, tags or [])
        encoded = self.model.tokenizer(
            [text],
            padding=True,
            truncation=True,
            max_length=self.model.router_config.max_length,
            return_tensors="pt",
        )
        inputs = {
            "input_ids": encoded["input_ids"].to(self.model.device_ref),
            "attention_mask": encoded["attention_mask"].to(self.model.device_ref),
        }
        logits = self.model(**inputs)[0]
        probabilities = torch.softmax(logits, dim=-1)
        order = torch.argsort(probabilities, descending=True).tolist()
        ranked = [
            {"worker_id": self.model.worker_ids[index], "probability": float(probabilities[index])}
            for index in order
        ]
        return {"selected_worker": ranked[0]["worker_id"], "ranking": ranked}

    async def answer(
        self, prompt: str, domain: str = "general", tags: list[str] | None = None
    ) -> dict:
        routing = self.route(prompt, domain, tags)
        worker_map = {worker.id: worker for worker in self.pool.workers}
        task = TaskRecord(
            task_id="inference",
            prompt=prompt,
            domain=domain,
            tags=tags or [],
        )
        errors = []
        for candidate in routing["ranking"]:
            worker = worker_map[candidate["worker_id"]]
            try:
                result = await self.runner.complete(worker, task)
                return {"routing": routing, "result": result.model_dump(), "fallback_errors": errors}
            except WorkerCallError as exc:
                errors.append({"worker_id": worker.id, "error": str(exc)})
        raise WorkerCallError(f"Every routed worker failed: {errors}")

    async def close(self) -> None:
        await self.runner.close()