"""
Serviço de integração com Modelos de Linguagem
"""
import os
import aiohttp
import asyncio
from typing import AsyncGenerator, Tuple, Optional, Dict, Any
from dataclasses import dataclass

from config.settings import LLMConfig


@dataclass
class LLMResponse:
    """Resposta do modelo"""
    content: str
    tokens_used: int
    model: str
    finish_reason: str


class LLMService:
    """
    Interface unificada para diferentes provedores de LLM
    Suporta: OpenAI, Anthropic, e modelos locais
    """
    
    def __init__(self, config: LLMConfig):
        self.config = config
        self.provider = config.provider
        self.session: Optional[aiohttp.ClientSession] = None
    
    async def _get_session(self) -> aiohttp.ClientSession:
        """Obtém ou cria sessão HTTP"""
        if self.session is None or self.session.closed:
            self.session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            )
        return self.session
    
    async def generate(
        self, 
        prompt: str, 
        system: Optional[str] = None,
        max_tokens: Optional[int] = None
    ) -> Tuple[str, int]:
        """
        Gera texto completo (não-streaming)
        Retorna: (texto, tokens_usados)
        """
        if self.provider == "openai":
            return await self._openai_generate(prompt, system, max_tokens)
        elif self.provider == "anthropic":
            return await self._anthropic_generate(prompt, system, max_tokens)
        elif self.provider == "local":
            return await self._local_generate(prompt, system, max_tokens)
        else:
            raise ValueError(f"Provedor não suportado: {self.provider}")
    
    async def generate_stream(
        self,
        prompt: str,
        system: Optional[str] = None
    ) -> AsyncGenerator[str, None]:
        """
        Gera texto em streaming (chunk por chunk)
        """
        if self.provider == "openai":
            async for chunk in self._openai_stream(prompt, system):
                yield chunk
        else:
            # Fallback: gera completo e yield por palavras
            text, _ = await self.generate(prompt, system)
            words = text.split(" ")
            for word in words:
                yield word + " "
                await asyncio.sleep(0.01)
    
    async def _openai_generate(
        self, 
        prompt: str, 
        system: Optional[str],
        max_tokens: Optional[int]
    ) -> Tuple[str, int]:
        """Integração com OpenAI"""
        session = await self._get_session()
        
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json"
        }
        
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        
        payload = {
            "model": self.config.model,
            "messages": messages,
            "temperature": self.config.temperature,
            "max_tokens": max_tokens or self.config.max_tokens
        }
        
        base_url = self.config.base_url or "https://api.openai.com/v1"
        
        async with session.post(
            f"{base_url}/chat/completions",
            headers=headers,
            json=payload
        ) as response:
            if response.status != 200:
                error = await response.text()
                raise Exception(f"OpenAI API error: {error}")
            
            data = await response.json()
            content = data["choices"][0]["message"]["content"]
            tokens = data.get("usage", {}).get("total_tokens", 0)
            
            return content.strip(), tokens
    
    async def _openai_stream(
        self,
        prompt: str,
        system: Optional[str]
    ) -> AsyncGenerator[str, None]:
        """Streaming com OpenAI"""
        session = await self._get_session()
        
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json"
        }
        
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        
        payload = {
            "model": self.config.model,
            "messages": messages,
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "stream": True
        }
        
        base_url = self.config.base_url or "https://api.openai.com/v1"
        
        async with session.post(
            f"{base_url}/chat/completions",
            headers=headers,
            json=payload
        ) as response:
            async for line in response.content:
                line = line.decode("utf-8").strip()
                if line.startswith("data: "):
                    data_str = line[6:]
                    if data_str == "[DONE]":
                        break
                    try:
                        import json
                        data = json.loads(data_str)
                        delta = data["choices"][0]["delta"].get("content", "")
                        if delta:
                            yield delta
                    except:
                        pass
    
    async def _anthropic_generate(
        self,
        prompt: str,
        system: Optional[str],
        max_tokens: Optional[int]
    ) -> Tuple[str, int]:
        """Integração com Anthropic Claude"""
        session = await self._get_session()
        
        headers = {
            "x-api-key": self.config.api_key,
            "anthropic-version": "2023-06-01",
            "Content-Type": "application/json"
        }
        
        payload = {
            "model": self.config.model,
            "max_tokens": max_tokens or self.config.max_tokens,
            "messages": [
                {"role": "user", "content": prompt}
            ]
        }
        
        if system:
            payload["system"] = system
        
        async with session.post(
            "https://api.anthropic.com/v1/messages",
            headers=headers,
            json=payload
        ) as response:
            if response.status != 200:
                error = await response.text()
                raise Exception(f"Anthropic API error: {error}")
            
            data = await response.json()
            content = data["content"][0]["text"]
            tokens = data.get("usage", {}).get("input_tokens", 0)
            tokens += data.get("usage", {}).get("output_tokens", 0)
            
            return content.strip(), tokens
    
    async def _local_generate(
        self,
        prompt: str,
        system: Optional[str],
        max_tokens: Optional[int]
    ) -> Tuple[str, int]:
        """Integração com modelo local (Ollama, etc)"""
        session = await self._get_session()
        
        payload = {
            "model": self.config.model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": self.config.temperature
            }
        }
        
        base_url = self.config.base_url or "http://localhost:11434"
        
        async with session.post(
            f"{base_url}/api/generate",
            json=payload
        ) as response:
            if response.status != 200:
                error = await response.text()
                raise Exception(f"Local model error: {error}")
            
            data = await response.json()
            return data.get("response", "").strip(), 0
    
    async def close(self):
        """Fecha conexões"""
        if self.session and not self.session.closed:
            await self.session.close()