import asyncio
import logging
from typing import List, Optional, Dict, Any
from ..handlers.handlers import get_global_var

from langchain.messages import SystemMessage
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser

logger = logging.getLogger(__name__)

class MemoryManager:
    def __init__(self):
        self.supabase_client = get_global_var("supabase_client")
        config = get_global_var("config")
        
        self.max_memory_messages = config.MAX_CONTEXT_MESSAGES
        self.min_memory_messages = config.HISTORY_MIN_MESSAGES if config.HISTORY_MIN_MESSAGES else 4
        self.token_limit = config.HISTORY_MAX_TOKENS if config.HISTORY_MAX_TOKENS else 5000
        
        self.chat_model = ChatOpenAI(model="gpt-5-mini", api_key=config.OPENAI_API_KEY) | StrOutputParser()
        
        # Словарь для отслеживания активных фоновых задач суммаризации по session_id
        self._active_summarization_tasks: Dict[str, asyncio.Task] = {}
        

    async def get_memory_messages(self, session_id: str) -> List[Dict[str, Any]]:
        """Возвращает историю сообщений в формате OpenAI (список словарей с role и content)"""
        chat_messages: List[Dict[str, Any]] = []
        
        logger.debug(f"[MemoryManager] Запрос истории для сессии {session_id}")
        session_info = await self.supabase_client.get_session_info(session_id)
        messages_len = session_info.get('messages_len', self.min_memory_messages)
        logger.info(
            f"[MemoryManager] Текущий messages_len={messages_len}, min={self.min_memory_messages}, max={self.max_memory_messages}"
        )
        
        stored_summary = session_info.get('summary', '')
        if stored_summary:
            logger.info("[MemoryManager] Найдена сохраненная суммаризация — добавляем в контекст")
            formatted_summary = self._format_summary(stored_summary)
            logger.debug(f"[MemoryManager] Форматированная суммаризация: {formatted_summary[:200]}...")
            chat_messages.append({"role": "system", "content": formatted_summary})
        else:
            logger.info("[MemoryManager] Суммаризация не найдена в сессии")
        
        messages = await self.supabase_client.get_chat_history(session_id, limit=messages_len)
        logger.info(f"[MemoryManager] Получено {len(messages)} сообщений из истории (limit={messages_len})")
        
        added_count = 0
        for msg in messages:
            if msg['role'] in ('user', 'assistant'):
                chat_messages.append({"role": msg['role'], "content": msg['content']})
                added_count += 1
            else:
                logger.debug(f"[MemoryManager] Пропущено сообщение с ролью: {msg['role']}")
        
        logger.info(f"[MemoryManager] Добавлено {added_count} сообщений из истории в контекст")
        
        total_tokens = self._count_tokens(chat_messages)
        logger.info(
            f"[MemoryManager] Подготовлено сообщений: {len(chat_messages)}, оценка токенов: {total_tokens}"
        )
        
        # Проверяем, нужно ли обрезать (исключаем суммаризацию из подсчета для проверки лимита сообщений)
        has_summary = chat_messages and chat_messages[0].get('role') == 'system' and self._is_summary_message(chat_messages[0])
        effective_messages_count = len(chat_messages) - (1 if has_summary else 0)
        
        # Если нужно обрезать, делаем быструю обрезку без суммаризации и запускаем фоновую задачу
        if total_tokens > self.token_limit or effective_messages_count > self.max_memory_messages - 1:
            logger.warning(
                f"[MemoryManager] История превышает лимиты (tokens>{self.token_limit} или messages>{effective_messages_count}>{self.max_memory_messages - 1}). Делаем быструю обрезку и запускаем фоновую суммаризацию."
            )
            # Быстрая обрезка: просто берем последние сообщения
            history_tail_size = max(self.min_memory_messages - 1, 0)
            if has_summary:
                # Сохраняем суммаризацию и берем хвост истории
                summary_msg = chat_messages[0]
                history_messages = chat_messages[1:]
                history_tail = history_messages[-history_tail_size:] if history_tail_size else []
                chat_messages = [summary_msg] + history_tail
            else:
                # Берем только хвост истории
                chat_messages = chat_messages[-history_tail_size:] if history_tail_size else []
            
            logger.info(f"[MemoryManager] Быстрая обрезка выполнена: {len(chat_messages)} сообщений")
            
            # Очищаем завершенные задачи из словаря (предотвращаем утечки памяти)
            completed_sessions = [
                sid for sid, task in self._active_summarization_tasks.items()
                if task.done()
            ]
            for sid in completed_sessions:
                del self._active_summarization_tasks[sid]
                logger.debug(f"[MemoryManager] Удалена завершенная задача для сессии {sid}")
            
            # Проверяем, не запущена ли уже фоновая суммаризация для этой сессии
            if session_id in self._active_summarization_tasks:
                existing_task = self._active_summarization_tasks[session_id]
                if not existing_task.done():
                    logger.info(
                        f"[MemoryManager] Фоновая суммаризация уже выполняется для сессии {session_id}, "
                        f"пропускаем запуск новой. Существующая задача обработает все сообщения из БД."
                    )
                    # Не запускаем новую задачу, существующая уже обработает все сообщения из БД
                else:
                    # Задача завершена (не должна быть здесь после очистки, но на всякий случай)
                    del self._active_summarization_tasks[session_id]
                    task = asyncio.create_task(self._background_summarize(session_id))
                    self._active_summarization_tasks[session_id] = task
                    logger.info(f"[MemoryManager] Запущена новая фоновая задача суммаризации для сессии {session_id}")
            else:
                # Нет активной задачи, запускаем новую
                task = asyncio.create_task(self._background_summarize(session_id))
                self._active_summarization_tasks[session_id] = task
                logger.info(f"[MemoryManager] Запущена фоновая задача суммаризации для сессии {session_id}")
            
        # Извлекаем суммаризацию из финального списка сообщений
        summary_for_storage = self._extract_summary(chat_messages)
        messages_len = self._calculate_messages_len(chat_messages)
        
        # Логируем финальное состояние перед возвратом
        logger.info(
            f"[MemoryManager] Финальное состояние: {len(chat_messages)} сообщений, messages_len={messages_len}, summary={'есть (' + str(len(summary_for_storage)) + ' символов)' if summary_for_storage else 'нет'}"
        )
        
        # Логируем первые несколько сообщений для отладки
        for idx, msg in enumerate(chat_messages[:3]):
            role = msg.get('role', 'unknown')
            content_preview = (msg.get('content', '') or '')[:100].replace("\n", " ")
            logger.info(f"[MemoryManager] Сообщение #{idx + 1}: {role} - {content_preview}...")
        
        # Обновляем сессию без ожидания (быстрое обновление)
        await self.supabase_client.update_session(
            session_id, {"messages_len": messages_len, "summary": summary_for_storage}
        )
        logger.info(f"[MemoryManager] Сессия {session_id} обновлена в БД")
        
        return chat_messages
    
    async def _background_summarize(self, session_id: str):
        """
        Фоновая задача для создания суммаризации истории диалога.
        Выполняется асинхронно, не блокирует основной поток.
        """
        try:
            logger.info(f"[MemoryManager] 🔄 Начало фоновой суммаризации для сессии {session_id}")
            
            # Получаем полную историю из БД для суммаризации
            session_info = await self.supabase_client.get_session_info(session_id)
            messages_len = session_info.get('messages_len', self.min_memory_messages)
            
            # Получаем больше сообщений для суммаризации (берем больше, чем обычно)
            full_messages = await self.supabase_client.get_chat_history(session_id, limit=messages_len * 2)
            
            # Формируем список сообщений для суммаризации
            messages_for_summary: List[Dict[str, Any]] = []
            existing_summary = session_info.get('summary', '')
            
            if existing_summary:
                formatted_summary = self._format_summary(existing_summary)
                messages_for_summary.append({"role": "system", "content": formatted_summary})
            
            for msg in full_messages:
                if msg['role'] in ('user', 'assistant'):
                    messages_for_summary.append({"role": msg['role'], "content": msg['content']})
            
            logger.info(f"[MemoryManager] 📚 Получено {len(messages_for_summary)} сообщений для суммаризации")
            
            # Создаем суммаризацию
            trimmed_messages = await self._trim_messages(messages_for_summary, existing_summary)
            
            # Извлекаем новую суммаризацию
            new_summary = self._extract_summary(trimmed_messages)
            new_messages_len = self._calculate_messages_len(trimmed_messages)
            
            # Сохраняем результат в БД
            await self.supabase_client.update_session(
                session_id, {"messages_len": new_messages_len, "summary": new_summary}
            )
            
            logger.info(
                f"[MemoryManager] ✅ Фоновая суммаризация завершена для сессии {session_id}. "
                f"Новая суммаризация: {len(new_summary)} символов, messages_len={new_messages_len}"
            )
            
        except Exception as e:
            logger.error(f"[MemoryManager] ❌ Ошибка в фоновой суммаризации для сессии {session_id}: {e}")
            logger.exception("Полный стек ошибки:")
        finally:
            # Удаляем задачу из словаря активных задач после завершения (успешного или с ошибкой)
            if session_id in self._active_summarization_tasks:
                del self._active_summarization_tasks[session_id]
                logger.debug(f"[MemoryManager] Задача суммаризации для сессии {session_id} удалена из активных")
    
    async def _trim_messages(
        self,
        messages: List[Dict[str, Any]],
        summary: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Обрезает сообщения, создавая суммаризацию. Работает со словарями OpenAI."""
        existing_summary = summary or self._extract_summary(messages)
        logger.info(f"[MemoryManager] Извлечена существующая суммаризация: {len(existing_summary) if existing_summary else 0} символов")
        
        messages_history: List[Dict[str, Any]] = (
            messages[1:] if messages and messages[0].get('role') == 'system' and self._is_summary_message(messages[0]) else messages
        )
        logger.info(f"[MemoryManager] История для суммаризации: {len(messages_history)} сообщений (исходно было {len(messages)})")

        summary_prompt = SystemMessage(
            content=f"""
            Ты — ассистент для суммаризации диалогов.
            Твоя задача — объединить уже существующую суммаризацию с новыми сообщениями из истории.

            Прошлая суммаризация: (может быть пустой)
            {existing_summary} (если ее не было не пиши об этом и если есть тоже писать не надо)

            Новые сообщения будут добавлены после этого system message.
            Сохраняй все ключевые мысли и действия, избегай потерь информации, делай связную компактную суммаризацию.
            
            ВАЖНО:
            - не надо ничего советовать, не надо ничего объяснять - твоя задача только суммаризация
            
            Возвращай только суммаризацию, без других комментариев.
            Максимальный размер суммаризации - 500 токенов.
            """
        )

        # Конвертируем словари OpenAI в LangChain сообщения для вызова модели
        from ..handlers.handlers import openai_messages_to_langchain, langchain_messages_to_openai
        langchain_history = openai_messages_to_langchain(messages_history)
        prompt_messages = [summary_prompt] + langchain_history
        
        logger.info(f"[MemoryManager] Отправляем {len(messages_history)} сообщений на суммаризацию")
        new_summary_text = await self.chat_model.ainvoke(prompt_messages)
        logger.info(f"[MemoryManager] Получена новая суммаризация: {len(new_summary_text)} символов")
        
        formatted_summary = self._format_summary(new_summary_text)
        new_summary = {"role": "system", "content": formatted_summary}
        logger.info(f"[MemoryManager] Форматированная суммаризация: {len(formatted_summary)} символов")

        history_tail_size = max(self.min_memory_messages - 1, 0)
        history_tail = messages_history[-history_tail_size:] if history_tail_size else []
        logger.info(
            f"[MemoryManager] Новая суммаризация создана. Возвращаем {1 + len(history_tail)} сообщений (summary + хвост)."
        )

        return [new_summary] + history_tail

    def _calculate_messages_len(
        self, messages: List[Dict[str, Any]]
    ) -> int:
        """Вычисляет messages_len для словарей OpenAI"""
        effective_len = len(messages)

        if messages and messages[0].get('role') == 'system' and self._is_summary_message(messages[0]):
            effective_len -= 1

        result = max(effective_len, 0) + 2
        logger.debug(
            f"[MemoryManager] Расчет messages_len: исходно={len(messages)}, эффективно={effective_len}, итог={result}"
        )
        return result

    def _count_tokens(self, messages: List[Dict[str, Any]]) -> int:
        """Подсчитывает токены для словарей OpenAI"""
        tokens = 0
        for msg in messages:
            content = msg.get('content', '')
            tokens += max(1, len(content) // 4)
        logger.debug(f"[MemoryManager] Оценка токенов для {len(messages)} сообщений: {tokens}")
        return tokens

    def _format_summary(self, summary: str) -> str:
        header = "## Суммаризация истории диалога (предыдущие сообщения до текущего момента):"
        summary_body = summary.strip()
        if not summary_body:
            return header
        return f"{header}\n{summary_body}"

    def _is_summary_message(self, message: Dict[str, Any]) -> bool:
        """Проверяет, является ли сообщение суммаризацией (для словаря OpenAI)."""
        if message.get('role') != 'system':
            return False
        content = message.get('content', '').strip()
        return content.startswith("## Суммаризация истории диалога")

    def _extract_summary(self, messages: List[Dict[str, Any]]) -> str:
        """Извлекает суммаризацию из первого сообщения (для словарей OpenAI)."""
        if not messages or messages[0].get('role') != 'system' or not self._is_summary_message(messages[0]):
            return ""

        content = messages[0].get('content', '').strip()
        header = "## Суммаризация истории диалога (предыдущие сообщения до текущего момента):"
        if content.startswith(header):
            summary_text = content[len(header):].lstrip()
            # Возвращаем пустую строку, если суммаризация содержит только заголовок
            return summary_text if summary_text else ""
        return content