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

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
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

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

# ----------------------------------------------------------------------
# Tools
# ----------------------------------------------------------------------
search_tool = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY"))

@tool
def web_search(query: str) -> str:
    """
    Perform a web search using Tavily and return a concise summary of the top results.
    The LLM will ask for a short note based on this summary.
    """
    results = search_tool.run(query)
    # results is a list of dicts with 'url' and 'content'
    summaries = [r.get("content", "") for r in results[:3]]
    return "\n".join(summaries) if summaries else "No relevant results found."

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

# ----------------------------------------------------------------------
# Nodes
# ----------------------------------------------------------------------
def outline_node(state: BriefState) -> BriefState:
    """Generate a 4-5 item outline for the given topic."""
    prompt = (
        f"Create a concise outline (4-5 bullet points) for a short research brief on the topic:\n"
        f"\"{state['topic']}\"\n"
        "Each bullet should be a short phrase suitable as a section heading."
    )
    response = llm.invoke([HumanMessage(content=prompt)])
    outline_text = response.content.strip()
    # Split on newlines and strip bullet characters
    items = [line.lstrip("-• ").strip() for line in outline_text.splitlines() if line.strip()]
    return {
        **state,
        "outline": items,
        "step_index": 0,
        "notes": [],
    }

def research_step_node(state: BriefState) -> BriefState:
    """Research one outline item, produce a short note, and store it."""
    outline = state["outline"]
    idx = state["step_index"]
    if outline is None or idx >= len(outline):
        return state  # safety

    current_topic = outline[idx]
    # Use web search tool
    search_result = web_search(current_topic)

    # Prompt LLM to write a 5-8 sentence note based on search result
    note_prompt = (
        f"Based on the following web search summary, write a short note (5-8 sentences) "
        f"that could serve as a paragraph for a research brief about \"{state['topic']}\". "
        f"Focus on the aspect: \"{current_topic}\".\n\n"
        f"Search summary:\n{search_result}"
    )
    note_response = llm.invoke([HumanMessage(content=note_prompt)])
    note = note_response.content.strip()

    new_notes = state["notes"] + [f"## {current_topic}\n{note}"]
    return {
        **state,
        "notes": new_notes,
        "step_index": idx + 1,
    }

def synthesize_node(state: BriefState) -> BriefState:
    """Combine all notes into a coherent brief."""
    if not state["notes"]:
        final = "No notes were collected."
    else:
        combined = "\n\n".join(state["notes"])
        synthesis_prompt = (
            f"Combine the following sections into a single cohesive research brief (about half to one page). "
            f"Keep the headings, ensure logical flow, and add a brief introduction and conclusion.\n\n"
            f"{combined}"
        )
        synthesis_response = llm.invoke([HumanMessage(content=synthesis_prompt)])
        final = synthesis_response.content.strip()
    return {
        **state,
        "final_brief": final,
    }

# ----------------------------------------------------------------------
# Graph construction
# ----------------------------------------------------------------------
workflow = StateGraph(BriefState)

workflow.add_node("outline", outline_node)
workflow.add_node("research_step", research_step_node)
workflow.add_node("synthesize", synthesize_node)

workflow.add_edge(START, "outline")
workflow.add_conditional_edges(
    "outline",
    lambda state: "research_step" if state["outline"] else END,
)

def continue_condition(state: BriefState) -> str:
    if state["outline"] is None:
        return END
    if state["step_index"] < len(state["outline"]):
        return "research_step"
    return "synthesize"

workflow.add_edge("research_step", "research_step")
workflow.add_conditional_edges("research_step", continue_condition)
workflow.add_edge("synthesize", END)

graph = workflow.compile()

# ----------------------------------------------------------------------
# DeepAgent wrapper (required by the course)
# ----------------------------------------------------------------------
agent = create_deep_agent(
    model=llm,
    tools=[web_search],
    backend=backend,
    system_prompt="You are an AI research assistant that helps build short research briefs.",
)

# ----------------------------------------------------------------------
# Demo execution
# ----------------------------------------------------------------------
async def main():
    topic = "Как студенту безопасно подключать MCP к LangChain"
    # Initialize state
    init_state: BriefState = {
        "topic": topic,
        "outline": None,
        "step_index": 0,
        "notes": [],
        "final_brief": None,
    }

    # Run the graph
    result = await graph.ainvoke(
        init_state,
        config={"configurable": {"thread_id": "demo-1"}},
    )

    # Print results
    print("\n--- Outline ---")
    if result["outline"]:
        for i, item in enumerate(result["outline"], 1):
            print(f"{i}. {item}")

    print("\n--- Research Steps ---")
    for i, note in enumerate(result["notes"], 1):
        print(f"[Step {i}]")
        print(note)
        print()

    print("\n--- Final Brief ---")
    print(result["final_brief"] or "No brief generated.")

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