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 deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_tavily import TavilySearchResults

# Load environment variables
from dotenv import load_dotenv
load_dotenv()

# LLM configuration – OpenRouter
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 – local shell + filesystem
backend = CompositeBackend([
    LocalShellBackend(workspace_dir="./workspace"),
    FilesystemBackend(),
])

# Tavily search tool – real web search
@tool
def tavily_search(query: str) -> str:
    """Search the web using Tavily and return a short summary."""
    tavily = TavilySearchResults(max_results=3)
    results = tavily.run(query)
    # Concatenate titles and snippets
    return "\n".join(f"{r['title']}: {r['content']}" for r in results)

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

# Node: generate outline
async def outline_node(state: BriefState) -> BriefState:
    prompt = (
        f"Generate a concise outline of 4–5 research steps for the topic: {state['topic']}\n"
        "Return a JSON array of strings, each a single step."
    )
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    # Extract JSON array
    import json, re
    try:
        array_text = re.search(r"\[.*\]", response.content, re.S).group(0)
        outline = json.loads(array_text)
    except Exception:
        outline = ["Step 1: ...", "Step 2: ...", "Step 3: ...", "Step 4: ..."]
    state["outline"] = outline
    state["step_index"] = 0
    state["notes"] = []
    return state

# Node: research one step
async def research_step_node(state: BriefState) -> BriefState:
    step = state["outline"][state["step_index"]]
    # Use Tavily to gather info
    search_query = f"{state['topic']} {step}"
    search_result = tavily_search(search_query)
    # Summarize with LLM
    prompt = (
        f"Using the following search results, write a concise note (5–8 sentences) for the step: {step}\n"
        f"Search results:\n{search_result}\n"
        "Note: keep it factual and cite sources if possible."
    )
    note = await llm.ainvoke([HumanMessage(content=prompt)])
    state["notes"].append(note.content.strip())
    state["step_index"] += 1
    return state

# Node: synthesize final brief
async def synthesize_node(state: BriefState) -> BriefState:
    notes = state["notes"]
    outline = state["outline"]
    prompt = (
        "You are an academic writer. Using the following outline and notes, produce a cohesive research brief of ½–1 page.\n"
        f"Outline: {outline}\n"
        f"Notes: {notes}\n"
        "Structure the brief with headings matching the outline steps."
    )
    brief = await llm.ainvoke([HumanMessage(content=prompt)])
    state["final_brief"] = brief.content.strip()
    return state

# Build LangGraph
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")
# After outline, loop research_step until all steps processed
graph.add_conditional_edges(
    "outline",
    lambda _: "research_step",
)
# After each research_step, decide whether to continue or synthesize
graph.add_conditional_edges(
    "research_step",
    lambda state: "synthesize" if state["step_index"] >= len(state["outline"]) else "research_step",
)
# Final node
graph.add_edge("synthesize", END)

app = graph.compile()

# DeepAgent wrapper
agent = create_deep_agent(
    model=llm,
    tools=[tavily_search],
    backend=backend,
    system_prompt="You are a research assistant that builds a brief.",
)

async def main():
    # Default topic
    topic = os.getenv("DEFAULT_TOPIC", "Как студенту безопасно подключать MCP к LangChain")
    # Run LangGraph 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 notes per step
    print("\n=== Notes ===")
    for i, note in enumerate(state["notes"], 1):
        print(f"[Step {i}] {note}\n")
    # Print final brief
    print("\n=== Final Brief ===")
    print(state["final_brief"])

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