"""
Rile Chat - Gemini AI integration for persona-based conversations
"""

import json
import re
from pathlib import Path
from typing import Dict, List, Optional, Union

from .loader import load_persona
from .utils import build_system_prompt

# Try to import google.generativeai, but don't fail if not installed
try:
    import google.generativeai as genai
    GEMINI_AVAILABLE = True
except ImportError:
    GEMINI_AVAILABLE = False
    genai = None


class Rile:
    """
    Main class for chatting with Rile personas using Gemini AI.
    
    Example:
        >>> from rilerary import Rile
        >>> riley = Rile(api_key="your-gemini-api-key")
        >>> riley.use_persona("chef_rile")
        >>> response = riley.chat("How do I make pasta?")
        >>> print(riley.history)
    """
    
    def __init__(self, api_key: str, model: str = "gemini-2.5-flash", plain_text_mode: bool = False):
        """
        Initialize Rile with a Gemini API key.
        
        Args:
            api_key: Your Google Gemini API key
            model: Gemini model to use (default: gemini-2.5-flash)
            plain_text_mode: If True, strips markdown formatting from responses.
                           Default is False (keeps markdown for HTML rendering).
        
        Raises:
            ImportError: If google-generativeai is not installed
        """
        if not GEMINI_AVAILABLE:
            raise ImportError(
                "google-generativeai is required for Rile chat functionality. "
                "Install it with: pip install google-generativeai"
            )
        
        self._api_key = api_key
        self._model_name = model
        self._plain_text_mode = plain_text_mode
        self._persona_name: Optional[str] = None
        self._persona_data: Optional[Dict] = None
        self._system_prompt: str = ""
        self._history: List[Dict[str, str]] = []
        self._chat_session = None
        
        # Configure Gemini
        genai.configure(api_key=api_key)
        self._model = genai.GenerativeModel(model)
    
    def _strip_markdown(self, text: str) -> str:
        """
        Strip markdown formatting from text for plain text output.
        
        Removes: **bold**, *italic*, # headers, ``` code blocks, etc.
        """
        # Remove code blocks
        text = re.sub(r'```[\s\S]*?```', lambda m: m.group().replace('```', '').strip(), text)
        # Remove inline code
        text = re.sub(r'`([^`]+)`', r'\1', text)
        # Remove bold
        text = re.sub(r'\*\*([^*]+)\*\*', r'\1', text)
        # Remove italic
        text = re.sub(r'\*([^*]+)\*', r'\1', text)
        # Remove headers
        text = re.sub(r'^#+\s*', '', text, flags=re.MULTILINE)
        # Remove bullet points markers but keep content
        text = re.sub(r'^[\*\-]\s+', '• ', text, flags=re.MULTILINE)
        # Remove numbered list formatting
        text = re.sub(r'^\d+\.\s+', '', text, flags=re.MULTILINE)
        # Remove links but keep text
        text = re.sub(r'\[([^\]]+)\]\([^)]+\)', r'\1', text)
        
        return text.strip()
    
    def use_persona(self, persona_name: str) -> "Rile":
        """
        Set the persona to use for conversations.
        
        Args:
            persona_name: Name of the persona (e.g., 'chef_rile', 'tech_rile')
            
        Returns:
            Self for method chaining
        """
        self._persona_name = persona_name
        self._persona_data = load_persona(persona_name)
        self._system_prompt = build_system_prompt(self._persona_data)
        self._chat_session = None  # Reset chat session for new persona
        return self
    
    def _get_chat_session(self):
        """Get or create a chat session with conversation history."""
        if self._chat_session is None:
            # Build conversation history for Gemini
            gemini_history = []
            for entry in self._history:
                gemini_history.append({
                    "role": "user",
                    "parts": [entry["question"]]
                })
                gemini_history.append({
                    "role": "model", 
                    "parts": [entry["answer"]]
                })
            
            self._chat_session = self._model.start_chat(history=gemini_history)
        
        return self._chat_session
    
    def chat(self, text: str, image: Optional[str] = None) -> str:
        """
        Send a message and get a response from the persona.
        
        The question and answer are automatically stored in history.
        
        Args:
            text: Your question or message
            image: Optional path to an image file for visual analysis
            
        Returns:
            The persona's response
            
        Raises:
            ValueError: If no persona has been set
        """
        if self._persona_data is None:
            raise ValueError(
                "No persona set. Call use_persona() first. "
                "Example: rile.use_persona('chef_rile')"
            )
        
        # Build the full prompt with system instructions
        full_prompt = f"{self._system_prompt}\n\nUser: {text}"
        
        # Prepare content parts
        content_parts = [text]
        
        # Add image if provided
        if image:
            import PIL.Image
            img = PIL.Image.open(image)
            content_parts = [text, img]
        
        # If we have history, use chat session for context
        if self._history and not image:
            chat = self._get_chat_session()
            response = chat.send_message(text)
        else:
            # First message or image message - generate directly
            if image:
                response = self._model.generate_content([full_prompt, PIL.Image.open(image)])
            else:
                response = self._model.generate_content(full_prompt)
        
        answer = response.text
        
        # Apply plain text mode if enabled
        if self._plain_text_mode:
            answer = self._strip_markdown(answer)
        
        # Store in history
        self._history.append({
            "question": text,
            "answer": answer,
            "image": image if image else None
        })
        
        # Reset chat session so next call includes updated history
        self._chat_session = None
        
        return answer
    
    def auto_persona(self, question: str) -> str:
        """
        Automatically select the best persona based on the question.
        
        Uses AI to analyze the question and choose the most appropriate persona.
        
        Args:
            question: The question to analyze
            
        Returns:
            Name of the selected persona
        """
        from .loader import get_available_personas, get_persona_info
        
        # Get all available personas with their descriptions
        personas = get_available_personas()
        persona_descriptions = []
        for p in personas:
            info = get_persona_info(p)
            persona_descriptions.append(f"- {p}: {info['description']}")
        
        persona_list = "\n".join(persona_descriptions)
        
        # Ask AI to select the best persona
        selection_prompt = f"""Based on the following question, select the most appropriate persona from this list.
        
Available personas:
{persona_list}

Question: "{question}"

Respond with ONLY the persona name (e.g., 'chef_rile' or 'tech_rile'). No explanation, just the name."""

        response = self._model.generate_content(selection_prompt)
        selected = response.text.strip().lower().replace("'", "").replace('"', '')
        
        # Validate the selection
        if selected not in personas:
            # Default to knowledge_rile if invalid selection
            selected = "knowledge_rile"
        
        return selected
    
    def auto_chat(self, text: str, image: Optional[str] = None) -> str:
        """
        Automatically select the best persona and respond to the question.
        
        Combines auto_persona() and chat() into a single call.
        
        Args:
            text: Your question or message
            image: Optional path to an image file for visual analysis
            
        Returns:
            The persona's response
        """
        selected_persona = self.auto_persona(text)
        self.use_persona(selected_persona)
        return self.chat(text, image=image)
    
    @property
    def history(self) -> List[Dict[str, str]]:
        """
        Get the conversation history.
        
        Returns:
            List of dictionaries with 'question' and 'answer' keys
        """
        return self._history.copy()
    
    @property
    def persona(self) -> Optional[Dict]:
        """Get the current persona data."""
        return self._persona_data
    
    @property
    def persona_name(self) -> Optional[str]:
        """Get the current persona name."""
        return self._persona_name
    
    def clear_history(self) -> "Rile":
        """
        Clear the conversation history.
        
        Returns:
            Self for method chaining
        """
        self._history = []
        self._chat_session = None
        return self
    
    def save_history(self, filepath: str) -> None:
        """
        Save conversation history to a JSON file.
        
        Args:
            filepath: Path to save the history file
        """
        data = {
            "persona": self._persona_name,
            "history": self._history
        }
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=2, ensure_ascii=False)
    
    def load_history(self, filepath: str) -> "Rile":
        """
        Load conversation history from a JSON file.
        
        Args:
            filepath: Path to the history file
            
        Returns:
            Self for method chaining
        """
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # Load persona if specified in history file
        if data.get("persona") and data["persona"] != self._persona_name:
            self.use_persona(data["persona"])
        
        self._history = data.get("history", [])
        self._chat_session = None  # Reset to rebuild with loaded history
        
        return self
    
    def __repr__(self) -> str:
        persona_info = self._persona_name or "None"
        return f"Riley(persona={persona_info}, history_length={len(self._history)})"
