"""Урок 1 - перший самописний агент (тріаж підтримки).

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

import json
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",
}
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) - термінальні: після дії агент
# зупиняється. Усі повертають рядок-спостереження.

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


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


TERMINAL = {"assign", "request_details", "close_ticket"}


@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"]),
}


# --- Граф інструментів -----------------------------------------------------
# Що має сенс далі, зважаючи на останній інструмент. Зверни увагу на петлю
# check_system -> check_system: саме тут самописний агент і зависатиме.

TOOL_GRAPH: dict[str, list[str]] = {
    "read_ticket": ["get_subscription", "check_system"],
    "get_subscription": ["check_system", "assign", "request_details", "close_ticket"],
    "check_system": ["check_system", "get_subscription", "assign", "request_details", "close_ticket"],
    "assign": [],
    "request_details": [],
    "close_ticket": [],
}


# --- State ------------------------------------------------------------------

@dataclass
class State:
    calls: list[str] = field(default_factory=list)
    facts: dict[str, str] = field(default_factory=dict)
    action: str | None = None


# --- 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) -> list[dict]:
    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"]


# --- Вибір наступного інструмента ------------------------------------------
# Передаємо лише дозволену підмножину схем; граф, а не промт, обмежує вибір.

def choose(ticket: str, allowed: list[str], state: State) -> tuple[str, dict]:
    r = client.messages.create(
        model=MODEL,
        max_tokens=1024,
        tools=[TOOLS[name].schema for name in allowed],
        tool_choice={"type": "any"},
        messages=[{"role": "user", "content": (
            "You are a support triage agent. Decide the single most useful next step.\n"
            f"Ticket: {ticket}\n"
            f"Known facts: {state.facts}\n"
            "When you have enough to act, pick an action (assign / request_details / close_ticket). "
            "Do not repeat identical calls."
        )}],
    )
    block = next(b for b in r.content if b.type == "tool_use")
    return block.name, block.input


# --- Цикл ------------------------------------------------------------------

def run(ticket: str, max_steps: int = 8) -> State:
    state = State()
    allowed = ["read_ticket"]
    for _ in range(max_steps):
        name, args = choose(ticket, allowed, state)
        sig = f"{name}({args})"
        if sig in state.calls:            # guard: дослівний повтор = застрягли
            break
        state.calls.append(sig)
        obs = TOOLS[name].fn(**args)
        for f in extract(obs):
            state.facts[f["key"]] = f["value"]
        if name in TERMINAL:
            state.action = f"{name} -> {obs}"
            break
        allowed = TOOL_GRAPH[name]
    return state


if __name__ == "__main__":
    state = run("T-1042")
    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 or "STUCK: stopped without a decision")