import json
from services.ollama_service import OllamaService
from services.external_model_service import ExternalModelService

class ModelRouter:
    def __init__(self):
        self.ollama = OllamaService()
        self.external = ExternalModelService()

        with open("models.json", "r") as f:
            self.models = json.load(f)

    def pick_model(self, intent: str, mode: str):
        # regras simples (pode ser expandido)
        if intent == "document" and mode == "advogado":
            return self.models["local"]["llama3.1-70b"]

        if intent == "document" and mode == "corporativo":
            return self.models["local"]["qwen2.5-7b"]

        if intent == "search":
            return self.models["external"]["claude-3.5-sonnet"]

        if intent == "monitoring":
            return self.models["external"]["mistral-large"]

        # fallback
        return self.models["local"]["llama3.1-8b"]

    def generate(self, intent: str, mode: str, prompt: str):
        model = self.pick_model(intent, mode)

        if model["provider"] == "ollama":
            return self.ollama.generate(prompt)

        if model["provider"] in ["openai", "anthropic", "mistral"]:
            return self.external.generate(model["provider"], model["model"], prompt)

        return "Erro: modelo não encontrado."