import os
from typing import TypedDict, List, Optional, Dict, Any

from langchain_openai import ChatOpenAI
from langchain_tavily import TavilySearchResults
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode
from langgraph.prebuilt import create_react_agent
from dotenv import load_dotenv

load_dotenv()

# ---------- 1. Состояние ----------
class BriefState(TypedDict):
    topic: str
    outline: List[str] | None
    step_index: int
    notes: List[str]
    final_brief: str | None

# ---------- 2. Инструменты ----------
search_tool = TavilySearchResults()

# ---------- 3. Узлы ----------
def outline_node(state: BriefState) -> Dict[str, Any]:
    """Создаёт план из 4–5 пунктов по теме."""
    llm = ChatOpenAI(model="gpt-4o-mini")
    prompt = (
        f"Составь план исследования по теме: {state['topic']}. "
        "План должен содержать 4–5 пунктов, каждый пункт – заголовок без описания."
    )
    plan = llm.invoke(prompt).content.strip()
    outline = [line.strip("- ").strip() for line in plan.splitlines() if line.strip()]
    return {"outline": outline, "step_index": 0, "notes": []}

def research_step_node(state: BriefState) -> Dict[str, Any]:
    """Для текущего пункта плана делает поиск и формирует заметку."""
    llm = ChatOpenAI(model="gpt-4o-mini")
    current_index = state["step_index"]
    if state["outline"] is None or current_index >= len(state["outline"]):
        return {"step_index": current_index, "notes": state["notes"]}

    topic = state["outline"][current_index]
    # Поиск
    search_query = f"{state['topic']} - {topic}"
    results = search_tool.invoke({"query": search_query})
    # Сводка
    summary_prompt = (
        f"Сделай краткую заметку (5–8 предложений) по теме '{topic}'. "
        f"Используй найденные источники: {results}"
    )
    note = llm.invoke(summary_prompt).content.strip()
    notes = state["notes"] + [f"**{topic}**\n{note}"]
    return {"step_index": current_index + 1, "notes": notes}

def synthesize_node(state: BriefState) -> Dict[str, Any]:
    """Объединяет все заметки в связный бриф."""
    llm = ChatOpenAI(model="gpt-4o-mini")
    combined = "\n\n".join(state["notes"])
    prompt = (
        f"Собери из следующих заметок связный исследовательский бриф (½–1 страница):\n{combined}"
    )
    brief = llm.invoke(prompt).content.strip()
    return {"final_brief": brief}

# ---------- 4. Граф ----------
def build_graph() -> StateGraph[BriefState]:
    graph = StateGraph(BriefState)

    # Добавляем узлы
    graph.add_node("outline", outline_node)
    graph.add_node("research_step", research_step_node)
    graph.add_node("synthesize", synthesize_node)

    # Переходы
    graph.set_entry_point("outline")
    graph.add_edge("outline", "research_step")
    graph.add_conditional_edges(
        "research_step",
        lambda state: "synthesize" if state["step_index"] >= len(state["outline"] or []) else "research_step",
        {
            "synthesize": "synthesize",
            "research_step": "research_step",
        },
    )
    graph.add_edge("synthesize", END)

    return graph

# ---------- 5. Запуск ----------
if __name__ == "__main__":
    topic = input("Введите тему исследования: ").strip()
    initial_state: BriefState = {
        "topic": topic,
        "outline": None,
        "step_index": 0,
        "notes": [],
        "final_brief": None,
    }

    graph = build_graph()
    chain = graph.compile()

    result = chain.invoke(initial_state)
    print("\n=== Итоговый исследовательский бриф ===\n")
    print(result["final_brief"])