from typing import TypedDict, Any
import random

# LLM setup – use placeholder values if no specific provider is mentioned
from langchain_openai import ChatOpenAI
from pydantic import SecretStr

llm = ChatOpenAI(
    model="openai/gpt-oss-20b",
    base_url='https://platform.brojs.ru/jrnl-bh/api/inference/v1',
    api_key=SecretStr("jrnl_30283ab953615cbb6846ff9940a1eedce0b76d7b2f59a2394f29e74643e6a90d"),
    temperature=0.2,
)

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

# ---------- Tool ----------
def unreliable_tool(task: str) -> str:
    """Simulate a tool that fails ~30% of the time."""
    if random.random() < 0.3:
        raise ValueError("Tool failure")
    # Very simple evaluation: just return the task string for demo
    return f"Result of '{task}'"

# ---------- Nodes ----------
def execute_task(state: AgentState) -> AgentState:
    try:
        result = unreliable_tool(state["task"])
        state.update(result=result, error=None)
    except Exception as e:
        state.update(result="", error=str(e))
    state.update(status="pending")
    return state

def verify_result(state: AgentState) -> AgentState:
    if state["error"]:
        # If tool failed, skip verification
        state.update(status="failed")
        return state
    prompt = f"Task result: {state['result']}. Is this correct? Respond with 'success' or 'failed'."
    verdict_obj = llm.invoke(prompt)
    # Depending on the LLM implementation, the response may be a string or an object with `content`
    if hasattr(verdict_obj, "content"):
        verdict = verdict_obj.content.strip().lower()
    else:
        verdict = str(verdict_obj).strip().lower()
    if "success" in verdict:
        state.update(status="success")
    else:
        state.update(status="failed")
    return state

def handle_error(state: AgentState) -> AgentState:
    state["attempts"] += 1
    if state["attempts"] >= state["max_attempts"]:
        state.update(status="max_attempts")
    else:
        state.update(status="pending")
    return state

# ---------- Graph ----------
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import InMemorySaver

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")

def _next(state: AgentState) -> str:
    status = state["status"]
    if status == "success":
        return END
    if status == "failed" and state["attempts"] < state["max_attempts"]:
        return "handle_error"
    return END

builder.add_conditional_edges("verify_result", _next)
builder.add_edge("handle_error", "execute_task")

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

# ---------- CLI ----------
def main():
    task = input("Задача: ").strip()
    if not task:
        print("Нет задачи")
        return
    initial_state: AgentState = {
        "task": task,
        "result": "",
        "attempts": 0,
        "status": "pending",
        "error": None,
        "max_attempts": 5,
    }
    state = graph.invoke(initial_state)
    attempts = state["attempts"] + (1 if state["status"] != "failed" else 0)
    print(f"\nИтог: {state['status']} за {attempts} попытки(й)")
    if state["result"]:
        print(f"Результат: {state['result']}")
    if state["error"]:
        print(f"Ошибка: {state['error']}")

if __name__ == "__main__":
    main()