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

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

# --- Tavily search tool ---
@tool
def tavily_search(query: str) -> str:
    """Search the web using Tavily and return a short summary of the top result."""
    from tavily import TavilyClient
    client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
    results = client.search(query, max_results=1)
    if not results:
        return "No relevant information found."
    # Return the first result's content (title + snippet)
    first = results[0]
    return f"{first.title}\n{first.snippet}"

# --- State definition for LangGraph (used inside a tool) ---
class BriefState(TypedDict):
    topic: str
    outline: List[str] | None
    step_index: int
    notes: List[str]
    final_brief: str | None

# --- LangGraph graph implementation ---
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

# Helper to format outline as a numbered list
def format_outline(outline: List[str]) -> str:
    return "\n".join(f"{i+1}. {p}" for i, p in enumerate(outline))

# Node: generate outline
async def outline_node(state: BriefState) -> BriefState:
    prompt = f"Create a concise 4–5 point outline for a research brief on the topic: {state['topic']}"
    response = await llm.ainvoke([HumanMessage(content=prompt)])
    outline_text = response.content.strip()
    # Assume the LLM returns a numbered list; split into lines
    outline = [line.strip() for line in outline_text.splitlines() if line.strip()]
    state['outline'] = outline
    state['step_index'] = 0
    state['notes'] = []
    return state

# Node: research one step
async def research_step_node(state: BriefState) -> BriefState:
    current_point = state['outline'][state['step_index']]
    # Use Tavily to get info
    search_query = f"{current_point}"
    search_result = tavily_search(search_query)
    # Summarize the result with LLM
    prompt = f"Summarize the following information in 5–8 sentences, focusing on the key points relevant to the research brief: {search_result}"
    summary = await llm.ainvoke([HumanMessage(content=prompt)])
    state['notes'].append(f"**{current_point}**\n{summary.content.strip()}")
    state['step_index'] += 1
    return state

# Node: synthesize final brief
async def synthesize_node(state: BriefState) -> BriefState:
    notes_text = "\n\n".join(state['notes'])
    prompt = f"Using the following notes, write a cohesive ½–1 page research brief. Include headings for each section.\n\n{notes_text}"
    brief = await llm.ainvoke([HumanMessage(content=prompt)])
    state['final_brief'] = brief.content.strip()
    return state

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

# Define transitions
graph.set_entry_point("outline")
graph.add_conditional_edges(
    "outline",
    lambda _: "research_step",
)

def research_cond(state: BriefState):
    return "research_step" if state['step_index'] < len(state['outline']) else "synthesize"

graph.add_conditional_edges("research_step", research_cond)
graph.add_edge("synthesize", END)

# Compile the graph into a tool
from langgraph.prebuilt import create_react_agent

# The graph will be used as a tool inside deepagents
@tool
async def run_brief(topic: str) -> str:
    """Generate a research brief for the given topic."""
    # Initialize state
    state: BriefState = {
        "topic": topic,
        "outline": None,
        "step_index": 0,
        "notes": [],
        "final_brief": None,
    }
    # Run the graph
    async for partial_state in graph.astream(state):
        pass  # we just wait for completion
    # After completion, return the brief
    return state['final_brief']

# --- DeepAgent setup ---
agent = create_deep_agent(
    model=llm,
    tools=[tavily_search, run_brief],
    backend=backend,
    system_prompt="You are a research assistant that creates concise research briefs. Use the provided tools to gather information and synthesize a brief.",
)

# --- Demo execution ---
async def main():
    default_topic = "Как студенту безопасно подключать MCP к LangChain"
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=f"Create a research brief on: {default_topic}")]},
        {"configurable": {"thread_id": "session-1"}},
    )
    print("\n--- Outline ---")
    # The outline is part of the first tool call; extract it
    for msg in result["messages"]:
        if msg.type == "tool":
            if "outline" in msg.content.lower():
                print(msg.content)
    print("\n--- Final Brief ---")
    print(result["messages"][-1].content)

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