"""Classe principal do IA Cost Tracker."""

import logging
from typing import Optional, Dict, Any, List
from datetime import datetime
import time

from .database import DatabaseManager
from .models import ModeloIA, UsoToken
from .providers import MaritacaProvider, AnthropicProvider
from .exceptions import ModelNotFoundError, ConfigurationError

logger = logging.getLogger(__name__)


class IATracker:
    """
    Classe principal para rastreamento de custos de IA.
    
    Esta classe gerencia o rastreamento de uso e custos de chamadas
    a APIs de IA, incluindo sincronização de preços e registro de uso.
    
    Example:
        >>> tracker = IATracker(
        ...     db_connection_string="postgresql://user:pass@localhost/db",
        ...     maritaca_api_key="sua-chave"
        ... )
        >>> 
        >>> # Sincronizar preços da Maritaca
        >>> tracker.sincronizar_precos_maritaca()
        >>> 
        >>> # Registrar uso (Anthropic busca preços do banco automaticamente)
        >>> tracker.registrar_chamada(
        ...     provedor="anthropic",
        ...     modelo="claude-3-5-sonnet-20241022",
        ...     tokens_input=100,
        ...     tokens_output=50,
        ...     aplicacao="meu-app",
        ...     tag_funcionalidade="chat"
        ... )
    """
    
    def __init__(self,
                 db_connection_string: str,
                 maritaca_api_key: Optional[str] = None,
                 maritaca_base_url: str = "https://chat.maritaca.ai"):
        """
        Inicializa o IA Cost Tracker.
        
        Args:
            db_connection_string: String de conexão PostgreSQL
            maritaca_api_key: Chave de API da Maritaca (opcional)
            maritaca_base_url: URL base da API Maritaca
        
        Raises:
            ConfigurationError: Se a configuração for inválida
        """
        self.db = DatabaseManager(db_connection_string)
        
        # Maritaca Provider (precisa de API key)
        self.maritaca = None
        if maritaca_api_key:
            self.maritaca = MaritacaProvider(maritaca_api_key, maritaca_base_url)
        
        # Anthropic Provider (usa o banco de dados)
        self.anthropic = AnthropicProvider(db_manager=self.db)
        
        logger.info("IATracker inicializado com sucesso")
    
    def sincronizar_precos_maritaca(self) -> int:
        """
        Sincroniza os preços dos modelos da Maritaca.
        
        Busca os modelos e preços da API da Maritaca e
        atualiza o banco de dados.
        
        Returns:
            Número de modelos sincronizados
        
        Raises:
            ConfigurationError: Se a Maritaca não estiver configurada
            ProviderError: Se houver erro ao buscar os modelos
        """
        if not self.maritaca:
            raise ConfigurationError(
                "Maritaca não configurada. Forneça maritaca_api_key na inicialização."
            )
        
        modelos = self.maritaca.buscar_modelos_disponiveis()
        
        count = 0
        for modelo in modelos:
            self.db.inserir_ou_atualizar_modelo(modelo)
            count += 1
        
        logger.info(f"{count} modelos da Maritaca sincronizados")
        return count
    
    def listar_modelos_anthropic(self) -> List[ModeloIA]:
        """
        Lista os modelos da Anthropic disponíveis no banco de dados.
        
        Returns:
            Lista de ModeloIA da Anthropic
        
        Note:
            Os preços da Anthropic devem ser mantidos manualmente no banco,
            pois não há API pública para consulta.
        """
        modelos = self.anthropic.buscar_modelos_disponiveis()
        logger.info(f"{len(modelos)} modelos da Anthropic encontrados no banco")
        return modelos
    
    def registrar_chamada(self,
                        provedor: str,
                        modelo: str,
                        tokens_input: int,
                        tokens_output: int,
                        aplicacao: Optional[str] = None,
                        tag_funcionalidade: Optional[str] = None,
                        usuario: Optional[str] = None,
                        tempo_resposta_ms: Optional[int] = None,
                        erro: bool = False,
                        mensagem_erro: Optional[str] = None) -> int:
        """
        Registra uma chamada de IA no banco de dados.
        
        Args:
            provedor: Nome do provedor (maritaca, anthropic)
            modelo: Nome do modelo utilizado
            tokens_input: Quantidade de tokens de entrada
            tokens_output: Quantidade de tokens de saída
            aplicacao: Nome da aplicação que fez a chamada
            tag_funcionalidade: Tag para identificar a funcionalidade
            usuario: Usuário que realizou a chamada
            tempo_resposta_ms: Tempo de resposta em milissegundos
            erro: Se houve erro na chamada
            mensagem_erro: Mensagem de erro, se houver
        
        Returns:
            ID do registro de uso criado
        
        Raises:
            ModelNotFoundError: Se o modelo não for encontrado no banco
        
        Example:
            >>> tracker.registrar_chamada(
            ...     provedor="anthropic",
            ...     modelo="claude-3-5-sonnet-20241022",
            ...     tokens_input=150,
            ...     tokens_output=75,
            ...     aplicacao="chat-app",
            ...     tag_funcionalidade="suporte"
            ... )
        """
        # Buscar modelo no banco
        modelo_db = self.db.buscar_modelo(provedor, modelo)
        
        if not modelo_db:
            raise ModelNotFoundError(
                f"Modelo {provedor}/{modelo} não encontrado no banco de dados. "
                f"Para Maritaca, execute sincronizar_precos_maritaca(). "
                f"Para Anthropic, adicione o modelo manualmente no banco."
            )
        
        # ✅ CORREÇÃO: Calcular custos diretamente (sem precisar de provider!)
        custo_input = (tokens_input / 1000) * modelo_db.custo_input_por_1k
        custo_output = (tokens_output / 1000) * modelo_db.custo_output_por_1k
        
        custos = {
            'custo_input': round(custo_input, 6),
            'custo_output': round(custo_output, 6),
            'custo_total': round(custo_input + custo_output, 6)
        }
        
        # Criar registro de uso
        uso = UsoToken(
            id_modelo=modelo_db.id_modelo,
            tokens_input=tokens_input,
            tokens_output=tokens_output,
            custo_input=custos['custo_input'],
            custo_output=custos['custo_output'],
            aplicacao=aplicacao,
            tag_funcionalidade=tag_funcionalidade,
            usuario=usuario,
            tempo_resposta_ms=tempo_resposta_ms,
            erro=erro,
            mensagem_erro=mensagem_erro
        )
        
        return self.db.registrar_uso(uso)
    
    def _get_provider(self, provedor: str):
        """
        Retorna o provider correspondente ao nome.
        
        Args:
            provedor: Nome do provedor
        
        Returns:
            Instância do provider
        
        Raises:
            ConfigurationError: Se o provedor não for suportado
        """
        provedor_lower = provedor.lower()
        
        if provedor_lower == "maritaca":
            if not self.maritaca:
                raise ConfigurationError(
                    "Maritaca não configurada. Forneça maritaca_api_key na inicialização."
                )
            return self.maritaca
        elif provedor_lower == "anthropic":
            return self.anthropic
        else:
            raise ConfigurationError(
                f"Provedor desconhecido: {provedor}. "
                f"Provedores suportados: maritaca, anthropic"
            )
    
    def track_call(self, 
                   provedor: str,
                   modelo: str,
                   aplicacao: Optional[str] = None,
                   tag_funcionalidade: Optional[str] = None,
                   usuario: Optional[str] = None):
        """
        Decorator para rastrear automaticamente chamadas de IA.
        
        Args:
            provedor: Nome do provedor (maritaca, anthropic)
            modelo: Nome do modelo utilizado
            aplicacao: Nome da aplicação
            tag_funcionalidade: Tag para identificar a funcionalidade
            usuario: Usuário que realizou a chamada
        
        Returns:
            Decorator function
        
        Example:
            >>> @tracker.track_call(
            ...     provedor="anthropic",
            ...     modelo="claude-3-5-sonnet-20241022",
            ...     aplicacao="chat-app"
            ... )
            ... def processar_mensagem(mensagem):
            ...     response = client.messages.create(
            ...         model="claude-3-5-sonnet-20241022",
            ...         messages=[{"role": "user", "content": mensagem}]
            ...     )
            ...     return response
        """
        def decorator(func):
            def wrapper(*args, **kwargs):
                start_time = time.time()
                erro = False
                mensagem_erro = None
                result = None
                
                try:
                    result = func(*args, **kwargs)
                    return result
                except Exception as e:
                    erro = True
                    mensagem_erro = str(e)
                    raise
                finally:
                    tempo_resposta_ms = int((time.time() - start_time) * 1000)
                    
                    # Extrair tokens da resposta
                    tokens_input = 0
                    tokens_output = 0
                    
                    # Tentar diferentes formatos de resposta
                    if result:
                        # Formato Anthropic
                        if hasattr(result, 'usage'):
                            tokens_input = getattr(result.usage, 'input_tokens', 0)
                            tokens_output = getattr(result.usage, 'output_tokens', 0)
                        # Formato OpenAI-like (Maritaca pode usar)
                        elif isinstance(result, dict) and 'usage' in result:
                            usage = result['usage']
                            tokens_input = usage.get('prompt_tokens', 0)
                            tokens_output = usage.get('completion_tokens', 0)
                    
                    try:
                        if tokens_input > 0 or tokens_output > 0:
                            self.registrar_chamada(
                                provedor=provedor,
                                modelo=modelo,
                                tokens_input=tokens_input,
                                tokens_output=tokens_output,
                                aplicacao=aplicacao,
                                tag_funcionalidade=tag_funcionalidade,
                                usuario=usuario,
                                tempo_resposta_ms=tempo_resposta_ms,
                                erro=erro,
                                mensagem_erro=mensagem_erro
                            )
                        else:
                            logger.warning(
                                "Não foi possível extrair tokens da resposta. "
                                "Registro não criado."
                            )
                    except Exception as e:
                        logger.error(f"Erro ao registrar chamada: {str(e)}")
            
            return wrapper
        return decorator
    
    def obter_estatisticas_aplicacao(self,
                                     aplicacao: str,
                                     data_inicio: Optional[datetime] = None,
                                     data_fim: Optional[datetime] = None) -> Dict[str, Any]:
        """
        Obtém estatísticas de uso para uma aplicação.
        
        Args:
            aplicacao: Nome da aplicação
            data_inicio: Data inicial do período
            data_fim: Data final do período
        
        Returns:
            Dicionário com estatísticas de uso e custo
        
        Example:
            >>> from datetime import datetime, timedelta
            >>> data_inicio = datetime.now() - timedelta(days=7)
            >>> stats = tracker.obter_estatisticas_aplicacao(
            ...     "meu-app",
            ...     data_inicio=data_inicio
            ... )
            >>> print(f"Custo total: R$ {stats['custo_total']:.2f}")
        """
        return self.db.buscar_custos_por_aplicacao(aplicacao, data_inicio, data_fim)
    
    def listar_modelos(self, provedor: Optional[str] = None) -> List[ModeloIA]:
        """
        Lista modelos disponíveis no banco de dados.
        
        Args:
            provedor: Filtrar por provedor específico (opcional)
        
        Returns:
            Lista de ModeloIA
        
        Example:
            >>> # Listar todos os modelos
            >>> todos = tracker.listar_modelos()
            >>> 
            >>> # Listar apenas Anthropic
            >>> anthropic = tracker.listar_modelos(provedor="anthropic")
            >>> for m in anthropic:
            ...     print(f"{m.nome_modelo}: R$ {m.custo_input_por_1k}/1k")
        """
        return self.db.listar_modelos(provedor=provedor)