import os
import json
import asyncio
from dotenv import load_dotenv
from typing import TypedDict, List, Optional

from langchain_openai import ChatOpenAI
from langchain_tavily import TavilySearchResults
from langchain.tools import tool

from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend

from langgraph.graph import StateGraph, START, END

# Load environment variables
load_dotenv()

# LLM configuration
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 for deepagents
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# TypedDict for graph state
class BriefState(TypedDict):
    topic: str
    outline: List[str] | None
    step_index: int
    notes: List[str]
    final_brief: str | None

# Tool: Web search using Tavily
@tool
def tavily_search(query: str) -> str:
    """Search the web for the query and return results."""
    search = TavilySearchResults()
    results = search.run(query)
    return results

# Outline node: generate research plan
def outline_node(state: BriefState) -> BriefState:
    if state["outline"] is None:
        prompt = (
            f"Generate 4-5 bullet points outlining a research plan for the topic: "
            f"{state['topic']}. Return as a JSON array of strings."
        )
        response = llm.invoke(prompt)
        try:
            outline = json.loads(response.content)
            if isinstance(outline, list):
                state["outline"] = outline
        except Exception:
            state["outline"] = []
    return state

# Research step node: one search per iteration
def research_step_node(state: BriefState) -> BriefState:
    if state["step_index"] < len(state["outline"] or []):
        current_step = state["outline"][state["step_index"]]
        # Perform web search
        search_results = tavily_search(current_step)
        # Summarize results
        summary_prompt = (
            f"Summarize the following search results into 5-8 sentences:\n{search_results}"
        )
        summary = llm.invoke(summary_prompt).content.strip()
        state["notes"].append(summary)
        state["step_index"] += 1
    return state

# Synthesize node: combine notes into final brief
def synthesize_node(state: BriefState) -> BriefState:
    if state["step_index"] >= len(state["outline"] or []):
        notes_text = "\n\n".join(state["notes"])
        synth_prompt = (
            f"Combine the following notes into a concise research brief (about one page). "
            f"Use headings for each point.\n\n{notes_text}"
        )
        brief = llm.invoke(synth_prompt).content.strip()
        state["final_brief"] = brief
    return state

# Build the LangGraph
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_conditional_edges(
    "outline",
    lambda state: "research_step" if state["outline"] is not None else END,
)
graph.add_conditional_edges(
    "research_step",
    lambda state: (
        "research_step"
        if state["step_index"] < len(state["outline"] or [])
        else "synthesize"
    ),
)
graph.add_edge("synthesize", END)

compiled_graph = graph.compile()

# Tool: Generate brief using the graph
@tool
def generate_brief(topic: str) -> str:
    """Generate a research brief for the given topic."""
    initial_state: BriefState = {
        "topic": topic,
        "outline": None,
        "step_index": 0,
        "notes": [],
        "final_brief": None,
    }
    final_state = compiled_graph.invoke(initial_state)
    return final_state["final_brief"] or ""

# Create deepagents agent
agent = create_deep_agent(
    model=llm,
    tools=[generate_brief, tavily_search],
    backend=backend,
    system_prompt="You are a helpful research assistant. Use the provided tools to generate a research brief.",
)

# Demo execution
async def main():
    topic = "Как студенту безопасно подключать MCP к LangChain"
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=f"Generate a brief on: {topic}")]},
        {"configurable": {"thread_id": "session-1"}},
    )
    print(result["messages"][-1].content)

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