from abc import ABC, abstractmethod
import ollama
from pydantic import BaseModel
from pathlib import Path
from google import genai
from google.genai import types
from mistralai import Mistral


class AbstractModelManager(ABC):
    def __init__(self, model_name, system_prompt_file="system.prompt"):
        self.model_name = model_name
        script_dir = Path(__file__).parent
        self.system_prompt_file = script_dir / system_prompt_file

    @abstractmethod
    def is_model_loaded(self, model):
        pass

    @abstractmethod
    def create_model(self, base_model, context_window=4096, temperature=0):
        pass

    @abstractmethod
    def request(self, prompt):
        pass

    @abstractmethod
    def delete(self):
        pass

class OllamaModelManager(AbstractModelManager):
    def is_model_loaded(self, model):
        loaded_models = [m.model for m in ollama.list().models]
        return model in loaded_models or f'{model}:latest' in loaded_models

    def create_model(self, base_model, context_window=4096, temperature=0):
        with open(self.system_prompt_file, 'r') as f:
            system = f.read()
        
        if not self.is_model_loaded(self.model_name):
            print(f"Creating model {self.model_name}")
            ollama.create(
                model=self.model_name,
                from_=base_model,
                system=system,
                parameters={
                    "num_ctx": context_window,
                    "temperature": temperature
                }
            )

    def request(self, prompt):
        response = ollama.chat(
            model=self.model_name, 
            messages=[{"role": "user", "content": prompt}],
        )
        response = response['message']['content']
        return response
    
    def delete(self):
        if self.is_model_loaded("C2Rust:latest"):
            print(f"Deleting model {self.model_name}")
            ollama.delete("C2Rust:latest")
        else:
            print(f"Model {self.model_name} not found, skipping deletion.")

class GeminiModelManager(AbstractModelManager):
    def __init__(self, api_key):
        super().__init__()
        self.client = genai.Client(api_key=api_key)
        self.model = "gemini-2.0-flash"
        # read system prompt from file
        with open(self.system_prompt_file, 'r') as f:
            self.system_instruction = f.read()


    def is_model_loaded(self, model):
        # Check if the specified model is the one set in the manager
        return model == self.model

    def create_model(self, base_model=None, context_window=4096, temperature=0):
        # Initialize the Gemini model settings (if applicable)
        self.model = base_model if base_model else "gemini-2.0-flash"

    def request(self, prompt, temperature=0, context_window=4096):
        # Request response from the Gemini model
        response = self.client.models.generate_content(
            model=self.model,
            contents=prompt,
            config=types.GenerateContentConfig(
                temperature=temperature,
                max_output_tokens=context_window,
                system_instruction=self.system_instruction,
            )
        )
        return response.text

    def delete(self):
        # Implement model deletion logic (if applicable)
        self.model = None

class MistralModelManager(AbstractModelManager):
    def __init__(self, api_key, model_name="mistral-small-latest", system_prompt_file="system.prompt"):
        super().__init__()
        self.client = Mistral(api_key=api_key)
        self.model = model_name
        # read system prompt from file
        with open(self.system_prompt_file, 'r') as f:
            self.system_instruction = f.read()

    def is_model_loaded(self, model):
        # Check if the specified model is the one set in the manager
        return model == self.model

    def create_model(self, base_model=None, context_window=4096, temperature=0):
        # Initialize the Mistral model settings (if applicable)
        self.model = base_model if base_model else "mistral-small-latest"

    def request(self, prompt, temperature=0, context_window=4096):
        # Request response from the Mistral model
        response = self.client.chat.complete(
            messages=[
            {
                "role":"user",
                "content": self.system_instruction + "\n" + prompt,
            }
            ],
            model=self.model,
            temperature=temperature,
            max_tokens=context_window,
        )
        return response.text

    def delete(self):
        # Implement model deletion logic (if applicable)
        self.model = None