from __future__ import annotations

from costrouter.models import CostTier
from costrouter.router.base import BaseRouter
from costrouter.router.heuristic import HeuristicRouter
from costrouter.router.llm import LLMRouter


class CascadeRouter(BaseRouter):
    """Heuristic-first, LLM fallback router.

    Runs the free heuristic classifier first. If confidence is high enough,
    uses that result directly. Otherwise, falls back to the LLM classifier.
    """

    def __init__(
        self,
        screener_model: str = "gpt-4o-mini",
        api_key: str | None = None,
        confidence_threshold: float = 0.7,
    ) -> None:
        self.heuristic = HeuristicRouter()
        self.llm = LLMRouter(screener_model=screener_model, api_key=api_key)
        self.confidence_threshold = confidence_threshold

    def classify(self, messages: list[dict]) -> tuple[CostTier, float, float]:
        tier, confidence, _ = self.heuristic.classify(messages)

        if confidence >= self.confidence_threshold:
            return tier, confidence, 0.0

        return self.llm.classify(messages)

    async def aclassify(self, messages: list[dict]) -> tuple[CostTier, float, float]:
        tier, confidence, _ = await self.heuristic.aclassify(messages)

        if confidence >= self.confidence_threshold:
            return tier, confidence, 0.0

        return await self.llm.aclassify(messages)
