"""
Self-Correcting Agent implementation using LangGraph.

This module defines a simple LangGraph that:
1. Generates an answer to a user question.
2. Checks the quality of the answer.
3. Corrects the answer if needed.
4. Returns the final answer.

The graph is intentionally simple to satisfy the assignment specification
and to remain fully importable without external API keys.
"""

from dataclasses import dataclass, field
from typing import Any, Dict

# Import LangGraph components
try:
    from langgraph.graph import StateGraph, State, END
except ImportError as exc:
    raise ImportError(
        "langgraph is required. Install it via 'pip install langgraph==0.0.1'"
    ) from exc

# --------------------------------------------------------------------------- #
# State definition
# --------------------------------------------------------------------------- #
@dataclass
class AgentState(State):
    """
    Holds the state of the agent during execution.
    """
    question: str = ""
    answer: str = ""
    feedback: str = ""
    final_answer: str = ""

# --------------------------------------------------------------------------- #
# Node implementations
# --------------------------------------------------------------------------- #
def ask(state: AgentState) -> AgentState:
    """
    Generates an answer to the provided question.
    """
    # In a real implementation, this would call an LLM.
    # Here we use a deterministic placeholder.
    state.answer = f"Answer to: {state.question}"
    return state

def check(state: AgentState) -> AgentState:
    """
    Checks the quality of the generated answer.
    """
    # Simple heuristic: if the answer contains the word 'bad', flag it.
    if "bad" in state.answer.lower():
        state.feedback = "Needs correction"
    else:
        state.feedback = "Good"
    return state

def correct(state: AgentState) -> AgentState:
    """
    Corrects the answer if the feedback indicates a problem.
    """
    if state.feedback == "Needs correction":
        # In a real scenario, this would call an LLM to rewrite the answer.
        state.final_answer = f"Corrected: {state.answer}"
    else:
        state.final_answer = state.answer
    return state

def final(state: AgentState) -> str:
    """
    Returns the final answer to the user.
    """
    return state.final_answer

# --------------------------------------------------------------------------- #
# Graph construction
# --------------------------------------------------------------------------- #
def build_agent_graph() -> StateGraph:
    """
    Builds and returns the LangGraph for the self-correcting agent.
    """
    graph = StateGraph(AgentState)

    # Add nodes
    graph.add_node("ask", ask)
    graph.add_node("check", check)
    graph.add_node("correct", correct)
    graph.add_node("final", final)

    # Define edges
    graph.set_entry_point("ask")
    graph.add_edge("ask", "check")

    # Conditional transition from check to either correct or final
    def check_transition(state: AgentState) -> str:
        return "correct" if state.feedback != "Good" else "final"

    graph.add_conditional_edges("check", check_transition)

    graph.add_edge("correct", "final")
    graph.add_edge("final", END)

    return graph

# --------------------------------------------------------------------------- #
# Public API
# --------------------------------------------------------------------------- #
def run_agent(question: str) -> str:
    """
    Runs the self-correcting agent on the given question.

    Parameters
    ----------
    question : str
        The user question to answer.

    Returns
    -------
    str
        The final answer produced by the agent.
    """
    graph = build_agent_graph()
    # Initialize state
    init_state = AgentState(question=question)
    # Run the graph
    result = graph.invoke(init_state)
    # The result is the final answer string
    return result

__all__ = [
    "AgentState",
    "ask",
    "check",
    "correct",
    "final",
    "build_agent_graph",
    "run_agent",
]