"""Урок 5 - multi-agent: оркестратор і чесна ціна команди агентів.

delegate - звичайний інструмент планувальника з Уроку 2; спеціаліст - той
самий агентний цикл з вузьким промтом і своїм корпусом; факти - на спільній
дошці з provenance. Офлайн; ключ ANTHROPIC_API_KEY - лише для demo-тріажу."""

import json
import os
import re
from dataclasses import dataclass, field

from anthropic import Anthropic

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

# Корпуси по доменах: billing-агент фізично не бачить tech-статей, і навпаки.
# TECH_KB - знайомі статті kb-1xx з бази Уроку 3; BILLING_KB - нові білінгові.

BILLING_KB = {
    "kb-201": "Duplicate charge after a plan upgrade: the old subscription stays active until its invoice closes. Refund the smaller charge and cancel the stale subscription, otherwise the customer is billed twice again next cycle.",
    "kb-202": "Refunds take 5-10 business days to reach the customer statement. Do not promise a faster window.",
    "kb-103": "Invoices are generated on the 1st of each month. A missing invoice usually means the subscription was paused, not a failed payment.",
}

TECH_KB = {
    "kb-101": "Export produces an empty file when the report window contains no closed rows. Ask the customer to widen the date range or check the status filter.",
    "kb-102": "Export button spins forever if the export queue is degraded. Check system status before debugging the account.",
    "kb-105": "API rate limits: pro plan allows 600 requests per minute, 429 responses include a retry-after header.",
}

TICKETS = {
    "T-3001": "Upgraded my plan yesterday, there is a duplicate charge on my card, and export gives an empty file.",
    "T-3002": "There is no invoice this month, did my payment fail?",
    "T-3003": "Export button keeps spinning since this morning.",
    "T-3005": "Our legal team needs a GDPR data processing addendum signed.",
}

GOLD = {  # тікет -> факти, без яких рішення неповне
    "T-3001": {"kb-201", "kb-101"},
    "T-3002": {"kb-103"},
    "T-3003": {"kb-102"},
}

def tokenize(text: str) -> list[str]:
    return [t for t in re.findall(r"[a-z0-9]+", text.lower()) if len(t) >= 3]

def search_kb(corpus: dict[str, str], query: str, k: int = 2) -> list[str]:
    # спрощений пошук з Уроку 3: збіг слів плюс поріг релевантності
    q = set(tokenize(query))
    scored = sorted(((len(q & set(tokenize(text))), cid)
                     for cid, text in corpus.items()), reverse=True)
    return [cid for score, cid in scored if score >= 2][:k]


# Звіт замість діалогу: переказ своїми словами губить факти, дошка - ні.

@dataclass
class Fact:
    key: str
    value: str
    source: str  # хто знайшов і де: "billing:kb-201"

@dataclass
class Report:
    specialist: str
    facts: list[Fact]
    recommendation: str
    confidence: float

@dataclass
class Board:
    facts: dict[str, Fact] = field(default_factory=dict)

    def merge(self, report: Report) -> int:
        new = [f for f in report.facts if f.key not in self.facts]
        self.facts.update({f.key: f for f in new})
        return len(new)

    def pretty(self) -> str:
        return "; ".join(f"{f.key}={f.value}" for f in self.facts.values())

@dataclass
class CostMeter:
    llm_calls: int = 0
    tokens: int = 0

    def add(self, prompt: str, output: str) -> None:
        # 1 токен ~ 4 символи: груба оцінка, але для порівняння архітектур чесна
        self.llm_calls += 1
        self.tokens += (len(prompt) + len(output)) // 4


REPORT_RULES = (
    "Report what you found and stop, do not guess the rest. Always return a "
    "structured report: facts with sources, exactly one recommendation, confidence "
    '0..1. Example: {"facts": [{"key": "kb-000", "value": "...", "source": '
    '"agent:kb-000"}], "recommendation": "one action", "confidence": 0.9}'
)

BILLING_PROMPT = (
    "You are the billing specialist. Scope: charges, refunds, invoices, proration, "
    "subscription state. Work only from your own knowledge base and billing tools, "
    "product behavior is not your scope. Quote amounts and invoice ids exactly as "
    "the tools return them, never promise refund timelines shorter than the "
    "policy, escalate charges older than 90 days to a human. " + REPORT_RULES
)

TECH_PROMPT = (
    "You are the tech specialist. Scope: product behavior, exports, api limits, "
    "caching, logs. Work only from your own knowledge base and diagnostic tools, "
    "all of them read-only, billing is not your scope. Prefer the cheapest "
    "diagnostic first; if logs and system status disagree, trust the logs and say "
    "so. Never suggest restarting services, that requires a human. " + REPORT_RULES
)

ORCH_PROMPT = (
    "You are the support triage orchestrator. Read the ticket, split it into "
    "independent domain tasks, delegate each task to one specialist, merge their "
    "reports on the fact board, then make exactly one final decision. Specialists: "
    "billing (charges, refunds, invoices), tech (product behavior, exports, api). "
    "Never forward one specialist's conclusions to another: independent looks "
    "only. If a specialist reports nothing new twice, stop delegating and "
    "escalate with the board as the report. The delegation budget is shared."
)

SINGLE_PROMPT = (
    # агент з Уроку 4; у проді цей промт розпух до трьох екранів
    "You are the support agent. You handle billing, invoices, refunds, exports, "
    "api and logs. Collect facts, then make exactly one decision."
)

@dataclass
class Specialist:
    name: str
    prompt: str
    corpus: dict[str, str]
    tools: tuple[str, ...]

SPECIALISTS = {
    "billing": Specialist("billing", BILLING_PROMPT, BILLING_KB,
                          ("search_kb", "get_subscription", "get_invoices")),
    "tech": Specialist("tech", TECH_PROMPT, TECH_KB,
                       ("search_kb", "check_system", "read_logs")),
}

def can_use(agent: str, tool: str) -> bool:
    # межі доступу з Уроку 1 стають архітектурою: least privilege по агентах
    return agent in SPECIALISTS and tool in SPECIALISTS[agent].tools

LOOKUPS = {
    "get_subscription": "subscription: plan=pro since yesterday, status=active, seats=12, previous=starter, previous_status=still active until its invoice closes, next_invoice=2026-09-01",
    "get_invoices": "invoices: inv-7702 79.00 pro 2026-08-09 paid; inv-7701 49.00 starter 2026-08-09 paid; inv-7640 49.00 starter 2026-07-01 paid",
    "check_system": "systems: export=healthy, export_queue=empty, api=healthy, billing=healthy, auth=healthy, email=healthy",
    "read_logs": "logs, pattern-compressed as in lesson 3: ExportEmpty customer={ID} x3 in 24h; PlanCacheStale customer={ID} x2 after upgrade; RateLimited key={ID} x41 since 09:00; no 5xx in 24h",
}

def observe(agent: str, tool: str, task: str) -> str:
    if not can_use(agent, tool):
        return f"denied by policy: {tool} is outside the {agent} scope"
    if tool != "search_kb":
        return LOOKUPS[tool]
    corpus = SPECIALISTS[agent].corpus
    return ("\n".join(f"[{cid}] {corpus[cid]}" for cid in search_kb(corpus, task))
            or "no relevant documents found")

def run_specialist(spec: Specialist, task: str, meter: CostMeter) -> Report:
    # той самий цикл, що в Уроках 1-2, тільки вузький. Назовні - звіт, не діалог.
    facts: list[Fact] = []
    history = ""
    for tool in spec.tools:
        obs = observe(spec.name, tool, task)
        meter.add(spec.prompt + task + history, obs)   # choose
        meter.add(spec.prompt + obs, "facts")          # extract
        history += "\n" + obs
        if tool == "search_kb":
            facts = [Fact(cid, spec.corpus[cid], f"{spec.name}:{cid}")
                     for cid in search_kb(spec.corpus, task)]
    if not facts:
        return Report(spec.name, [], "out of my scope, return to orchestrator", 0.2)
    meter.add(spec.prompt + task + history, str(facts))  # зібрати звіт
    return Report(spec.name, facts, f"apply {facts[0].key}", 0.9)


@dataclass
class Budget:
    max_delegations: int = 4
    max_no_progress: int = 2
    delegations: int = 0
    no_progress: int = 0

def delegate(name: str, task: str, board: Board, budget: Budget, meter: CostMeter) -> str:
    # інструмент планувальника з Уроку 2; межі знову в коді, а не в промті
    if budget.delegations >= budget.max_delegations:
        return "denied by policy: delegation budget exhausted, escalate"
    if budget.no_progress >= budget.max_no_progress:
        return "denied by policy: two delegations without new facts, escalate"
    budget.delegations += 1
    # незалежність: спеціаліст бачить задачу, а не чужі висновки з дошки
    report = run_specialist(SPECIALISTS[name], task, meter)
    new = board.merge(report)
    budget.no_progress = 0 if new else budget.no_progress + 1
    return (f"{report.specialist}: +{new} facts, confidence "
            f"{report.confidence}, {report.recommendation}")


def route_fanout(ticket: str) -> list[str]:
    # v1: "хай глянуть обидва" - і за кожен погляд платиш (тріаж - у triage_llm)
    return list(SPECIALISTS)

def run_multi(ticket_id: str, router=route_fanout):
    # оркестратор - сам планувальник з Уроку 2: choose і extract щокроку
    ticket = TICKETS[ticket_id]
    board, budget, meter = Board(), Budget(), CostMeter()
    history = ""
    meter.add(ORCH_PROMPT + ticket, "plan")
    contribs = []
    for name in router(ticket):
        before = len(board.facts)
        ctx = ORCH_PROMPT + ticket + history + board.pretty()
        meter.add(ctx, f"delegate({name})")            # choose
        obs = delegate(name, ticket, board, budget, meter)
        meter.add(ORCH_PROMPT + obs, "facts")          # extract звіту
        history += "\n" + obs
        contribs.append((name, len(board.facts) - before))
    meter.add(ORCH_PROMPT + ticket + history + board.pretty(), "final decision")
    return board, meter, contribs


def run_single(ticket_id: str):
    # агент з Уроку 4: один промт, спільний корпус, стислі факти в контексті
    ticket = TICKETS[ticket_id]
    board, meter = Board(), CostMeter()
    corpus = BILLING_KB | TECH_KB
    hits = search_kb(corpus, ticket)
    for cid in hits:
        board.facts[cid] = Fact(cid, corpus[cid], f"single:{cid}")
    facts_line = ""
    for obs in (ticket, "\n".join(f"[{cid}] {corpus[cid]}" for cid in hits),
                LOOKUPS["get_subscription"]):
        meter.add(SINGLE_PROMPT + ticket + facts_line, obs)  # choose
        meter.add(SINGLE_PROMPT + obs, "facts")              # extract
        facts_line = ", ".join(board.facts)
    meter.add(SINGLE_PROMPT + facts_line, "final decision")
    return board, meter, [("single", len(board.facts))]


def evaluate(run_fn) -> tuple[float, float, float]:
    ok = calls = tokens = 0
    for tid, want in GOLD.items():
        board, meter, _ = run_fn(tid)
        ok += want <= set(board.facts)
        calls += meter.llm_calls
        tokens += meter.tokens
    return ok / len(GOLD), calls / len(GOLD), tokens / len(GOLD)


def relay_naive(report: Report) -> str:
    # АНТИПРИКЛАД: переказ без дошки - рекомендація доїжджає, факти ні
    return f"{report.specialist} says: {report.recommendation}"


ROUTE_SCHEMA = {
    "type": "object", "required": ["tasks"], "additionalProperties": False,
    "properties": {"tasks": {"type": "array", "items": {
        "type": "object", "required": ["specialist", "task"],
        "additionalProperties": False,
        "properties": {"specialist": {"type": "string", "enum": list(SPECIALISTS)},
                       "task": {"type": "string"}}}}},
}

def triage_llm(ticket_id: str) -> str:
    # LLM-тріаж: оркестратор ріже тікет на незалежні підзадачі (потрібен ключ)
    r = client.messages.create(
        model=MODEL, max_tokens=512,
        output_config={"format": {"type": "json_schema", "schema": ROUTE_SCHEMA}},
        messages=[{"role": "user", "content": (
            "Split this support ticket into independent domain tasks, one per specialist "
            f"that is actually needed. billing: charges, refunds, invoices. tech: product behavior, exports, api.\nTicket: {TICKETS[ticket_id]}")}],
    )
    data = json.loads(next(b.text for b in r.content if b.type == "text"))
    board, budget, meter = Board(), Budget(), CostMeter()
    for item in data["tasks"]:
        print(f"delegate({item['specialist']}): {item['task']}")
        print("  " + delegate(item["specialist"], item["task"], board, budget, meter))
    return f"facts on board: {', '.join(board.facts)}"


if __name__ == "__main__":
    print("=== Success і вартість: single проти multi на одному GOLD ===")
    _, _, base = evaluate(run_single)
    print(f"{'конфігурація':<16}{'success':>9}{'llm/тікет':>11}{'токени/тікет':>14}{'x single':>10}")
    for cname, fn in [("single", run_single), ("multi fan-out", run_multi)]:
        acc, calls, tokens = evaluate(fn)
        print(f"{cname:<16}{acc:>9.2f}{calls:>11.1f}{tokens:>14.0f}{tokens / base:>9.1f}x")

    print("\n=== Чи додав другий спеціаліст хоч один факт (fan-out) ===")
    for tid in GOLD:
        _, _, contribs = run_multi(tid)
        print(f"{tid}  " + "  ".join(f"{name} +{new}" for name, new in contribs))

    print("\n=== Ping-pong делегація: тікет поза обома доменами ===")
    board, budget, meter = Board(), Budget(), CostMeter()
    for name in ("billing", "tech", "billing", "tech"):
        print(f"delegate({name}) -> {delegate(name, TICKETS['T-3005'], board, budget, meter)}")
    print(f"витрачено даремно: {meter.llm_calls} LLM-викликів, ~{meter.tokens} токенів")

    print("\n=== Зіпсутий телефон: переказ проти дошки ===")
    report = run_specialist(SPECIALISTS["billing"], TICKETS["T-3001"], CostMeter())
    print("наївний переказ:", relay_naive(report))
    print("дошка фактів з provenance:")
    for f in report.facts:
        print(f"  {f.key} <- {f.source}: {f.value}")

    print("\n=== Межі доступу: спеціаліст проситься в чужі інструменти ===")
    print("billing -> read_logs:", observe("billing", "read_logs", ""))
    print("tech -> get_invoices:", observe("tech", "get_invoices", ""))

    if os.environ.get("ANTHROPIC_API_KEY"):
        print("\n=== LLM-тріаж тікета T-3001 ===")
        print(triage_llm("T-3001"))