"""Урок 2 - планувальник і контроль (тріаж підтримки).

Продовження агента з Уроку 1. Додано: план перед першим викликом інструмента,
прогрес замість лічильника повторів, переплани з бюджетом і ескалація зі
звітом замість мовчазного STOP. Інструменти макетні, тож цикл безпечно ганяти
офлайн. Потрібен ANTHROPIC_API_KEY.
"""

import json
import time
from dataclasses import dataclass, field
from typing import Callable

from anthropic import Anthropic

MODEL = "claude-opus-5"
client = Anthropic()


# --- Макетні дані ----------------------------------------------------------

TICKETS = {
    "T-1042": "Subject: Export stuck. Body: The export button spins forever. customer_id=C-4471",
    "T-2077": "Subject: Empty export file. Body: Export finishes but the file has no rows. customer_id=C-9130",
}
SUBSCRIPTIONS = {
    "C-4471": "plan=pro, status=active, seats=12",
}
SYSTEMS = {
    "billing": "up",
    "export": "degraded - elevated error rate since 14:02",
    "auth": "up",
    "email": "up",
}


# --- Інструменти -----------------------------------------------------------
# Пошук (lookups) читають світ. Дії (actions) - термінальні. Нова дія:
# escalate - явна передача людині зі звітом, а не мовчазний обрив.

def read_ticket(ticket_id: str) -> str:
    return TICKETS.get(ticket_id, "ticket not found")


def get_subscription(customer_id: str) -> str:
    return f"{customer_id}: {SUBSCRIPTIONS.get(customer_id, 'no subscription on file')}"


def check_system(name: str) -> str:
    return f"{name}: {SYSTEMS.get(name, 'unknown')}"


def assign(team: str) -> str:
    return f"assigned to {team}"


def request_details(question: str) -> str:
    return f"asked customer: {question}"


def close_ticket(resolution: str) -> str:
    return f"closed: {resolution}"


def escalate(reason: str) -> str:
    return f"escalated to human: {reason}"


LOOKUPS = ["read_ticket", "get_subscription", "check_system"]
TERMINAL = {"assign", "request_details", "close_ticket", "escalate"}


@dataclass
class Tool:
    fn: Callable[..., str]
    schema: dict


def _tool(fn, description, properties, required) -> Tool:
    return Tool(fn, {
        "name": fn.__name__,
        "description": description,
        "input_schema": {
            "type": "object",
            "properties": properties,
            "required": required,
            "additionalProperties": False,
        },
    })


TOOLS: dict[str, Tool] = {
    "read_ticket": _tool(read_ticket, "Read a support ticket by id.", {
        "ticket_id": {"type": "string"},
    }, ["ticket_id"]),
    "get_subscription": _tool(get_subscription, "Look up a customer's plan and status.", {
        "customer_id": {"type": "string"},
    }, ["customer_id"]),
    "check_system": _tool(check_system, "Check the status of one system (billing, export, auth, email).", {
        "name": {"type": "string"},
    }, ["name"]),
    "assign": _tool(assign, "Route the ticket to a team and stop.", {
        "team": {"type": "string"},
    }, ["team"]),
    "request_details": _tool(request_details, "Ask the customer for more info and stop.", {
        "question": {"type": "string"},
    }, ["question"]),
    "close_ticket": _tool(close_ticket, "Close the ticket with a resolution and stop.", {
        "resolution": {"type": "string"},
    }, ["resolution"]),
    "escalate": _tool(escalate, "Hand the ticket to a human with a reason and stop.", {
        "reason": {"type": "string"},
    }, ["reason"]),
}

# Мета-інструмент: не чіпає світ, а повертає керування планувальнику.
REPLAN = {
    "name": "replan",
    "description": "The current plan no longer fits the facts. Explain why and request a new plan.",
    "input_schema": {
        "type": "object",
        "properties": {"reason": {"type": "string"}},
        "required": ["reason"],
        "additionalProperties": False,
    },
}


# --- План -------------------------------------------------------------------

PLAN_SCHEMA = {
    "type": "object",
    "properties": {
        "steps": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "tool": {"type": "string", "enum": LOOKUPS},
                    "goal": {"type": "string"},
                },
                "required": ["tool", "goal"],
                "additionalProperties": False,
            },
        },
        "done_when": {"type": "string"},
    },
    "required": ["steps", "done_when"],
    "additionalProperties": False,
}


@dataclass
class Plan:
    steps: list[dict]
    done_when: str
    cursor: int = 0

    def current(self) -> str | None:
        if self.cursor >= len(self.steps):
            return None
        return self.steps[self.cursor]["tool"]

    def pretty(self) -> str:
        lines = [f"{i + 1}. {s['tool']} - {s['goal']}" for i, s in enumerate(self.steps)]
        return "\n".join(lines) + f"\ndone_when: {self.done_when}"


def make_plan(ticket_id: str, state: "State", why: str = "") -> Plan:
    state.llm_calls += 1
    r = client.messages.create(
        model=MODEL,
        max_tokens=1024,
        output_config={"format": {"type": "json_schema", "schema": PLAN_SCHEMA}},
        messages=[{"role": "user", "content": (
            f"Plan support triage for ticket {ticket_id}.\n"
            f"Known facts: {state.facts}\n"
            f"Lookup tools: {LOOKUPS}. Actions afterwards: {sorted(TERMINAL)}.\n"
            "List the lookup steps in order with a goal for each, and state "
            "done_when: the minimal evidence that is enough to act."
            + (f"\nThe previous plan failed: {why}" if why else "")
        )}],
    )
    data = json.loads(next(b.text for b in r.content if b.type == "text"))
    return Plan(steps=data["steps"], done_when=data["done_when"])


# --- Бюджети і state --------------------------------------------------------

@dataclass
class Budget:
    max_steps: int = 8         # стеля викликів інструментів
    max_replans: int = 2       # третій зрив плану = ескалація
    max_no_progress: int = 2   # два кроки без нових фактів = міняємо підхід
    deadline_s: float = 60.0   # тріаж не тримає чергу


@dataclass
class State:
    calls: list[str] = field(default_factory=list)
    facts: dict[str, str] = field(default_factory=dict)
    plans: list[Plan] = field(default_factory=list)
    action: str | None = None
    replans: int = 0
    no_progress: int = 0
    llm_calls: int = 0


# --- Compression: спостереження -> типізовані факти -------------------------

EXTRACT_SCHEMA = {
    "type": "object",
    "properties": {
        "facts": {
            "type": "array",
            "items": {
                "type": "object",
                "properties": {
                    "key": {"type": "string"},
                    "value": {"type": "string"},
                },
                "required": ["key", "value"],
                "additionalProperties": False,
            },
        },
    },
    "required": ["facts"],
    "additionalProperties": False,
}


def extract(observation: str, state: "State") -> list[dict]:
    state.llm_calls += 1
    r = client.messages.create(
        model=MODEL,
        max_tokens=512,
        output_config={"format": {"type": "json_schema", "schema": EXTRACT_SCHEMA}},
        messages=[{"role": "user", "content": f"Extract facts (key/value) from this tool output:\n{observation}"}],
    )
    text = next(b.text for b in r.content if b.type == "text")
    return json.loads(text)["facts"]


# --- Вибір кроку: меню = поточний крок плану + дії + replan -----------------

def choose(ticket_id: str, plan: Plan, state: State) -> tuple[str, dict]:
    state.llm_calls += 1
    allowed = ([plan.current()] if plan.current() else []) + sorted(TERMINAL)
    r = client.messages.create(
        model=MODEL,
        max_tokens=1024,
        tools=[TOOLS[name].schema for name in allowed] + [REPLAN],
        tool_choice={"type": "any"},
        messages=[{"role": "user", "content": (
            f"You are a support triage agent working ticket {ticket_id}.\n"
            f"Plan:\n{plan.pretty()}\n"
            f"Known facts: {state.facts}\n"
            "Follow the plan step by step. Act as soon as done_when is satisfied. "
            "If the facts contradict the plan, call replan and say why."
        )}],
    )
    block = next(b for b in r.content if b.type == "tool_use")
    return block.name, block.input


# --- Ескалація зі звітом (STOP і DONE - не одне й те саме) -------------------

def give_up(state: State, plan: Plan, reason: str) -> State:
    report = (
        f"reason: {reason}; "
        f"plan: {[s['tool'] for s in plan.steps]}; done_when: {plan.done_when}; "
        f"facts: {state.facts}; calls: {state.calls}"
    )
    state.action = f"escalate -> {escalate(report)}"
    return state


def replan(ticket_id: str, state: State, budget: Budget, why: str) -> Plan | None:
    if state.replans >= budget.max_replans:
        give_up(state, state.plans[-1], f"replan budget exhausted: {why}")
        return None
    state.replans += 1
    state.no_progress = 0
    plan = make_plan(ticket_id, state, why)
    state.plans.append(plan)
    return plan


# --- Процес -----------------------------------------------------------------

def run(ticket_id: str, budget: Budget = Budget()) -> State:
    state = State()
    started = time.monotonic()
    plan = make_plan(ticket_id, state)
    state.plans.append(plan)
    while True:
        if len(state.calls) >= budget.max_steps:
            return give_up(state, plan, "step budget exhausted")
        if time.monotonic() - started > budget.deadline_s:
            return give_up(state, plan, "deadline exceeded")
        if state.no_progress >= budget.max_no_progress:
            plan = replan(ticket_id, state, budget, "no new facts for two steps in a row")
            if plan is None:
                return state
            continue

        name, args = choose(ticket_id, plan, state)
        if name == "replan":
            plan = replan(ticket_id, state, budget, args.get("reason", "model requested"))
            if plan is None:
                return state
            continue

        state.calls.append(f"{name}({args})")
        obs = TOOLS[name].fn(**args)
        if name in TERMINAL:
            state.action = f"{name} -> {obs}"
            return state

        before = dict(state.facts)
        for f in extract(obs, state):
            state.facts[f["key"]] = f["value"]
        state.no_progress = 0 if state.facts != before else state.no_progress + 1
        if plan.current() == name:
            plan.cursor += 1


if __name__ == "__main__":
    for ticket_id in ("T-1042", "T-2077"):
        print(f"=== {ticket_id} ===")
        state = run(ticket_id)
        for i, plan in enumerate(state.plans):
            print(f"Plan v{i + 1}:")
            for line in plan.pretty().splitlines():
                print("  " + line)
        print("Call history:")
        for c in state.calls:
            print(" ", c)
        print("Facts:")
        for key, value in state.facts.items():
            print(f"  {key}: {value}")
        print("Action:", state.action)
        print(f"Spent: {len(state.calls)} tool calls, {state.llm_calls} LLM calls, {state.replans} replans")
        print()