# Models

from typing import Any, List, Dict, Optional
from langchain_community.chat_models import ChatOllama
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

model_configurations: Dict[str, Dict[str, Any]] = {
    "qwen3:32b": {
        "type": "Ollama",
        "model_name": "qwen3:32b",
        "temperature": 0.3,
    },
    "mistral-small:24b": {
        "type": "Ollama",
        "model_name": "mistral-small:24b",
        "temperature": 0.3,
    },
    "deepseek-r1:32b": {
        "type": "Ollama",
        "model_name": "deepseek-r1:32b",
        "temperature": 0.3,
    }
}

class OllamaWithDebug(ChatOllama):
    def __init__(self, model, **kwargs):
        super().__init__(model=model, **kwargs)
        print(f"--- DEBUG INIT: Model {model} ---")
        print(f"Injected parameters: {kwargs}")
        print("Actual instance attributes:")
        for attr in ['temperature', 'top_p', 'base_url']:
            print(f"  - {attr}: {getattr(self, attr, 'UNDEFINED')}")
        print("-" * 40)

    def invoke(self, messages: Any, **kwargs):
        print(f"\n[DEBUG INVOKE] Sending to {self.model}")
        temp = getattr(self, 'temperature', 'N/A')
        print(f"Active config -> Temp: {temp}")

        preview = str(messages)[:150]
        print(f"Input preview: {preview}...")

        try:
            response = super().invoke(messages, **kwargs)
            print(f"[DEBUG RESPONSE] Success! Response length: {len(str(response.content))} chars")
            return response

        except Exception as e:
            print(f"!!! DEBUG ERROR !!! Ollama call failed: {e}")
            raise

def get_model_instance(model_key: str) -> Any:
    config = model_configurations.get(model_key)
    if config is None:
        raise ValueError(f"Model not configured: {model_key}")

    if config.get("type", "").lower() == "ollama":
        return OllamaWithDebug(
            model=config["model_name"],
            base_url="http://localhost:11434",
            temperature=config.get("temperature", 0.3),
            top_p=config.get("top_p"),
        )
    else:
        raise ValueError(f"Unsupported or missing model type: {config.get('type')}")