import os
import asyncio
from typing import TypedDict, Annotated, List

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from langchain.tools import tool
from langchain_tavily import TavilySearchResults
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

# ----------------- LLM -----------------
llm = ChatOpenAI(
    model="openai/gpt-oss-20b:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
)

# ----------------- Backend -----------------
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# ----------------- Tavily Tool -----------------
search_tool = TavilySearchResults(max_results=3)

# ----------------- State -----------------
class BriefState(TypedDict):
    topic: str
    outline: List[str] | None
    step_index: int
    notes: List[str]
    final_brief: str | None

# ----------------- Nodes -----------------
async def outline_node(state: BriefState) -> BriefState:
    prompt = f"""Create a concise outline of 4–5 research steps for the topic: {state['topic']}. Return a JSON array of strings."""
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    # Parse JSON array
    import json
    try:
        outline = json.loads(response.content)
        if not isinstance(outline, list):
            raise ValueError
    except Exception:
        outline = ["Step 1: Define scope", "Step 2: Search", "Step 3: Analyze", "Step 4: Summarize"]
    state.update(outline=outline, step_index=0, notes=[], final_brief=None)
    return state

async def research_step_node(state: BriefState) -> BriefState:
    step = state['outline'][state['step_index']]
    # Perform web search via Tavily
    results = await search_tool.ainvoke(step)
    # Summarize results into 5–8 sentences
    summary_prompt = f"""Summarize the following search results into 5–8 concise sentences for the research step: {step}.

Results:
{results}"""
    summary = await llm.ainvoke([HumanMessage(content=summary_prompt)])
    state['notes'].append(f"{step}\n{summary.content}")
    state['step_index'] += 1
    return state

async def synthesize_node(state: BriefState) -> BriefState:
    # Combine notes into a coherent brief with headings
    heading_prompt = """Combine the following notes into a ½–1 page research brief. Use the step titles as headings and write in a clear, academic style.

Notes:
""" + "\n\n".join(state['notes'])
    brief = await llm.ainvoke([HumanMessage(content=heading_prompt)])
    state['final_brief'] = brief.content
    return state

# ----------------- Graph -----------------
graph = StateGraph(BriefState)
graph.add_node("outline", outline_node)
graph.add_node("research_step", research_step_node)
graph.add_node("synthesize", synthesize_node)

# Entry point
graph.set_entry_point("outline")

# Conditional edges
graph.add_conditional_edges(
    "outline",
    lambda _: "research_step",
)

graph.add_conditional_edges(
    "research_step",
    lambda state: "synthesize" if state['step_index'] >= len(state['outline']) else "research_step",
)

graph.add_edge("synthesize", END)

app = graph.compile()

# ----------------- DeepAgent -----------------
agent = create_deep_agent(
    model=llm,
    tools=[search_tool],
    backend=backend,
    system_prompt="You are a research assistant that builds a brief based on a topic.",
)

# ----------------- Main -----------------
async def main():
    topic = os.getenv("DEFAULT_TOPIC", "Как студенту безопасно подключать MCP к LangChain")
    # Run graph to get outline and notes
    state = await app.ainvoke({"topic": topic, "outline": None, "step_index": 0, "notes": [], "final_brief": None})
    # Print outline
    print("\n=== Outline ===")
    for i, step in enumerate(state['outline'], 1):
        print(f"{i}. {step}")
    # Print each research step note
    for i, note in enumerate(state['notes'], 1):
        print(f"\n[Step {i}] {note.splitlines()[0]}")
        print(note.splitlines()[1])
    # Print final brief
    print("\n=== Final Brief ===")
    print(state['final_brief'])

if __name__ == "__main__":
    asyncio.run(main())