import os
import asyncio
import sys
from typing import Dict, Optional
import logging
from lightrag import LightRAG
from psycopg2 import sql
import psycopg2
from lightrag.kg.shared_storage import initialize_pipeline_status
import yaml
import shutil

def load_config():
    with open('config.yaml', 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

config = load_config()

logger = logging.getLogger("lightrag")

def make_safe_db_name(project_id: str) -> str:
    """
    Создает безопасное имя базы данных из project_id
    PostgreSQL требует, чтобы имена БД содержали только буквы, цифры и подчеркивания
    """
    safe_id = project_id.replace('-', '_')
    # Убираем все остальные небезопасные символы
    # safe_id = ''.join(c if c.isalnum() or c == '_' else '_' for c in safe_id)

    return f"project_{safe_id}"

class ProjectRAG:
    """Полностью изолированный экземпляр LightRAG для проекта"""
    
    def __init__(self, project_id: str, 
                 db_host: str, db_port: int, db_user: str, db_password: str,
                 llm_func, embedding_func):
        self.project_id = project_id
        self.db_name = make_safe_db_name(project_id)
        self.db_host = db_host
        self.db_port = db_port
        self.db_user = db_user
        self.db_password = db_password
        self.working_dir = f"rag_storage/{self.db_name}"
        self.llm_func = llm_func
        self.embedding_func = embedding_func
        self.rag_instance = None
    
    async def initialize(self, create_if_not_exists=True):
        """Инициализирует RAG с опциональным созданием базы данных"""
        if create_if_not_exists:
            await self._ensure_database_exists()
        
        await self._initialize_in_clean_environment()
        
        return self.rag_instance
    
    async def delete_by_doc_id(self, doc_id: str):
        """Удаляет документ по ID в изолированном окружении"""
        if not self.rag_instance:
            await self.initialize()
        return await self.rag_instance.adelete_by_doc_id(doc_id)
    
    async def check_database_exists(self) -> bool:
        """Проверяет существование базы данных без её создания"""
        try:
            admin_db = os.getenv('POSTGRES_ADMIN_DATABASE')
            conn = psycopg2.connect(
                host=self.db_host,
                port=self.db_port,
                user=self.db_user,
                password=self.db_password,
                dbname=admin_db
            )
            conn.autocommit = True
            
            with conn.cursor() as cursor:
                cursor.execute("SELECT 1 FROM pg_database WHERE datname = %s",
                            (self.db_name,))
                exists = cursor.fetchone()
                
            conn.close()
            return exists is not None
            
        except Exception as e:
            logger.error(f"Ошибка при проверке существования базы PostgreSQL: {e}")
            return False

    async def _ensure_database_exists(self):
        """Создает базу данных если она не существует"""
        try:
            admin_db = os.getenv('POSTGRES_ADMIN_DATABASE')
            logger.info(f"🔍 Подключаемся к admin БД: {admin_db}, создаем БД: {self.db_name}")
            conn = psycopg2.connect(
                host=self.db_host,
                port=self.db_port,
                user=self.db_user,
                password=self.db_password,
                dbname=admin_db
            )
            conn.autocommit = True
            
            with conn.cursor() as cursor:
                cursor.execute("SELECT 1 FROM pg_database WHERE datname = %s",
                               (
                                   self.db_name,
                               ))
                exists = cursor.fetchone()
                
                if not exists:
                    logger.info(f"Создание базы данных PostgreSQL: {self.db_name}")
                    create_db_query = sql.SQL("CREATE DATABASE {}").format(
                        sql.Identifier(self.db_name)
                        )
                    cursor.execute(create_db_query)
                    logger.info(f"База данных PostgreSQL {self.db_name} успешно создана")
            
            conn.close()
            
        except Exception as e:
            logger.error(f"Ошибка при создании базы PostgreSQL: {e}")
            raise
    
    async def _initialize_in_clean_environment(self):
        """Инициализирует RAG с чистыми соединениями к нужной БД"""
        
        os.environ['POSTGRES_HOST'] = self.db_host
        os.environ['POSTGRES_PORT'] = str(self.db_port)
        os.environ['POSTGRES_USER'] = self.db_user
        os.environ['POSTGRES_PASSWORD'] = self.db_password
        os.environ['POSTGRES_DATABASE'] = self.db_name
        
        # ВАЖНО: Очищаем кэш модулей LightRAG, чтобы сбросить все соединения
        for module_name in list(sys.modules.keys()):
            if module_name.startswith('lightrag.kg.postgres'):
                del sys.modules[module_name]
        
        chunk_tokens = config.get('performance').get('chunk_token_size')
        self.rag_instance = LightRAG(
            working_dir=self.working_dir,
            chunk_token_size=chunk_tokens,
            chunk_overlap_token_size=chunk_tokens // 10,
            llm_model_func=self.llm_func,
            embedding_func=self.embedding_func,
            kv_storage="PGKVStorage",
            vector_storage="PGVectorStorage",
            graph_storage="PGGraphStorage",
            doc_status_storage="PGDocStatusStorage",
            namespace_prefix="",
            max_parallel_insert=config.get('performance').get('max_parallel_insert'),
            llm_model_max_async=config.get('performance').get('max_concurrent_llm_requests'),
            embedding_batch_num=config.get('performance').get('embedding_batch_num'),
            embedding_func_max_async=config.get('performance').get('embedding_func_max_async'),
        )
        
        for attempt in range(3):
            try:
                await self.rag_instance.initialize_storages()
                await initialize_pipeline_status()
                logger.info(f"RAG для проекта {self.project_id} успешно инициализирован в БД {self.db_name}")
                break
            except Exception as e:
                if "deadlock detected" in str(e) and attempt < 2:
                    logger.warning(f"Обнаружен deadlock, повторная попытка через 2 секунды...")
                    await asyncio.sleep(2 * (attempt + 1))
                else:
                    logger.error(f"Ошибка при инициализации хранилищ: {e}")
                    raise
    
    async def query(self, query, **kwargs):
        """Выполняет запрос в изолированном окружении"""
        if not self.rag_instance:
            await self.initialize()
        return await self.rag_instance.aquery(query, **kwargs)
    
    async def insert(self, input_text, ids=None, **kwargs):
        """Вставляет документ(ы) в изолированном окружении - поддерживает и одиночные и списки"""
        if not self.rag_instance:
            await self.initialize()
        
        if isinstance(input_text, list):
            logger.info(f"🚀 Батчевая вставка: {len(input_text)} документов")
        
        return await self.rag_instance.ainsert(input=input_text, ids=ids, **kwargs)
    
    async def drop_database(self):
        """Удаляет базу данных проекта"""
        try:
            logger.info(f"🗄️ Удаляем базу данных {self.db_name}")
            
            admin_db = os.getenv('POSTGRES_ADMIN_DATABASE')
            conn = psycopg2.connect(
                host=self.db_host,
                port=self.db_port,
                user=self.db_user,
                password=self.db_password,
                dbname=admin_db
            )
            conn.autocommit = True
            
            with conn.cursor() as cursor:
                cursor.execute("""
                    SELECT pg_terminate_backend(pid)
                    FROM pg_stat_activity
                    WHERE datname = %s AND pid <> pg_backend_pid()
                """, (self.db_name,))
                
                await asyncio.sleep(2)
                
                drop_db_query = sql.SQL("DROP DATABASE IF EXISTS {}").format(
                    sql.Identifier(self.db_name)
                )
                cursor.execute(drop_db_query)
                logger.info(f"✅ База данных {self.db_name} успешно удалена")
            
            conn.close()
            
        except Exception as e:
            logger.error(f"❌ Ошибка при удалении БД {self.db_name}: {e}")
            raise

    def cleanup_working_directory(self):
        """Удаляет рабочую директорию проекта"""
        try:
            if os.path.exists(self.working_dir):
                logger.info(f"🗂️ Удаляем рабочую директорию {self.working_dir}")
                shutil.rmtree(self.working_dir)
                logger.info(f"✅ Рабочая директория {self.working_dir} успешно удалена")
            else:
                logger.info(f"📁 Рабочая директория {self.working_dir} не найдена")
                
        except Exception as e:
            logger.error(f"❌ Ошибка при удалении рабочей директории {self.working_dir}: {e}")
            raise


class ProjectRAGManager:
    """Менеджер для управления изолированными RAG экземплярами"""
    
    def __init__(self, llm_func, embedding_func):
        self.project_instances: Dict[str, ProjectRAG] = {}
        self.llm_func = llm_func
        self.embedding_func = embedding_func
    
    async def get_rag_for_project(self, project_id: str, create_if_not_exists=False) -> Optional[ProjectRAG]:
        """
        Возвращает изолированный RAG для проекта
        
        Args:
            project_id: ID проекта
            create_if_not_exists: Создавать базу данных, если она не существует
        """
        if project_id in self.project_instances:
            return self.project_instances[project_id]
        
        db_host = os.getenv('POSTGRES_HOST')
        db_port = int(os.getenv('POSTGRES_PORT'))
        db_user = os.getenv('POSTGRES_USER')
        db_password = os.getenv('POSTGRES_PASSWORD')
        
        project_rag = ProjectRAG(
            project_id=project_id,
            db_host=db_host,
            db_port=db_port,
            db_user=db_user,
            db_password=db_password,
            llm_func=self.llm_func,
            embedding_func=self.embedding_func
        )
        
        if not create_if_not_exists:
            exists = await project_rag.check_database_exists()
            if not exists:
                logger.warning(f"Проект с ID {project_id} не существует")
                return None
        
        await project_rag.initialize(create_if_not_exists=create_if_not_exists)
        self.project_instances[project_id] = project_rag
        
        return project_rag
