"""LangGraph interrupt/resume with automatic Beacon trace continuation.

Demonstrates how BeaconLangGraphConfig automatically detects interrupts and
continues traces across interrupt/resume cycles using LangGraph checkpointers.

This example shows:
- Trace continuity: Same trace_id maintained across interrupt → resume
- Session continuity: Same session_id across multiple conversations in a thread
- Automatic interrupt detection from checkpoint state (state.next, state.interrupts)
- Human-in-the-loop workflow pattern with approval gates

The example uses a simple research assistant workflow:
1. Research - Gather information (runs to completion)
2. Plan - Create action plan (INTERRUPTED - requires human approval)
3. Execute - Execute the plan (runs after resume)

Requirements:
    pip install lumenova-beacon[langchain]
    pip install langchain-openai langgraph
"""

import uuid
from typing import Annotated, Sequence, TypedDict

import dotenv
from langchain_core.messages import BaseMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, END, add_messages

from lumenova_beacon.tracing.integrations.langchain import BeaconLangGraphConfig

dotenv.load_dotenv()


# === State Definition ===

class ResearchState(TypedDict):
    """State for the research assistant workflow."""
    messages: Annotated[Sequence[BaseMessage], add_messages]
    research_findings: str
    action_plan: str
    execution_result: str


# === Node Implementations ===

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)


def research_node(state: ResearchState):
    """Gather information about the topic."""
    query = state["messages"][-1].content

    # Call LLM to research the topic
    research_prompt = f"Provide 3 key facts about: {query}"
    findings = llm.invoke([HumanMessage(content=research_prompt)]).content

    print(f"\n[Research Node] Findings:\n{findings}")

    return {
        "research_findings": findings,
        "messages": [HumanMessage(content=f"Research completed: {findings[:100]}...")]
    }


def plan_node(state: ResearchState):
    """Create an action plan based on research findings."""
    findings = state.get("research_findings", "")

    # Call LLM to create action plan
    plan_prompt = f"Based on these findings, create a 3-step action plan:\n{findings}"
    plan = llm.invoke([HumanMessage(content=plan_prompt)]).content

    print(f"\n[Plan Node] Action Plan:\n{plan}")

    return {
        "action_plan": plan,
        "messages": [HumanMessage(content=f"Plan created: {plan[:100]}...")]
    }


def execute_node(state: ResearchState):
    """Execute the action plan."""
    plan = state.get("action_plan", "")

    # Call LLM to simulate execution
    execute_prompt = f"Simulate executing this plan and provide a brief summary:\n{plan}"
    result = llm.invoke([HumanMessage(content=execute_prompt)]).content

    print(f"\n[Execute Node] Execution Result:\n{result}")

    return {
        "execution_result": result,
        "messages": [HumanMessage(content=f"Execution complete: {result[:100]}...")]
    }


# === Graph Construction ===

def build_graph():
    """Build the research assistant graph with interrupt before planning."""
    workflow = StateGraph(ResearchState)

    # Add nodes
    workflow.add_node("research", research_node)
    workflow.add_node("plan", plan_node)
    workflow.add_node("execute", execute_node)

    # Add edges: research → plan → execute → END
    workflow.set_entry_point("research")
    workflow.add_edge("research", "plan")
    workflow.add_edge("plan", "execute")
    workflow.add_edge("execute", END)

    # Compile with checkpointer and interrupt before planning
    checkpointer = MemorySaver()
    return workflow.compile(
        checkpointer=checkpointer,
        interrupt_before=["plan"]  # Pause for human approval before planning
    )


# === Main Example ===

if __name__ == "__main__":
    # Build the graph once
    app = build_graph()

    # Generate a unique thread ID for this conversation
    thread_id = f"research-{uuid.uuid4()}"
    print(f"Thread ID: {thread_id}\n")

    # =========================================================================
    # PHASE 1: Initial Run (will interrupt before planning)
    # =========================================================================
    print("=" * 70)
    print("PHASE 1: Initial Research (will interrupt before planning)")
    print("=" * 70)

    # Create BeaconLangGraphConfig - this generates fresh IDs
    beacon = BeaconLangGraphConfig(
        graph=app,
        thread_id=thread_id,
        agent_name="research-assistant",
    )

    # Get config (triggers checkpoint resolution - will find no checkpoint)
    config = beacon.get_config()

    print("\nBeacon State (Initial):")
    print(f"  Trace ID:      {beacon.trace_id}")
    print(f"  Root Span ID:  {beacon.root_span_id}")
    print(f"  Session ID:    {beacon.handler._session_id}")
    print(f"  Is Resume:     {beacon.is_resume}")

    # Run the graph with a research query
    print("\nExecuting graph...")
    result = app.invoke(
        {"messages": [HumanMessage(content="AI safety best practices")]},
        config=config,
    )

    # Check checkpoint state
    state = app.get_state(config)
    print(f"\n✓ Graph interrupted: {bool(state.next)}")
    print(f"  Next nodes: {state.next}")
    print(f"  Checkpoint metadata keys: {list(state.metadata.keys())}")

    # Save trace ID for comparison
    phase1_trace_id = beacon.trace_id
    phase1_root_span_id = beacon.root_span_id
    phase1_session_id = beacon.handler._session_id

    # =========================================================================
    # PHASE 2: Resume (after human approval)
    # =========================================================================
    print("\n" + "=" * 70)
    print("PHASE 2: Resume After Human Approval")
    print("=" * 70)

    # Create NEW BeaconLangGraphConfig with SAME thread_id
    beacon_resume = BeaconLangGraphConfig(
        graph=app,
        thread_id=thread_id,  # Same thread!
        agent_name="research-assistant",
    )

    # Get config (triggers checkpoint resolution - will find interrupt)
    config_resume = beacon_resume.get_config()

    print("\nBeacon State (Resume):")
    print(f"  Trace ID:      {beacon_resume.trace_id}")
    print(f"  Root Span ID:  {beacon_resume.root_span_id}")
    print(f"  Session ID:    {beacon_resume.handler._session_id}")
    print(f"  Is Resume:     {beacon_resume.is_resume}")

    # Verify trace continuity
    print("\nTrace Continuity Verification:")
    print(f"  ✓ Same trace ID:      {beacon_resume.trace_id == phase1_trace_id}")
    print(f"  ✓ Same root span ID:  {beacon_resume.root_span_id == phase1_root_span_id}")
    print(f"  ✓ Same session ID:    {beacon_resume.handler._session_id == phase1_session_id}")

    # Resume execution (pass None to continue from checkpoint)
    print("\nResuming graph execution...")
    result = app.invoke(None, config=config_resume)

    print("\n✓ Graph completed")

    # =========================================================================
    # PHASE 3: New Query (demonstrates session continuity)
    # =========================================================================
    print("\n" + "=" * 70)
    print("PHASE 3: New Query in Same Thread")
    print("=" * 70)

    # Create NEW BeaconLangGraphConfig with SAME thread_id
    beacon_new = BeaconLangGraphConfig(
        graph=app,
        thread_id=thread_id,  # Same thread, but previous execution completed
        agent_name="research-assistant",
    )

    # Get config (triggers checkpoint resolution - will find completed state)
    config_new = beacon_new.get_config()

    print("\nBeacon State (New Query):")
    print(f"  Trace ID:      {beacon_new.trace_id}")
    print(f"  Root Span ID:  {beacon_new.root_span_id}")
    print(f"  Session ID:    {beacon_new.handler._session_id}")
    print(f"  Is Resume:     {beacon_new.is_resume}")

    # Verify session continuity but new trace
    print("\nSession Continuity Verification:")
    print(f"  ✓ Different trace ID:  {beacon_new.trace_id != phase1_trace_id}")
    print(f"  ✓ SAME session ID:     {beacon_new.handler._session_id == phase1_session_id}")

    # Run a new query (this will execute to completion without interrupt)
    print("\nExecuting new query...")
    result = app.invoke(
        {"messages": [HumanMessage(content="Machine learning trends")]},
        config=config_new,
    )

    print("\n✓ Graph completed without interruption\n")


# Check your Beacon dashboard to see the complete trace flow:
#
# PHASE 1 (Initial run until interrupt):
# ========================================
# Trace ID: {phase1_trace_id}
# - Root span (type=agent, name="research-assistant")
# - Child span: "research" node (type=chain)
#   - LLM generation span with token counts
# - Root span attributes:
#   - langgraph.interrupt = true
#   - langgraph.interrupt.id = interrupt ID
#   - gen_ai.agent.name = "research-assistant"
#
# PHASE 2 (Resume - SAME trace continues):
# ==========================================
# Trace ID: {phase1_trace_id} (SAME!)
# - Root span (continued from Phase 1)
# - Child span: "plan" node (type=chain)
#   - LLM generation span
# - Child span: "execute" node (type=chain)
#   - LLM generation span
# - Root span attributes:
#   - langgraph.resumed = true
#   - Complete workflow: research → plan → execute
#
# PHASE 3 (New conversation):
# ============================
# Trace ID: {new_trace_id} (DIFFERENT)
# Session ID: {session_id} (SAME as Phase 1 & 2)
# - Root span (type=agent, fresh trace)
# - Child spans: research → plan → execute
# - No interruption (runs to completion)
#
# Key observations:
# - Phase 1 & 2 share the same trace_id (trace continuity across interrupt)
# - All 3 phases share the same session_id (conversation thread tracking)
# - Phase 3 gets a new trace_id (fresh conversation after previous completed)
# - Checkpoint metadata automatically persists beacon_trace_id, beacon_root_span_id, beacon_session_id
