"""Gerenciador de conexões e operações com banco de dados."""

import psycopg2
from psycopg2.extras import RealDictCursor
from typing import List, Optional, Dict, Any
from contextlib import contextmanager
from decimal import Decimal
import logging

from .models import ModeloIA, UsoToken
from .exceptions import DatabaseError, ModelNotFoundError

logger = logging.getLogger(__name__)


class DatabaseManager:
    """
    Gerenciador de conexões e operações com PostgreSQL.
    
    Esta classe gerencia todas as interações com o banco de dados,
    incluindo operações CRUD para modelos e registros de uso.
    """
    
    def __init__(self, connection_string: str):
        """
        Inicializa o gerenciador de banco de dados.
        
        Args:
            connection_string: String de conexão PostgreSQL
                              (ex: "postgresql://user:pass@localhost/dbname")
        
        Raises:
            DatabaseError: Se houver erro na conexão com o banco
        """
        self.connection_string = connection_string
        self._test_connection()
    
    def _test_connection(self) -> None:
        """Testa a conexão com o banco de dados."""
        try:
            with self._get_connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("SELECT 1")
        except Exception as e:
            raise DatabaseError(f"Erro ao conectar ao banco de dados: {str(e)}")
    
    @contextmanager
    def _get_connection(self):
        """
        Context manager para conexões com o banco.
        
        Yields:
            psycopg2.connection: Conexão com o banco de dados
        """
        conn = None
        try:
            conn = psycopg2.connect(self.connection_string)
            yield conn
            conn.commit()
        except Exception as e:
            if conn:
                conn.rollback()
            raise DatabaseError(f"Erro na operação do banco: {str(e)}")
        finally:
            if conn:
                conn.close()
    
    @staticmethod
    def _convert_decimal_to_float(row_dict: Dict) -> Dict:
        """
        Converte valores Decimal para float no dicionário.
        
        Args:
            row_dict: Dicionário com dados do banco
        
        Returns:
            Dicionário com Decimals convertidos para float
        """
        converted = {}
        for key, value in row_dict.items():
            if isinstance(value, Decimal):
                converted[key] = float(value)
            else:
                converted[key] = value
        return converted
    
    def buscar_modelo(self, provedor: str, nome_modelo: str) -> Optional[ModeloIA]:
        """
        Busca um modelo específico no banco de dados.
        
        Args:
            provedor: Nome do provedor (ex: 'maritaca', 'anthropic')
            nome_modelo: Nome do modelo
        
        Returns:
            ModeloIA se encontrado, None caso contrário
        
        Raises:
            DatabaseError: Se houver erro na consulta
        """
        query = """
            SELECT id_modelo, provedor, nome_modelo, custo_input_por_1k,
                   custo_output_por_1k, ativo, criado_em
            FROM dim_modelo_ia
            WHERE provedor = %s AND nome_modelo = %s AND ativo = true
        """
        
        with self._get_connection() as conn:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(query, (provedor, nome_modelo))
                row = cur.fetchone()
                
                if row:
                    row_dict = self._convert_decimal_to_float(dict(row))
                    return ModeloIA(**row_dict)
                return None
    
    def listar_modelos(self, provedor: Optional[str] = None, 
                      apenas_ativos: bool = True) -> List[ModeloIA]:
        """
        Lista modelos do banco de dados.
        
        Args:
            provedor: Filtrar por provedor específico (opcional)
            apenas_ativos: Se True, retorna apenas modelos ativos
        
        Returns:
            Lista de ModeloIA
        
        Raises:
            DatabaseError: Se houver erro na consulta
        """
        query = """
            SELECT id_modelo, provedor, nome_modelo, custo_input_por_1k,
                   custo_output_por_1k, ativo, criado_em
            FROM dim_modelo_ia
            WHERE 1=1
        """
        params = []
        
        if provedor:
            query += " AND provedor = %s"
            params.append(provedor)
        
        if apenas_ativos:
            query += " AND ativo = true"
        
        query += " ORDER BY provedor, nome_modelo"
        
        with self._get_connection() as conn:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(query, params)
                rows = cur.fetchall()
                modelos = []
                for row in rows:
                    row_dict = self._convert_decimal_to_float(dict(row))
                    modelos.append(ModeloIA(**row_dict))
                return modelos
    
    def inserir_ou_atualizar_modelo(self, modelo: ModeloIA) -> int:
        """
        Insere um novo modelo ou atualiza se já existir.
        
        Args:
            modelo: Objeto ModeloIA a ser inserido/atualizado
        
        Returns:
            ID do modelo inserido/atualizado
        
        Raises:
            DatabaseError: Se houver erro na operação
        """
        query = """
            INSERT INTO dim_modelo_ia 
                (provedor, nome_modelo, custo_input_por_1k, custo_output_por_1k, ativo)
            VALUES (%s, %s, %s, %s, %s)
            ON CONFLICT (provedor, nome_modelo) 
            DO UPDATE SET
                custo_input_por_1k = EXCLUDED.custo_input_por_1k,
                custo_output_por_1k = EXCLUDED.custo_output_por_1k,
                ativo = EXCLUDED.ativo
            RETURNING id_modelo
        """
        
        with self._get_connection() as conn:
            with conn.cursor() as cur:
                cur.execute(query, (
                    modelo.provedor,
                    modelo.nome_modelo,
                    modelo.custo_input_por_1k,
                    modelo.custo_output_por_1k,
                    modelo.ativo
                ))
                id_modelo = cur.fetchone()[0]
                logger.info(f"Modelo {modelo.provedor}/{modelo.nome_modelo} "
                          f"inserido/atualizado com ID {id_modelo}")
                return id_modelo
    
    def registrar_uso(self, uso: UsoToken) -> int:
        """
        Registra um uso de tokens no banco de dados.
        
        Args:
            uso: Objeto UsoToken com os dados do uso
        
        Returns:
            ID do registro de uso criado
        
        Raises:
            DatabaseError: Se houver erro na inserção
        """
        query = """
            INSERT INTO fato_uso_tokens (
                id_modelo, tokens_input, tokens_output,
                custo_input, custo_output, tag_funcionalidade,
                aplicacao, usuario, tempo_resposta_ms,
                erro, mensagem_erro, timestamp_uso
            ) VALUES (
                %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
            ) RETURNING id_uso
        """
        
        with self._get_connection() as conn:
            with conn.cursor() as cur:
                cur.execute(query, (
                    uso.id_modelo,
                    uso.tokens_input,
                    uso.tokens_output,
                    uso.custo_input,
                    uso.custo_output,
                    uso.tag_funcionalidade,
                    uso.aplicacao,
                    uso.usuario,
                    uso.tempo_resposta_ms,
                    uso.erro,
                    uso.mensagem_erro,
                    uso.timestamp_uso
                ))
                id_uso = cur.fetchone()[0]
                logger.info(f"Uso registrado com ID {id_uso} - "
                          f"Custo total: R$ {uso.custo_total:.6f}")
                return id_uso
    
    def buscar_custos_por_aplicacao(self, aplicacao: str, 
                                    data_inicio: Optional[Any] = None,
                                    data_fim: Optional[Any] = None) -> Dict[str, Any]:
        """
        Busca estatísticas de custo para uma aplicação específica.
        
        Args:
            aplicacao: Nome da aplicação
            data_inicio: Data inicial do período (opcional)
            data_fim: Data final do período (opcional)
        
        Returns:
            Dicionário com estatísticas de uso e custo
        
        Raises:
            DatabaseError: Se houver erro na consulta
        """
        query = """
            SELECT 
                COUNT(*) as total_chamadas,
                SUM(tokens_input) as total_tokens_input,
                SUM(tokens_output) as total_tokens_output,
                SUM(tokens_total) as total_tokens,
                SUM(custo_total) as custo_total,
                AVG(custo_total) as custo_medio,
                AVG(tempo_resposta_ms) as tempo_medio_resposta_ms,
                SUM(CASE WHEN erro = true THEN 1 ELSE 0 END) as total_erros
            FROM fato_uso_tokens
            WHERE aplicacao = %s
        """
        params = [aplicacao]
        
        if data_inicio:
            query += " AND timestamp_uso >= %s"
            params.append(data_inicio)
        
        if data_fim:
            query += " AND timestamp_uso <= %s"
            params.append(data_fim)
        
        with self._get_connection() as conn:
            with conn.cursor(cursor_factory=RealDictCursor) as cur:
                cur.execute(query, params)
                resultado = cur.fetchone()
                if resultado:
                    return self._convert_decimal_to_float(dict(resultado))
                return {}