import ollama
import structlog

logger = structlog.get_logger()

class ModelManager:
    def __init__(self):
        self.model_capabilities = {
            "qwen2.5:7b": ["general_knowledge", "structured_output", "multilingual", "instruction_following", "structured_data"],
            "llama3.1:8b": ["general_knowledge", "reasoning", "tool_calling", "conversation", "multilingual", "instruction_following"],
            "qwen2.5-coder:7b": ["code_generation", "code_analysis", "instruction_following", "math_reasoning"],
            "llama3.2:3b": ["summarization", "instruction_following", "tool_calling", "multilingual"],
            "llava:7b": ["visual_reasoning", "visual_conversation", "visual_tool_calling", "vision", "ocr", "multimodal"],
        }
        logger.info("ModelManager initialized", model_capabilities=self.model_capabilities)

    def get_model_capabilities(self, model_name):
        capabilities = self.model_capabilities.get(model_name, [])
        logger.debug("Retrieved model capabilities", model=model_name, capabilities=capabilities)
        return capabilities

    def select_best_model(self, required_capability):
        suitable_models = [model for model, capabilities in self.model_capabilities.items() if required_capability in capabilities]
        selected_model = suitable_models[0] if suitable_models else list(self.model_capabilities.keys())[0]
        logger.info("Selected best model", required_capability=required_capability, selected_model=selected_model)
        return selected_model

    def generate_text(self, model_name, prompt, max_length=100, system="You are a helpful assistant."):
        logger.debug("Generating text", model=model_name, prompt=prompt, max_length=max_length)
        # Check if model exists
        try:
            ollama.pull(model_name)
            logger.debug("Model pulled successfully", model=model_name)
        except ollama.RequestError as e:
            if "not found" in str(e):
                logger.error("Model not found", model=model_name)
                return "Model not found"
            else:
                logger.exception("Error pulling model", model=model_name, error=str(e))
                raise e
        response = ollama.generate(model=model_name, prompt=prompt, system=system)
        logger.debug("Text generated", model=model_name, response=response['response'])
        return response['response']

model_manager = ModelManager()