"""
Main entry point for the self‑correcting LangGraph agent.

The program demonstrates a simple task – evaluating an arithmetic expression –
and shows how the agent retries until the LLM judge confirms success or the
maximum number of attempts is reached.

Usage:
    python main.py "2+2"

Three examples are printed in the README and can be run directly from this file.
"""

import os
import random
from typing import TypedDict, Dict, Any

from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage
from rich.console import Console
from rich.table import Table

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
MAX_ATTEMPTS = 5
LLM_MODEL = "openai/gpt-oss-20b:free"
BASE_URL = "https://platform.brojs.ru/jrnl-bh/api/inference/v1"
API_KEY_ENV = "JOURNAL_MCP_PAT"

console = Console()

# ---------------------------------------------------------------------------
# LLM instance – BroJS provider
# ---------------------------------------------------------------------------
llm = ChatOpenAI(
    model=LLM_MODEL,
    base_url=BASE_URL,
    api_key=os.getenv(API_KEY_ENV),
    temperature=0.0,
)

# ---------------------------------------------------------------------------
# State definition
# ---------------------------------------------------------------------------
class AgentState(TypedDict):
    task: str
    result: str | None
    attempts: int
    status: str  # pending | success | failed | max_attempts
    error: str | None
    max_attempts: int

# ---------------------------------------------------------------------------
# Tool – unreliable arithmetic evaluator
# ---------------------------------------------------------------------------
def unreliable_tool(expr: str) -> str:
    """Evaluate a simple arithmetic expression.

    With 30 % probability it raises ValueError to simulate an external
    failure.  The function is intentionally minimal – the goal is to
    demonstrate retry logic, not complex parsing.
    """
    if random.random() < 0.3:
        raise ValueError("Simulated evaluation error")
    try:
        # Safe eval: only arithmetic operators are allowed.
        result = eval(expr, {"__builtins__": None}, {})
    except Exception as exc:
        raise ValueError(f"Invalid expression: {expr}") from exc
    return str(result)

# ---------------------------------------------------------------------------
# Node implementations
# ---------------------------------------------------------------------------
async def execute_task(state: AgentState) -> AgentState:
    """Run the unreliable tool and store its output.

    The node updates ``result`` and ``error`` fields.  If the tool raises an
    exception, the error message is captured and ``status`` is set to
    ``failed`` – this will trigger a retry.
    """
    expr = state["task"]
    try:
        result = unreliable_tool(expr)
        state.update(result=result, error=None, status="pending")
    except Exception as exc:  # pragma: no cover – exercised via retries
        state.update(result=None, error=str(exc), status="failed")
    return state

async def verify_result(state: AgentState) -> AgentState:
    """Ask the LLM to judge whether the result is correct.

    The prompt explicitly asks for a single word answer – ``success`` or
    ``failed``.  Any other response is treated as failure.
    """
    if state["result"] is None:
        # No result – treat as failed to trigger retry logic.
        return state

    prompt = (
        f"You are a judge evaluating the correctness of an arithmetic\n"
        f"expression: {state['task']}\n"
        f"Result produced by the agent: {state['result']}\n"
        "Is this result correct? Respond with only one word: success or failed."
    )
    msg = HumanMessage(content=prompt)
    response = await llm.ainvoke([msg])
    verdict = response.content.strip().lower()
    if verdict == "success":
        state.update(status="success")
    else:
        state.update(status="failed")
    return state

async def handle_error(state: AgentState) -> AgentState:
    """Increment attempt counter and decide whether to retry.

    If the maximum number of attempts is reached, set status to
    ``max_attempts``.  Otherwise reset error and result for a fresh run.
    """
    state["attempts"] += 1
    if state["attempts"] >= state["max_attempts"]:
        state.update(status="max_attempts", error=None, result=None)
    else:
        # Prepare for retry: clear previous result and error.
        state.update(error=None, result=None, status="pending")
    return state

# ---------------------------------------------------------------------------
# Graph construction
# ---------------------------------------------------------------------------
builder = StateGraph(AgentState)
builder.add_node("execute_task", execute_task)
builder.add_node("verify_result", verify_result)
builder.add_node("handle_error", handle_error)

builder.set_entry_point("execute_task")
builder.add_edge("execute_task", "verify_result")
builder.add_conditional_edges(
    "verify_result",
    lambda x: {
        "success": END,
        "failed": "handle_error",
        "max_attempts": END,
    }[x["status"]],
)
builder.add_edge("handle_error", "execute_task")

graph = builder.compile(checkpointer=MemorySaver())

# ---------------------------------------------------------------------------
# CLI helper – run a single task and print progress table
# ---------------------------------------------------------------------------
async def run_single(expr: str) -> None:
    initial_state: AgentState = {
        "task": expr,
        "result": None,
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": MAX_ATTEMPTS,
    }

    state = initial_state
    table = Table(title=f"Self‑correcting agent – {expr}")
    table.add_column("Attempt", justify="right")
    table.add_column("Result")
    table.add_column("Status")
    table.add_column("Error")

    while state["status"] not in ("success", "max_attempts"):
        state = await graph.ainvoke(state)
        attempt_no = state["attempts"] + 1 if state["status"] == "failed" else state["attempts"]
        table.add_row(
            str(attempt_no),
            str(state.get("result", "")),
            state["status"],
            state.get("error", "") or "",
        )

    console.print(table)
    if state["status"] == "success":
        console.print(f"✅ Success in {state['attempts'] + 1} attempt(s). Result: {state['result']}\n")
    else:
        console.print("❌ Max attempts reached without success.\n")

# ---------------------------------------------------------------------------
# Main entry point – run three examples
# ---------------------------------------------------------------------------
if __name__ == "__main__":  # pragma: no cover – manual execution only
    import sys

    if len(sys.argv) > 1:
        expressions = [" ".join(sys.argv[1:])]
    else:
        expressions = ["2+2", "10/3", "5*6-7"]

    for expr in expressions:
        import asyncio

        asyncio.run(run_single(expr))