"""LLM integration for Superviseur bot."""

import json
import aiohttp
import logging
import re
from typing import Optional, Dict, Any

logger = logging.getLogger('Superviseur')


class LLMManager:
    """Manages LLM interactions with Ollama."""
    
    def __init__(
        self,
        ollama_api_url: str,
        ollama_model: str,
        ollama_timeout: int = 60,
        ollama_temperature: float = 0.7
    ):
        self.ollama_api_url = ollama_api_url
        self.ollama_model = ollama_model
        self.ollama_timeout = ollama_timeout
        self.ollama_temperature = ollama_temperature
        self.http_session: Optional[aiohttp.ClientSession] = None
    
    async def get_session(self) -> aiohttp.ClientSession:
        """Get or create HTTP session."""
        if not self.http_session or self.http_session.closed:
            self.http_session = aiohttp.ClientSession()
        return self.http_session
    
    async def close_session(self) -> None:
        """Close HTTP session."""
        if self.http_session and not self.http_session.closed:
            await self.http_session.close()
    
    def build_payload(
        self,
        prompt: str,
        system_prompt: Optional[str] = None,
        vision_model: bool = False,
        attachments: Optional[list] = None
    ) -> Dict[str, Any]:
        """Build the API payload for Ollama."""
        payload = {
            "model": self.ollama_model,
            "prompt": prompt,
            "stream": True,
            "options": {"temperature": self.ollama_temperature}
        }
        
        if system_prompt:
            payload["system"] = system_prompt
        
        # Vision support (simplified)
        if vision_model and attachments:
            for attachment in attachments:
                if attachment.content_type and attachment.content_type.startswith('image/'):
                    if any(v in self.ollama_model.lower() for v in 
                           ['llava', 'bakllava', 'moondream', 'llama3.2-vision', 'minicpm-v']):
                        # Note: actual image handling would require download
                        pass
        
        return payload
    
    async def call_ollama(self, payload: Dict[str, Any]) -> str:
        """Call Ollama API and accumulate response."""
        session = await self.get_session()

        async with session.post(
            self.ollama_api_url,
            json=payload,
            timeout=aiohttp.ClientTimeout(total=self.ollama_timeout)
        ) as response:
            response.raise_for_status()
            
            # Si on ne streame pas, on récupère tout d'un coup
            if not payload.get("stream", True):
                data = await response.json()
                return data.get("response", "")

            accumulated_reply = ""
            async for line in response.content:
                line = line.decode('utf-8').strip()
                if not line:
                    continue
                    
                try:
                    data = json.loads(line)
                    
                    if 'thinking' in data and data['thinking']:
                        print(f"🤔 Le superviseur réfléchit : {data['thinking']}", flush=True)
                        continue

                    if 'response' in data:
                        chunk = data['response']
                        accumulated_reply += chunk
                        if chunk.strip():
                            print(f"🤖 Réponse du Superviseur (fragment) : {chunk}", flush=True)
                    
                    elif 'action' in data:
                        accumulated_reply = line
                        break

                except json.JSONDecodeError:
                    accumulated_reply += line

            return accumulated_reply
    
    def extract_json_actions(self, text: str) -> Optional[str]:
        """Extract JSON action objects from text more robustly."""
        actions = []
        
        # Regex plus laxiste pour trouver des blocs JSON potentiels
        # On cherche des blocs commençant par {"action" ou {'action'
        matches = re.finditer(r'\{[\s\n]*[\"\']action[\"\'].*?\}', text, re.DOTALL)
        
        for match in matches:
            json_str = match.group(0)
            try:
                # Tentative de parsing standard
                parsed = json.loads(json_str)
                if isinstance(parsed, dict) and 'action' in parsed:
                    actions.append(parsed)
            except json.JSONDecodeError:
                # Si erreur, on tente de remplacer les simples quotes par des doubles 
                # (Cas fréquent si le LLM mélange les styles)
                try:
                    # Remplacement très basique (peut échouer si contenu complexe)
                    json_corrected = json_str.replace("'", '"')
                    parsed = json.loads(json_corrected)
                    if isinstance(parsed, dict) and 'action' in parsed:
                        actions.append(parsed)
                        continue
                except: pass

                # Méthode récursive des accolades si le bloc est tronqué ou complexe
                start = match.start()
                brace_count = 0
                for i in range(start, len(text)):
                    if text[i] == '{': brace_count += 1
                    elif text[i] == '}': 
                        brace_count -= 1
                        if brace_count == 0:
                            try:
                                candidate = text[start:i+1]
                                try:
                                    parsed = json.loads(candidate)
                                except:
                                    parsed = json.loads(candidate.replace("'", '"'))
                                
                                if isinstance(parsed, dict) and 'action' in parsed:
                                    actions.append(parsed)
                                    break
                            except: pass
        
        if actions:
            return json.dumps(actions)
        return None