"""
02_sequential_chain.py — Classic linear pipeline: A → B → C → D

Each agent declares its upstream with llm_input=["tag"] — that is the ONLY
way edges are drawn. Agents without a list llm_input appear as isolated nodes.

Pipeline: researcher → writer → editor → publisher

WHAT THE DASHBOARD SHOULD SHOW:
  Live Graph:
    • 4 nodes left-to-right: researcher → writer → editor → publisher
    • Edges numbered: [1]→[2]→[3]→[4]
    • All green on success

  State View:
    • Step 1 — researcher: INPUT = {topic}, OUTPUT = research notes
    • Step 2 — writer:     INPUT = {notes}, OUTPUT = draft blog post
    • Step 3 — editor:     INPUT = {draft}, OUTPUT = edited post
    • Step 4 — publisher:  INPUT = {edited}, OUTPUT = published article

  Time Travel:
    • 4-step scrubber — scrub back to see research notes before draft existed

  Token Dashboard:
    • 4 bars, one per agent

  Run Comparison:
    • 4 rows: each agent with duration and token count
"""
import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from agentradar import radar, reset_run

load_dotenv()

llm = ChatOpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url=os.getenv("DEEPSEEK_BASE_URL"),
    model=os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
    temperature=0.7,
)


@radar(tag="researcher", llm_input="Research the topic and return 5 key bullet-point facts")
def researcher(topic: str) -> str:
    """Step 1: Research the topic. No upstream — floats as root node."""
    message = llm.invoke(
        f"Research this topic and return exactly 5 key bullet-point facts:\n{topic}"
    )
    return message.content


@radar(tag="writer", llm_input=["researcher"])
def writer(research_notes: str) -> str:
    """Step 2: Write a draft. Declares researcher as upstream → edge drawn."""
    message = llm.invoke(
        f"Write a short 3-paragraph blog post draft based on these research notes:\n{research_notes}"
    )
    return message.content


@radar(tag="editor", llm_input=["writer"])
def editor(draft: str) -> str:
    """Step 3: Edit the draft. Declares writer as upstream → edge drawn."""
    message = llm.invoke(
        f"Edit this blog post draft for clarity, grammar, and flow. "
        f"Return the improved version:\n{draft}"
    )
    return message.content


@radar(tag="publisher", llm_input=["editor"])
def publisher(edited_post: str) -> dict:
    """Step 4: Publish. Declares editor as upstream → edge drawn."""
    message = llm.invoke(
        f"Format this blog post with a catchy title, subheadings, and a conclusion paragraph:\n{edited_post}"
    )
    return {
        "title": "AI-Generated Article",
        "content": message.content,
        "word_count": len(message.content.split()),
        "status": "published",
    }


def main():
    reset_run()
    topic = "How neural networks learn: backpropagation explained simply"

    print("=== Sequential Chain Pipeline ===\n")

    print("[1/4] Researching...")
    notes = researcher(topic)
    print(f"  → Got {len(notes.split())} words of research notes\n")

    print("[2/4] Writing draft...")
    draft = writer(notes)
    print(f"  → Draft: {len(draft.split())} words\n")

    print("[3/4] Editing...")
    edited = editor(draft)
    print(f"  → Edited: {len(edited.split())} words\n")

    print("[4/4] Publishing...")
    article = publisher(edited)
    print(f"  → Published! Word count: {article['word_count']}")
    print(f"\nPreview: {article['content'][:80]}...")

    print("\n✓ Dashboard: 4-node chain via llm_input list, no auto-detection")

    print("=== Sequential Chain Pipeline ===\n")

    print("[1/4] Researching...")
    notes = researcher(topic)
    print(f"  → Got {len(notes.split())} words of research notes\n")

    print("[2/4] Writing draft...")
    draft = writer(notes)
    print(f"  → Draft: {len(draft.split())} words\n")

    print("[3/4] Editing...")
    edited = editor(draft)
    print(f"  → Edited: {len(edited.split())} words\n")

    print("[4/4] Publishing...")
    article = publisher(edited)
    print(f"  → Published! Word count: {article['word_count']}")
    print(f"\nPreview: {article['content'][:80]}...")

    print("\n✓ Dashboard: 4-node chain via llm_input list, no auto-detection")



if __name__ == "__main__":
    main()
    input("\nPress Enter to view dashboard (http://localhost:7111) or Ctrl+C to exit...")
