"""
research_pipeline.py -- Complete multi-agent research and writing pipeline.

Install:
    pip install langgraph langchain-anthropic langfuse duckduckgo-search

Run:
    python research_pipeline.py
"""

import operator
from pathlib import Path
from typing import TypedDict, Annotated

from langchain_anthropic import ChatAnthropic
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_core.messages import HumanMessage, BaseMessage
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, START, END
from langgraph.types import interrupt, Command
from langfuse.callback import CallbackHandler

# -- State --

class State(TypedDict):
    topic: str
    research: str
    draft: str
    approved: bool
    saved_to: str
    messages: Annotated[list[BaseMessage], operator.add]


# -- Setup --

model = ChatAnthropic(model="claude-sonnet-4-6")
search = DuckDuckGoSearchRun()
langfuse = CallbackHandler()
OUTPUT_DIR = Path.home() / "ai-workspace" / "published"
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

CALLBACKS = {"callbacks": [langfuse]}


# -- Nodes --

def researcher_node(state: State) -> State:
    print(f"\n[Researcher] Researching: {state['topic']}")
    raw = search.run(state["topic"])
    response = model.invoke([
        HumanMessage(content=
            f"Synthesise these search results about '{state['topic']}' into structured notes:\n\n{raw}\n\n"
            f"Cover: key concepts, recent developments, practical examples, limitations."
        )
    ], config=CALLBACKS)
    return {"research": response.content}


def writer_node(state: State) -> State:
    print("[Writer] Writing draft...")
    response = model.invoke([
        HumanMessage(content=
            f"Write a 500-word blog post about '{state['topic']}' using these notes:\n\n{state['research']}\n\n"
            f"Requirements: specific title, hook opening, 3-4 body paragraphs, actionable conclusion. "
            f"Tone: practical and honest, no hype words."
        )
    ], config=CALLBACKS)
    return {"draft": response.content}


def approval_node(state: State) -> State:
    print("\n[Approval Gate] Pausing for human review...")
    print("-" * 50)
    print(state["draft"][:400], "\n...[truncated]")
    print("-" * 50)

    decision = interrupt({
        "message": "Review the draft above. Approve to publish?",
        "preview": state["draft"][:400],
    })
    return {"approved": decision.get("approved", False)}


def publisher_node(state: State) -> State:
    slug = state["topic"].lower().replace(" ", "-")[:50]
    path = OUTPUT_DIR / f"{slug}.md"
    path.write_text(state["draft"], encoding="utf-8")
    status = "published" if state["approved"] else "rejected"
    print(f"[Publisher] Saved ({status}) -> {path}")
    return {"saved_to": str(path)}


# -- Graph --

def build_pipeline():
    memory = MemorySaver()
    g = StateGraph(State)
    g.add_node("researcher", researcher_node)
    g.add_node("writer", writer_node)
    g.add_node("approval", approval_node)
    g.add_node("publisher", publisher_node)
    g.add_edge(START, "researcher")
    g.add_edge("researcher", "writer")
    g.add_edge("writer", "approval")
    g.add_edge("approval", "publisher")
    g.add_edge("publisher", END)
    return g.compile(checkpointer=memory)


# -- Run --

def run(topic: str):
    app = build_pipeline()
    thread = {"configurable": {"thread_id": f"run-{topic[:20].replace(' ', '-')}"}}
    initial = {"topic": topic, "research": "", "draft": "", "approved": False, "saved_to": "", "messages": []}

    print(f"\n{'='*50}")
    print(f"Topic: {topic}")
    print(f"{'='*50}")

    for _ in app.stream(initial, config=thread):
        pass

    print("\n>>> Human reviewing (auto-approving for demo)...")
    for event in app.stream(Command(resume={"approved": True}), config=thread):
        pass

    print("\n>>> Done. Check Langfuse for full traces.")


if __name__ == "__main__":
    run("How AI agents work and why they matter in 2025")

export ANTHROPIC_API_KEY="your-key"
export LANGFUSE_SECRET_KEY="sk-lf-..."
export LANGFUSE_PUBLIC_KEY="pk-lf-..."

python research_pipeline.py