import os
import json
import asyncio
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# LLM configuration – OpenRouter via langchain_openai
from langchain_openai import ChatOpenAI

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,
)

# Tavily search tool
from langchain_tavily import TavilySearchResults

tavily = TavilySearchResults(tavily_api_key=os.getenv("TAVILY_API_KEY"))

# LangGraph imports
from langgraph.graph import StateGraph, START, END
from typing import TypedDict, Annotated, List
from langgraph.graph.message import add_messages

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

# Outline node – generate 4‑5 bullet points
async def outline_node(state: BriefState) -> BriefState:
    prompt = (
        f"Given the research topic: {state['topic']}\n"
        "Provide a concise outline with 4–5 bullet points. Return the outline as a JSON array of strings."
    )
    response = await llm.ainvoke(prompt)
    try:
        outline = json.loads(response)
        if not isinstance(outline, list):
            raise ValueError
    except Exception:
        # Fallback: split lines
        outline = [line.strip('- • ') for line in response.splitlines() if line.strip()]
    return {
        "topic": state['topic'],
        "outline": outline,
        "step_index": 0,
        "notes": [],
        "final_brief": None,
    }

# Research step node – one web search per step
async def research_step_node(state: BriefState) -> BriefState:
    step = state['outline'][state['step_index']]
    query = f"{state['topic']} {step}"
    search_results = tavily.run(query)
    prompt = (
        f"Using the following search results, write a concise note (5–8 sentences) about the step: {step}.\n"
        f"Search results:\n{search_results}\n"
        "Note:")
    note = await llm.ainvoke(prompt)
    notes = state['notes'] + [note]
    step_index = state['step_index'] + 1
    return {
        "topic": state['topic'],
        "outline": state['outline'],
        "step_index": step_index,
        "notes": notes,
        "final_brief": None,
    }

# Synthesize node – produce final brief
async def synthesize_node(state: BriefState) -> BriefState:
    notes_text = "\n\n".join(state['notes'])
    prompt = (
        f"Based on the following notes, write a cohesive research brief about the topic: {state['topic']}\n"
        "Include headings for each point and keep the brief ½–1 page long.\n"
        f"Notes:\n{notes_text}\n"
        "Final brief:")
    final = await llm.ainvoke(prompt)
    return {
        "topic": state['topic'],
        "outline": state['outline'],
        "step_index": state['step_index'],
        "notes": state['notes'],
        "final_brief": final,
    }

# 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)

# Conditional edge after research_step
def condition(state: BriefState):
    if state['step_index'] < len(state['outline']):
        return "research_step"
    return "synthesize"

graph.add_edge(START, "outline")
graph.add_edge("outline", "research_step")
graph.add_conditional_edges("research_step", condition)
graph.add_edge("synthesize", END)

# Compile the graph into a function
compiled_graph = graph.compile()

# Function to run the whole brief generation
async def run_brief(topic: str) -> str:
    init_state: BriefState = {
        "topic": topic,
        "outline": None,
        "step_index": 0,
        "notes": [],
        "final_brief": None,
    }
    final_state = await compiled_graph.ainvoke(init_state)
    return final_state["final_brief"]

# DeepAgents integration
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from langchain.tools import tool
from langchain_core.messages import HumanMessage

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

@tool
async def generate_brief(topic: str) -> str:
    """Generate a research brief for the given topic."""
    return await run_brief(topic)

agent = create_deep_agent(
    model=llm,
    tools=[generate_brief],
    backend=backend,
    system_prompt="You are a research assistant. Use the provided tools to generate briefs.",
)

async def main():
    default_topic = "Как студенту безопасно подключать MCP к LangChain"
    response = await agent.ainvoke(
        {"messages": [HumanMessage(content=f"Generate brief on '{default_topic}'")],
         "configurable": {"thread_id": "session-1"}},
    )
    # The tool output will be in the last message
    print("\n=== Research Brief ===\n")
    print(response["messages"][-1].content)

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