"""
state.py - Defines the global state schema for the Multi-Agent Blog Generator.
"""
from pydantic import BaseModel


class BlogState(BaseModel):
    # User Inputs
    topic: str = ""
    audience: str = "General Readers"

    # Researcher Agent Outputs & Feedback
    research: str = ""
    research_feedback: str = ""

    # Writer Agent Outputs & Feedback
    draft: str = ""
    draft_feedback: str = ""

    # Editor Agent Final Output
    final_blog: str = ""

    # Metadata & Loop Safeguards
    revision_count: int = 0

"""
agents.py - Prompt templates, LLM initialization, and agent definitions.
"""
import os
from langchain_core.prompts import ChatPromptTemplate
from langchain_groq import ChatGroq


def get_llm(model_name: str = "llama-3.3-70b-versatile", temperature: float = 0.5):
    """Initializes and returns the ChatGroq model instance."""
    api_key = os.getenv("GROQ_API_KEY")
    return ChatGroq(
        model=model_name,
        temperature=temperature,
        groq_api_key=api_key
    )


# =====================================================================
# 1. RESEARCHER AGENT
# =====================================================================
RESEARCHER_PROMPT = ChatPromptTemplate.from_messages([
    (
        "system",
        "You are an expert Researcher Agent. Given a blog topic and target audience, "
        "produce a clear, structured research outline.\n\n"
        "Include:\n"
        "- 5 to 7 key points the blog should cover\n"
        "- Important facts, stats, and real-world examples for each point\n"
        "- Suggested angle or narrative hook\n"
        "- Be concise, use bullet points, and DO NOT write the full blog post."
    ),
    (
        "user",
        "Topic: {topic}\n"
        "Target Audience: {audience}\n"
        "Revision Context: {revision_hints}\n\n"
        "Produce the research outline now."
    )
])


def researcher_agent(llm: ChatGroq, topic: str, audience: str, feedback: str = "") -> str:
    """Executes the Researcher Agent to generate or revise an outline."""
    if feedback:
        revision_hints = f"The human reviewer requested these changes on your previous research:\n'{feedback}'\nPlease address all feedback."
    else:
        revision_hints = "This is your initial research attempt. Provide a comprehensive outline."

    chain = RESEARCHER_PROMPT | llm
    result = chain.invoke({
        "topic": topic,
        "audience": audience,
        "revision_hints": revision_hints
    })
    return str(result.content)


# =====================================================================
# 2. WRITER AGENT
# =====================================================================
WRITER_PROMPT = ChatPromptTemplate.from_messages([
    (
        "system",
        "You are a professional Blog Writer Agent. Using the approved research notes provided, "
        "write a complete, highly engaging blog post.\n\n"
        "Guidelines:\n"
        "- Length: 500 to 800 words\n"
        "- Structure: Catchy title, strong introductory hook, 3 to 5 H2 sections, and a conclusion\n"
        "- Tone: Clear, friendly, and tailored specifically to the target audience\n"
        "- Formatting: Clean GitHub Flavored Markdown\n"
        "- Do not append word counts or meta commentary."
    ),
    (
        "user",
        "Topic: {topic}\n"
        "Target Audience: {audience}\n"
        "Approved Research Notes:\n{research}\n\n"
        "Revision Context: {revision_hints}\n\n"
        "Write the full blog post now."
    )
])


def writer_agent(llm: ChatGroq, topic: str, audience: str, research: str, feedback: str = "") -> str:
    """Executes the Writer Agent to generate or revise a full markdown draft."""
    if feedback:
        revision_hints = f"The human reviewer requested these changes on your previous draft:\n'{feedback}'\nApply these revisions strictly."
    else:
        revision_hints = "This is your first drafting attempt. Transform the research into a compelling post."

    chain = WRITER_PROMPT | llm
    result = chain.invoke({
        "topic": topic,
        "audience": audience,
        "research": research,
        "revision_hints": revision_hints
    })
    return str(result.content)


# =====================================================================
# 3. EDITOR AGENT
# =====================================================================
EDITOR_PROMPT = ChatPromptTemplate.from_messages([
    (
        "system",
        "You are an Editor Agent representing the final quality checkpoint before publication.\n\n"
        "Your task:\n"
        "- Fix grammar, spelling, and awkward phrasing\n"
        "- Improve narrative flow and transitions between sections\n"
        "- Strip out generic AI clichés and overly fancy words (e.g., 'delve', 'tapestry', 'testament')\n"
        "- Ensure the writing sounds human, authoritative, and natural\n"
        "- Preserve the original markdown formatting\n"
        "- Output ONLY the final published blog post with ZERO introductory or concluding remarks."
    ),
    (
        "user",
        "Topic: {topic}\n"
        "Raw Draft:\n{draft}\n\n"
        "Produce the final edited blog post now."
    )
])


def editor_agent(llm: ChatGroq, topic: str, draft: str) -> str:
    """Executes the Editor Agent for final editorial cleanup."""
    chain = EDITOR_PROMPT | llm
    result = chain.invoke({
        "topic": topic,
        "draft": draft
    })
    return str(result.content)

"""
graph.py - Orchestrates nodes, interrupt checkpoints, conditional routes, and graph compilation.
"""
from typing import Literal
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.types import interrupt

from state import BlogState
from agents import get_llm, researcher_agent, writer_agent, editor_agent

MAX_REVISIONS = 3
llm = get_llm()


# ---------------------------------------------------------------------
# Graph Execution Nodes
# ---------------------------------------------------------------------
def researcher_node(state: BlogState) -> dict:
    """Executes the researcher agent and stores outline in state."""
    research_data = researcher_agent(
        llm=llm,
        topic=state.topic,
        audience=state.audience,
        feedback=state.research_feedback
    )
    return {
        "research": research_data,
        "research_feedback": ""  # Reset feedback once processed
    }


def human_review_research_node(state: BlogState) -> dict:
    """Interrupts execution to obtain human review for the research outline."""
    decision = interrupt({
        "stage": "research_review",
        "research": state.research,
        "instructions": [
            "Reply with 'approve' to proceed to writing.",
            "Provide textual feedback to request research revisions."
        ]
    })

    # Parse dictionary or string inputs from Command(resume=...)
    if isinstance(decision, dict):
        action = decision.get("action", "approve").lower()
        feedback = decision.get("feedback", "")
    else:
        text = str(decision).strip()
        if text.lower() in ["approve", "approved", "ok", "yes", "proceed", ""]:
            action = "approve"
            feedback = ""
        else:
            action = "revise"
            feedback = text

    if action == "revise" and feedback:
        return {
            "research_feedback": feedback,
            "revision_count": state.revision_count + 1
        }
    return {"research_feedback": ""}


def writer_node(state: BlogState) -> dict:
    """Executes the writer agent and stores draft in state."""
    draft_data = writer_agent(
        llm=llm,
        topic=state.topic,
        audience=state.audience,
        research=state.research,
        feedback=state.draft_feedback
    )
    return {
        "draft": draft_data,
        "draft_feedback": ""  # Reset feedback once processed
    }


def human_review_draft_node(state: BlogState) -> dict:
    """Interrupts execution to obtain human review for the full draft."""
    decision = interrupt({
        "stage": "draft_review",
        "draft": state.draft,
        "instructions": [
            "Reply with 'approve' to proceed to final editing.",
            "Provide textual feedback to request drafting revisions."
        ]
    })

    if isinstance(decision, dict):
        action = decision.get("action", "approve").lower()
        feedback = decision.get("feedback", "")
    else:
        text = str(decision).strip()
        if text.lower() in ["approve", "approved", "ok", "yes", "proceed", ""]:
            action = "approve"
            feedback = ""
        else:
            action = "revise"
            feedback = text

    if action == "revise" and feedback:
        return {
            "draft_feedback": feedback,
            "revision_count": state.revision_count + 1
        }
    return {"draft_feedback": ""}


def editor_node(state: BlogState) -> dict:
    """Executes the editor agent for final publication polish."""
    final_output = editor_agent(
        llm=llm,
        topic=state.topic,
        draft=state.draft
    )
    return {"final_blog": final_output}


# ---------------------------------------------------------------------
# Conditional Routing Logic
# ---------------------------------------------------------------------
def route_after_research_review(state: BlogState) -> Literal["researcher", "writer"]:
    """Routes back to researcher if feedback exists and budget permits; else to writer."""
    if state.research_feedback and state.revision_count <= MAX_REVISIONS:
        return "researcher"
    return "writer"


def route_after_draft_review(state: BlogState) -> Literal["writer", "editor"]:
    """Routes back to writer if feedback exists and budget permits; else to editor."""
    if state.draft_feedback and state.revision_count <= MAX_REVISIONS:
        return "writer"
    return "editor"


# ---------------------------------------------------------------------
# Graph Assembly
# ---------------------------------------------------------------------
def build_blog_graph():
    """Assembles and compiles the LangGraph Multi-Agent Workflow."""
    builder = StateGraph(BlogState)

    # Register Nodes
    builder.add_node("researcher", researcher_node)
    builder.add_node("review_research", human_review_research_node)
    builder.add_node("writer", writer_node)
    builder.add_node("review_draft", human_review_draft_node)
    builder.add_node("editor", editor_node)

    # Define Linear & Conditional Transitions
    builder.add_edge(START, "researcher")
    builder.add_edge("researcher", "review_research")
    
    builder.add_conditional_edges(
        "review_research",
        route_after_research_review,
        {"researcher": "researcher", "writer": "writer"}
    )

    builder.add_edge("writer", "review_draft")
    
    builder.add_conditional_edges(
        "review_draft",
        route_after_draft_review,
        {"writer": "writer", "editor": "editor"}
    )

    builder.add_edge("editor", END)

    # Compile with In-Memory Checkpointing
    memory = MemorySaver()
    return builder.compile(checkpointer=memory)

from langgraph.types import Command
from graph import build_blog_graph

graph = build_blog_graph()
config = {"configurable": {"thread_id": "thread-session-001"}}

# 1. Start execution with user inputs
initial_input = {
    "topic": "What is Model Context Protocol (MCP)? A 2026 Developer Guide",
    "audience": "Agentic AI Engineers & Python Developers"
}

graph.invoke(initial_input, config=config)

# 2. Inspect state and extract the interrupt payload
state_snapshot = graph.get_state(config)
interrupt_payload = state_snapshot.tasks[0].interrupts[0].value

print(f"Current Stage: {interrupt_payload['stage']}")
print("\n--- Proposed Research Outline ---\n")
print(interrupt_payload["research"])

# Resume execution with an 'approve' command
graph.invoke(
    Command(resume={"action": "approve", "feedback": ""}),
    config=config
)

# Execution proceeds through writer_node and halts at review_draft
state_snapshot = graph.get_state(config)
draft_payload = state_snapshot.tasks[0].interrupts[0].value

print(f"Current Stage: {draft_payload['stage']}")
print("\n--- Generated Draft Blog ---\n")
print(draft_payload["draft"][:600] + "\n... [Draft Truncated]")

# Resume execution to trigger editor_node and finish at END
graph.invoke(
    Command(resume={"action": "approve", "feedback": ""}),
    config=config
)

final_state = graph.get_state(config).values
print("\n================ FINAL PUBLISHED BLOG POST ================\n")
print(final_state["final_blog"])

config_session_2 = {"configurable": {"thread_id": "thread-session-002"}}

# 1. Start initial run
graph.invoke({
    "topic": "Building a Q&A Bot with LangChain and OpenAI",
    "audience": "Beginner Developers"
}, config=config_session_2)

# 2. Approve research outline directly
graph.invoke(Command(resume={"action": "approve"}), config=config_session_2)

# 3. Draft arrives -> Human requests specific revisions
critique = "The draft is too long. Condense to under 400 words and remove all raw code snippets."
graph.invoke(
    Command(resume={"action": "revise", "feedback": critique}),
    config=config_session_2
)

# 4. Check the revised draft at the draft review stage
revised_snapshot = graph.get_state(config_session_2)
revised_payload = revised_snapshot.tasks[0].interrupts[0].value

print("--- Revised Draft (Compliant with Feedback) ---")
print(revised_payload["draft"])