"""
Определяет логику рабочих процессов агента,
используя LangGraph для управления состояниями
и переходами между узлами обработки.
"""

from typing import Dict, Any, Optional, Set
from langgraph.graph import StateGraph, END
from graph_history_manager import GraphHistoryManager
from models import AgentState
from nodes import parse_command_node, execute_command_node, call_llm_node, generate_images_node, analyze_image_node, get_meet_subtitles_node, get_teams_subtitles_node, summarize_history_node, handle_error_node, help_node
import uuid # Добавлено для генерации UUID
from llm_client import get_llm
# --- LangGraph: Построение графа ---

graph_history_manager = GraphHistoryManager()

workflow = StateGraph(AgentState)

# Добавляем узлы
workflow.add_node("parse_command", parse_command_node)
workflow.add_node("execute_command",
                  execute_command_node)  # Это по сути роутер
workflow.add_node("call_llm", call_llm_node)
workflow.add_node("generate_images", generate_images_node)
workflow.add_node("analyze_image", analyze_image_node)
workflow.add_node("get_meet_subtitles", get_meet_subtitles_node)
workflow.add_node("get_teams_subtitles", get_teams_subtitles_node)
workflow.add_node("summarize_history", summarize_history_node)
workflow.add_node("handle_error", handle_error_node)
workflow.add_node("help", help_node)

# Устанавливаем точку входа
workflow.set_entry_point("parse_command")


# Определяем маршрутизацию после parse_command
def route_command(state: AgentState) -> str:
    """Маршрутизирует выполнение на основе распознанной команды."""
    if state.error:
        return "handle_error"  # Если была ошибка парсинга команды
    elif state.command == "imagine":
        return "generate_images"
    elif state.command == "analyze":
        return "analyze_image"
    elif state.command == "subtitles_meet":
        return "get_meet_subtitles"
    elif state.command == "subtitles_teams":
        return "get_teams_subtitles"
    elif state.command == "summarize":
        return "summarize_history"
    elif state.command == "help":
        return "help"
    elif state.command == "chat":
        return "call_llm"
    elif state.command == "error":  # Команда была, но не известная
        return "handle_error"
    else:
        # Fallback на LLM, если команда не распознана или пуста (хотя parse_command должен был бы ее поймать)
        return "call_llm"


workflow.add_conditional_edges(
    "parse_command",  # Откуда
    route_command,  # Функция, определяющая куда идти
    {
        "generate_images": "generate_images",
        "analyze_image": "analyze_image",
        "get_meet_subtitles": "get_meet_subtitles",
        "get_teams_subtitles": "get_teams_subtitles",
        "summarize_history": "summarize_history",
        "help": "help",
        "call_llm": "call_llm",
        "handle_error":
        "handle_error",  # Для ошибок, найденных в parse_command
    },
)

# Все узлы операций (кроме ошибок) ведут к END после завершения
workflow.add_edge("call_llm", END)
workflow.add_edge("generate_images", END)
workflow.add_edge("analyze_image", END)
workflow.add_edge("get_meet_subtitles", END)
workflow.add_edge("get_teams_subtitles", END)
workflow.add_edge("summarize_history", END)
workflow.add_edge("help", END)
workflow.add_edge("handle_error",
                  END)  # Ошибка - это тоже конечный пункт для текущего запроса

# Компилируем граф
app = workflow.compile()


def run_agent(user_input: str,
              existing_graph_id: Optional[str] = None,
              parent_node_id: Optional[str] = None,
              system_prompt: Optional[str] = None,
              model: Optional[str] = None) -> Dict[str, Any]:
    """
    Запускает агента с заданным пользовательским вводом.
    Может продолжить существующий граф по graph_id.
    Возвращает обновленное состояние и ID графа.
    """
    
    # ---------------------------------------------------- Initial State Setup ----------------------------------------------------------------
    initial_state = AgentState(
        input=user_input, 
        parent_node_id=parent_node_id, 
        system_prompt=system_prompt,
        selected_model=model or "gemini-2.5-flash"
    )
        
    # Сохраняем исходные ID узлов и ребер для определения новых после выполнения графа
    original_node_ids: Set[str] = set()
    original_edge_ids: Set[str] = set()

    # Если есть существующий ID графа, загружаем его историю и структуру
    loaded_graph_data = None
    if existing_graph_id:
        # Получаем данные графа, включая сгенерированные сообщения до parent_node_id (либо current_node_id из БД)
        # Если parent_node_id передан, используем его для построения истории.
        loaded_graph_data = graph_history_manager.get_graph(
            existing_graph_id, target_node_id=parent_node_id)
        
        if loaded_graph_data:
            print(f"Продолжаю существующий граф {existing_graph_id}")
            # Загружаем узлы и ребра для продолжения графа
            initial_state.graph_nodes = loaded_graph_data.get("graph_nodes", [])
            initial_state.graph_edges = loaded_graph_data.get("graph_edges", [])
            
            # Сохраняем ID загруженных узлов и ребер
            original_node_ids = {node['id'] for node in initial_state.graph_nodes}
            original_edge_ids = {edge['id'] for edge in initial_state.graph_edges}

            # Устанавливаем parent_node_id для нового входного узла
            # Если parent_node_id был передан в запросе, используем его,
            # иначе берем current_node_id из загруженного графа.
            initial_state.parent_node_id = parent_node_id or loaded_graph_data.get("current_node_id")
            
            # История чата для текущего раунда работы LLM формируется из загруженных сообщений.
            initial_state.temporary_chat_history = loaded_graph_data.get("messages", [])

    # Определяем graph_id, который будет использоваться для сохранения.
    # Если существующий ID не передан, GraphHistoryManager сгенерирует новый.
    graph_id_to_save = existing_graph_id


    # ---------------------------------------------------- run_agent ----------------------------------------------------------------
    # Запускаем граф
    result = app.invoke(initial_state)
    final_state = AgentState(**result)

    # ---------------------------------------------------- Save Graph ----------------------------------------------------------------
    # Определяем новые узлы и ребра для сохранения
    newly_added_nodes = [node for node in final_state.graph_nodes if node['id'] not in original_node_ids]
    newly_added_edges = [edge for edge in final_state.graph_edges if edge['id'] not in original_edge_ids]

    # Сохраняем обновленное состояние графа.
    # Обратите внимание: chat_history больше не сохраняется напрямую,
    # она является частью данных узлов в graph_nodes.
    final_graph_id = graph_history_manager.save_graph_changes(
        graph_id_to_save,
        newly_added_nodes,
        newly_added_edges,
        final_state.current_node_id
    )

    # ---------------------------------------------------- prepare response ----------------------------------------------------------------
    # Для `messages` в `response_data` используем `final_state.temporary_chat_history`,
    # так как она отражает только сообщения текущей ветки, добавленные в ходе этого выполнения.
    response_data = {
        "graph_id": final_graph_id,
        "messages": final_state.temporary_chat_history,
        "llm_response": final_state.llm_response,
        "image_urls": final_state.image_urls,
        "analysis_result": final_state.analysis_result,
        "subtitles": final_state.subtitles,
        "summarized_history_text": final_state.summarized_history_text,
        "error": final_state.error,
        "graph_visualization_data": {
            "nodes": final_state.graph_nodes,
            "edges": final_state.graph_edges,
            "current_node_id": final_state.current_node_id
        }
    }
    return response_data

def run_agent_streaming(graph_id: str, user_node_id: str, assistant_node_id: str, system_prompt: Optional[str] = None, model: Optional[str] = None):
    """
    Генератор для стримингового ответа LLM.
    Возвращает чанки контента по мере их получения.
    """
    try:
        # Загружаем граф и историю до user_node_id
        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", [])
        
        # Формируем сообщения для LLM
        from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
        
        messages_for_llm = []
        if system_prompt:
            messages_for_llm.append(SystemMessage(content=system_prompt))
        
        for msg in messages:
            if msg["role"] == "user":
                messages_for_llm.append(HumanMessage(content=msg["content"]))
            elif msg["role"] == "assistant":
                messages_for_llm.append(AIMessage(content=msg["content"]))

        # Получаем LLM и стримим ответ
        llm = get_llm(model or "gemini-2.5-flash")
        
        for chunk in llm.stream2(messages_for_llm):
            content = chunk.content if hasattr(chunk, 'content') else str(chunk)
            yield {"type": "chunk", "content": content}
            
    except Exception as e:
        print(f"Ошибка при стриминге: {e}")
        yield {"type": "error", "error": str(e)}