"""Phi-3.5-mini model wrapper using llama-cpp-python."""

from typing import Optional, List, Dict, Any
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

from src.config import ModelConfig


class PhiModel:
    """Wrapper for Phi-3.5-mini model."""
    
    def __init__(self, config: Optional[ModelConfig] = None):
        """Initialize the model wrapper.
        
        Args:
            config: Model configuration. Uses defaults if not provided.
        """
        self.config = config or ModelConfig()
        self._model: Optional[Llama] = None
        self._model_path: Optional[str] = None
    
    @property
    def model(self) -> Llama:
        """Lazy load the model."""
        if self._model is None:
            self._load_model()
        return self._model
    
    def _load_model(self) -> None:
        """Download and load the model."""
        print(f"Downloading model from {self.config.repo_id}...")
        self._model_path = hf_hub_download(
            repo_id=self.config.repo_id,
            filename=self.config.filename
        )
        
        print("Loading model into memory...")
        self._model = Llama(
            model_path=self._model_path,
            n_ctx=self.config.n_ctx,
            n_threads=self.config.n_threads,
            verbose=self.config.verbose
        )
        print("Model loaded successfully!")
    
    def generate(self, prompt: str, max_tokens: Optional[int] = None) -> str:
        """Generate text completion.
        
        Args:
            prompt: Input prompt.
            max_tokens: Maximum tokens to generate.
            
        Returns:
            Generated text.
        """
        output = self.model(
            prompt,
            max_tokens=max_tokens or self.config.max_tokens,
            temperature=self.config.temperature,
            echo=False
        )
        return output["choices"][0]["text"].strip()
    
    def chat(
        self, 
        messages: List[Dict[str, str]], 
        max_tokens: Optional[int] = None
    ) -> str:
        """Generate chat completion.
        
        Args:
            messages: List of message dicts with 'role' and 'content'.
            max_tokens: Maximum tokens to generate.
            
        Returns:
            Assistant's response.
        """
        output = self.model.create_chat_completion(
            messages=messages,
            max_tokens=max_tokens or self.config.max_tokens,
            temperature=self.config.temperature
        )
        return output["choices"][0]["message"]["content"].strip()
    
    def chat_with_context(
        self, 
        query: str, 
        context: str,
        system_prompt: Optional[str] = None
    ) -> str:
        """Generate response with RAG context.
        
        Args:
            query: User's question.
            context: Retrieved context from documents.
            system_prompt: Optional system prompt.
            
        Returns:
            Generated response.
        """
        if system_prompt is None:
            system_prompt = (
                "You are a helpful assistant. Answer the user's question based on "
                "the provided context. If the context doesn't contain relevant "
                "information, say so honestly. Be concise and accurate."
            )
        
        user_message = f"""Context:
{context}

Question: {query}

Please answer based on the context provided above."""
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ]
        
        return self.chat(messages)