from __future__ import annotations
from services import hf_provider, featherless_provider

class ModelRouter:
    """
    Routes tasks to the correct model and provider.
    understand() → MiniCPM-V (HF) — all document tasks
    reason()     → Nemotron 4B (Featherless) — all reasoning tasks
    translate()  → Tiny Aya (HF) — multilingual
    transcribe() → Whisper (HF) — speech to text
    """
    def __init__(self, ocr_model: str, reasoning_model: str,
                 multilingual_model: str, speech_model: str,
                 hf_api_key: str = "", featherless_api_key: str = "",
                 max_tokens: int = 1024, temperature: float = 0.3):
        self._ocr_model = ocr_model
        self._reason_model = reasoning_model
        self._multi_model = multilingual_model
        self._speech_model = speech_model
        self._hf_key = hf_api_key
        self._fl_key = featherless_api_key
        self._max_tokens = max_tokens
        self._temperature = temperature

    def understand(self, prompt: str, image_b64: str = "") -> str:
        """MiniCPM-V via HuggingFace — OCR, diagram reading, concept extraction."""
        if image_b64:
            return hf_provider.vision_generate(
                self._ocr_model, image_b64, prompt,
                self._hf_key, self._max_tokens)
        return hf_provider.generate(
            self._ocr_model, prompt, api_key=self._hf_key,
            max_tokens=self._max_tokens, temperature=self._temperature)

    def reason(self, prompt: str, system: str = "") -> str:
        """Nemotron 3 Nano 4B via Featherless — quests, questions, tutor."""
        return featherless_provider.generate(
            self._reason_model, prompt, system,
            self._fl_key, self._max_tokens, self._temperature)

    def translate(self, prompt: str, system: str = "") -> str:
        """Tiny Aya 3.3B via HuggingFace — multilingual explanations."""
        return hf_provider.generate(
            self._multi_model, prompt, system,
            self._hf_key, self._max_tokens, self._temperature)

    def transcribe(self, audio_bytes: bytes) -> str:
        """Whisper via HuggingFace — speech to text."""
        return hf_provider.transcribe(self._speech_model, audio_bytes, self._hf_key)