import os
import base64
import json
from typing import Dict, Any, Optional

from langchain_core.messages import HumanMessage, AIMessage, SystemMessage, AIMessageChunk
from react_agent import build_react_agent
# Убедитесь, что graph_history_manager импортирован. 
# Если код в workflows.py, то graph_history_manager создается/импортируется там.

from graph_history_manager import GraphHistoryManager

graph_history_manager = GraphHistoryManager()


def run_agent_streaming(graph_id: str,
                        user_node_id: str,
                        assistant_node_id: str,
                        system_prompt: Optional[str] = None,
                        model: Optional[str] = None,
                        cache_folder: Optional[str] = None,
                        temperature: float = 0.7,
                        max_tokens: Optional[int] = None,
                        agency_mode: bool = False,
                        obsidian_settings: dict = None):
    """
    Основной генератор: Запускает LangGraph / ReAct агента.
    Поддерживает мультимодальность, чтение файлов из Vault и исполнение Python-тулов.
    Отдает словари, которые api.py запакует в SSE-события.
    """
    try:
        loaded_graph_data = graph_history_manager.get_graph(
            graph_id, target_node_id=user_node_id)

        if not loaded_graph_data:
            yield {"type": "error", "error": "Граф не найден."}
            return

        messages = loaded_graph_data.get("messages", [])

        # 2. Формируем контекст LangChain с учетом файлов и изображений (вложений)
        messages_for_llm = []
        if system_prompt:
            messages_for_llm.append(SystemMessage(content=system_prompt))

        # Обрабатываем вложения
        for i, msg in enumerate(messages):
            is_last = (i == len(messages) - 1)
            role = msg.get("role")

            if role == "user":
                text_content = msg.get("content", "")
                attachments = msg.get("attachments", [])

                # Проверяем, есть ли мультимодальный контент
                has_multimodal = False
                content_parts = [{"type": "text", "text": text_content}]

                # Обрабатываем вложения
                if attachments:
                    for att in attachments:
                        # Прикрепляем файл только если он перманентный (@@/##) или это последнее сообщение
                        if att.get("permanent") or is_last:
                            attachment_content = prepare_attachment_for_llm(
                                att, cache_folder)

                            # Если это изображение, то это мультимодальный контент
                            if attachment_content.get("type") == "image_url":
                                has_multimodal = True

                            content_parts.append(attachment_content)

                # Если есть картинки, шлем list of dicts. Иначе склеиваем в один текст.
                if has_multimodal:
                    # Мультимодальное сообщение
                    messages_for_llm.append(
                        HumanMessage(content=content_parts))
                else:
                    # Текстовое сообщение (объединяем все текстовые части)
                    full_text = text_content
                    for part in content_parts[
                            1:]:  # Пропускаем первую часть (основной текст)
                        if part.get("type") == "text":
                            full_text += part.get("text", "")
                    messages_for_llm.append(HumanMessage(content=full_text))

            elif role == "assistant":
                messages_for_llm.append(AIMessage(content=msg.get("content", "")))

        # Локальное хранилище для отладки текущего запуска
        current_run_debug = []
        
        def capture_debug(req, res):
            current_run_debug.append({"req": req, "res": res})

        # 3. Создаем агента. Если agency_mode == False, он соберет 0 тулов и отработает как обычная LLM.
        agent_executor = build_react_agent(
            model_name=model, 
            temperature=temperature, 
            max_tokens=max_tokens, 
            agency_mode=agency_mode, 
            obsidian_settings=obsidian_settings or {},
            debug_callback=capture_debug
        )
        
        # 4. Выполняем цикл с потоковой передачей событий
        step_counter = 0

        # LangGraph Stream Mode 'messages': отдает чанки токенов и вызовы функций
        for chunk, metadata in agent_executor.stream({"messages": messages_for_llm}, stream_mode="messages"):
            
            if isinstance(chunk, AIMessageChunk):
                # Агент запросил вызов инструмента (Формирование JSON аргументов)
                if hasattr(chunk, "tool_call_chunks") and chunk.tool_call_chunks:
                    for tc in chunk.tool_call_chunks:
                        if "name" in tc and tc["name"]:
                            step_counter += 1
                            yield {
                                "type": "tool_start", 
                                "name": tc["name"], 
                                "step": step_counter,
                                "node_id": assistant_node_id
                            }

                # Обычный стриминг текста финального ответа
                if chunk.content:
                    yield {
                        "type": "chunk", 
                        "content": chunk.content,
                        "node_id": assistant_node_id
                    }
            
            # Тул отработал и принес ответ
            elif hasattr(chunk, "type") and chunk.type == "tool":
                yield {
                    "type": "tool_end", 
                    "name": chunk.name, 
                    "step": step_counter,
                    "node_id": assistant_node_id
                }

        # В САМОМ КОНЦЕ, после того как агент закончил работу (вышел из цикла)
        if current_run_debug and obsidian_settings.get('showDebugInfo'):
            debug_section = "\n\n---\n" + "\n\n---\n".join([
                f"**JSON Request:**\n```json\n{d['req']}\n```\n**Result:**\n```text\n{d['res']}\n```"
                for d in current_run_debug
            ])
            yield {
                "type": "chunk", 
                "content": debug_section, 
                "node_id": assistant_node_id
            }

    except Exception as e:
        print(f"Ошибка при стриминге: {e}")
        import traceback
        traceback.print_exc()
        yield {"type": "error", "error": str(e)}


def prepare_attachment_for_llm(attachment: Dict[str, Any],
                               cache_folder: str) -> Dict[str, Any]:
    """Подготавливает вложение для отправки в LLM."""
    mime_type = attachment.get("mimeType", "application/octet-stream")
    cached_path = attachment.get("cachedPath")
    att_cache_folder = attachment.get("cacheFolder", cache_folder)  # ИЗМЕНЕНО
    source_path = attachment.get("sourcePath")
    att_type = attachment.get("type")
    name = attachment.get("name")

    # Получаем полный путь к кешированному файлу
    full_cached_path = os.path.join(att_cache_folder, cached_path)

    # Проверяем существование кеша
    if not os.path.exists(full_cached_path):
        print(f"❌ Кеш не найден: {full_cached_path}")
        return {"type": "text", "text": f"\n[Файл {name} недоступен]\n"}

    # Проверяем тип файла
    if mime_type.startswith('image/'):
        # Для изображений отправляем как base64
        with open(full_cached_path, 'rb') as f:
            image_data = base64.b64encode(f.read()).decode('utf-8')

        return {
            "type": "image_url",
            "image_url": {
                "url": f"data:{mime_type};base64,{image_data}"
            }
        }
    elif mime_type.startswith('text/') or \
         mime_type in ['application/json', 'application/xml'] or \
         mime_type.startswith('application/x-') or \
         mime_type == 'application/octet-stream': # TODO: Спорно, но пока вводим - фронтенд не смог сопоставить расширению mime_type 
        # неразрывно связано с методом getMimeType() на фронтенде
        # Для текстовых файлов (или потенциально текстовых с неизвестным типом)
        try:
            with open(full_cached_path, 'r', encoding='utf-8') as f:
                content = f.read()
        except UnicodeDecodeError:
            # Если не получилось прочитать как текст
            return {
                "type": "text",
                "text": f"\n[Файл {name} не может быть прочитан как текст]\n"
            }

        att_type_ru = "промпта" if att_type == "prompt" else "файла"
        header = f"\n\n--- Содержимое {att_type_ru}: {name} ---\n"
        footer = f"\n--- Конец {att_type_ru}: {name} ---\n\n"

        return {"type": "text", "text": header + content.strip() + footer}
    else:
        # Для других типов файлов возвращаем информацию
        return {
            "type":
            "text",
            "text":
            f"\n[Файл {name} ({mime_type}) не может быть отображен в чате]\n"
        }


def restore_cache_from_source(source_path: str, att_type: str,
                              cache_folder: str, cached_path: str) -> str:
    """Восстанавливает кеш из исходного файла."""
    full_cached_path = os.path.join(cache_folder, cached_path)

    # Создаем папку кеша, если не существует
    os.makedirs(cache_folder, exist_ok=True)

    if att_type == 'external':
        # Для внешних файлов source_path - абсолютный
        if not os.path.exists(source_path):
            raise FileNotFoundError(f"Исходный файл не найден: {source_path}")

        import shutil
        shutil.copy2(source_path, full_cached_path)
    else:
        # Для внутренних файлов и промптов нужно получить полный путь
        # Это требует доступа к vault, что на бэкенде недоступно
        # Поэтому это нужно делать на фронтенде
        raise NotImplementedError(
            "Восстановление кеша для внутренних файлов должно происходить на фронтенде"
        )

    return full_cached_path


def format_attachment_for_llm(attachment: Dict[str, Any]) -> str:
    """Форматирует вложение для вставки в сообщение LLM."""
    att_type = "промпта" if attachment["type"] == "prompt" else "файла"
    header = f"\n\n--- Содержимое {att_type}: {attachment['name']} ---\n"
    footer = f"\n--- Конец {att_type}: {attachment['name']} ---\n\n"
    return header + attachment.get("content", "").strip() + footer

