"""Модуль для логирования генераций и метрик в PostgreSQL базу данных."""
import json
import logging
import os
import uuid
from typing import Dict, List, Optional
import psycopg2

logger = logging.getLogger(__name__)


class DBLogger:
    """Класс для логирования данных о генерациях в PostgreSQL."""

    def __init__(self) -> None:
        """Инициализирует соединение с базой данных PostgreSQL."""
        self.conn = None
        self.db_params = {
            'dbname': os.getenv('POSTGRES_DB'),
            'user': os.getenv('POSTGRES_USER'),
            'password': os.getenv('POSTGRES_PASSWORD'),
            'host': os.getenv('POSTGRES_HOST'),
            'port': os.getenv('POSTGRES_PORT'),
        }

    def connect(self) -> None:
        """Устанавливает соединение с базой данных."""
        try:
            self.conn = psycopg2.connect(**self.db_params)
            logger.info('Подключение к PostgreSQL установлено')
        except Exception as e:
            logger.error(f'Ошибка подключения к PostgreSQL: {str(e)}')
            self.conn = None

    def _ensure_connection(self) -> bool:
        """
        Проверяет и восстанавливает соединение при необходимости.
        
        Returns:
            bool: True если соединение активно, False при ошибке
        """
        try:
            # Шаг 1: Проверка на None
            if self.conn is None:
                logger.info("Соединение отсутствует, устанавливаем новое...")
                self.connect()
                return self.conn is not None
            
            # Шаг 2: Проверка состояния соединения
            if self.conn.closed != 0:
                logger.warning(f"Соединение закрыто (closed={self.conn.closed}), переподключаемся...")
                self.conn = None
                self.connect()
                return self.conn is not None
            
            # Шаг 3: Проверка активности соединения
            try:
                with self.conn.cursor() as cursor:
                    cursor.execute("SELECT 1")
                    result = cursor.fetchone()
                    if result and result[0] == 1:
                        logger.debug("Соединение активно и работает корректно")
                        return True
                    else:
                        logger.warning("Test query вернул неожиданный результат")
                        raise psycopg2.OperationalError("Invalid test query result")
                        
            except (psycopg2.OperationalError, psycopg2.InterfaceError) as e:
                logger.warning(f"Test query failed: {e}, переподключаемся...")
                self.conn = None
                self.connect()
                return self.conn is not None
                
        except Exception as e:
            logger.error(f"Критическая ошибка при проверке соединения: {e}")
            self.conn = None
            return False
    
    def disconnect(self) -> None:
        """Закрывает соединение с базой данных."""
        if self.conn:
            self.conn.close()
            self.conn = None
            logger.info('Соединение с PostgreSQL закрыто')

    def log_generation_request(
        self,
        document_type: str,
        project: str,
        page_name: str,
        attached_link : str
    ) -> Optional[str]:
        """Логирует основную информацию о запросе на генерацию.
        
        Args:
            document_type: Тип генерируемого документа
            project: Название проекта
            page_name: Название страницы
            
        Returns:
            request_id: UUID запроса или None в случае ошибки
        """
        if not self._ensure_connection():
            logger.error("Не удалось установить соединение с БД для логирования запроса")
            return None

        request_id = str(uuid.uuid4())
        
        try:
            with self.conn.cursor() as cursor:
                cursor.execute(
                    """
                    INSERT INTO generation_requests 
                    (id, document_type, project, page_name, attached_link)
                    VALUES (%s, %s, %s, %s, %s)
                    """,
                    (
                        request_id,
                        document_type,
                        project,
                        page_name,
                        attached_link if attached_link else None
                    ),
                )
                self.conn.commit()
                logger.info(f'Запрос на генерацию {request_id} успешно залогирован')
                return request_id
        except Exception as e:
            logger.error(f'Ошибка при логировании запроса: {str(e)}')
            self.conn.rollback()
            return None

    def log_generation_inputs(
        self,
        request_id: str,
        user_question: str,
        jira_text: str,
        previous_generations_summary: str,
        few_shot_examples: List[Dict[str, str]],
    ) -> bool:
        """Логирует входные данные для генерации.
        
        Args:
            request_id: UUID запроса
            user_question: Вопрос пользователя
            jira_text: Текст из Jira
            previous_generations_summary: Суммаризация предыдущих генераций (опционально)
            few_shot_examples: Примеры для few-shot обучения (опционально)
            
        Returns:
            bool: True в случае успеха, False в случае ошибки
        """
        if not self.conn:
            self.connect()
            if not self.conn:
                return False

        try:
            with self.conn.cursor() as cursor:
                cursor.execute(
                    """
                    INSERT INTO generation_inputs 
                    (request_id, user_question, jira_text, previous_generations_summary, few_shot_examples)
                    VALUES (%s, %s, %s, %s, %s)
                    """,
                    (
                        request_id,
                        user_question,
                        jira_text,
                        previous_generations_summary,
                        json.dumps(few_shot_examples, ensure_ascii=False) if few_shot_examples else None,
                    ),
                )
                self.conn.commit()
                logger.info(f'Входные данные для запроса {request_id} успешно залогированы')
                return True
        except Exception as e:
            logger.error(f'Ошибка при логировании входных данных: {str(e)}')
            self.conn.rollback()
            return False

    def log_generation_contexts(
        self,
        request_id: str,
        document_context: List[Dict[str, str]],
        table_context: str,
        graph_context: str,
    ) -> bool:
        """Логирует контекстные данные для генерации.
        
        Args:
            request_id: UUID запроса
            document_context: Контекст из документов (опционально)
            table_context: Контекст из таблиц (опционально)
            
        Returns:
            bool: True в случае успеха, False в случае ошибки
        """
        if not self.conn:
            self.connect()
            if not self.conn:
                return False

        try:
            with self.conn.cursor() as cursor:
                cursor.execute(
                    """
                    INSERT INTO generation_contexts 
                    (request_id, document_context, table_context, graph_context)
                    VALUES (%s, %s, %s, %s)
                    """,
                    (
                        request_id,
                        json.dumps(document_context, ensure_ascii=False) if document_context else None,
                        table_context if table_context else None,
                        graph_context if graph_context else None,
                    ),
                )
                self.conn.commit()
                logger.info(f'Контекстные данные для запроса {request_id} успешно залогированы')
                return True
        except Exception as e:
            logger.error(f'Ошибка при логировании контекстных данных: {str(e)}')
            self.conn.rollback()
            return False

    def log_generation_result(
        self,
        request_id: str,
        generated_answer: str,
        timings: Dict[str, float],
        metrics: Dict[str, float],
    ) -> bool:
        """Логирует результат генерации.
        
        Args:
            request_id: UUID запроса
            generated_answer: Сгенерированный ответ
            timings: Время генераций по этапам
            metrics: Метрики качества (опционально)
            
        Returns:
            bool: True в случае успеха, False в случае ошибки
        """
        if not self.conn:
            self.connect()
            if not self.conn:
                return False

        try:
            with self.conn.cursor() as cursor:
                cursor.execute(
                    """
                    INSERT INTO generation_results 
                    (request_id, generated_answer, timings, metrics)
                    VALUES (%s, %s, %s, %s)
                    """,
                    (
                        request_id,
                        generated_answer,
                        json.dumps(timings, ensure_ascii=False) if timings else None,
                        json.dumps(metrics, ensure_ascii=False) if metrics else None,
                    ),
                )
                self.conn.commit()
                logger.info(f'Результат генерации для запроса {request_id} успешно залогирован')
                return True
        except Exception as e:
            logger.error(f'Ошибка при логировании результата: {str(e)}')
            self.conn.rollback()
            return False

# Singleton instance
db_logger = DBLogger()
