source .venv/bin/activate                      # the environment from Module 1
cd supportdesk
pip install scikit-learn==1.9.1                # used by the ensembling example in Part B
mkdir -p runs
PYTHONPATH=. python examples/m05_baselines.py

"""Shared helpers for Module 5: the refund-decision task, its rule verifier, prompts, and scoring.

The refund policy comes from the billing-refunds help-center article. Each
scenario is a customer message plus the facts a billing system would hold
about that customer. The gold decision is computed from those facts by
`decide`, which is also the external verifier used later in the module.
"""
from __future__ import annotations

import math
import re
from collections.abc import Callable
from dataclasses import dataclass
from datetime import date

from supportdesk.data import CATEGORIES, get_article
from supportdesk.llm import Usage

TODAY = date(2026, 9, 21)  # a Monday; every scenario is judged as of this day
LABELS = ("refund_duplicate", "refund_full", "no_refund", "needs_owner")
REFUND_ROLES = ("owner", "billing_admin")


@dataclass(frozen=True)
class Facts:
    """What the billing system knows about the charge in question."""
    role: str            # owner, billing_admin, admin, member
    billing: str         # monthly or annual
    charged_on: date     # date of the purchase or renewal being disputed
    duplicate: bool      # True if the same invoice was charged twice


@dataclass(frozen=True)
class RefundCase:
    id: str
    text: str
    facts: Facts

    @property
    def gold(self) -> str:
        return decide(self.facts)


def decide(f: Facts, today: date = TODAY) -> str:
    """The refund policy as code. Rules apply in order; the first match wins."""
    if f.role not in REFUND_ROLES:
        return "needs_owner"          # only owners and billing admins may request refunds
    if f.duplicate:
        return "refund_duplicate"     # duplicate charges are always refunded
    if f.billing == "annual" and (today - f.charged_on).days <= 14:
        return "refund_full"          # annual plans within 14 days of purchase or renewal
    return "no_refund"                # monthly plans, or annual after 14 days


def _c(i: int, text: str, role: str, billing: str, y: int, m: int, d: int, dup: bool = False) -> RefundCase:
    return RefundCase(f"R{i:02d}", text, Facts(role, billing, date(y, m, d), dup))


CASES: list[RefundCase] = [
    _c(1, "I'm the workspace owner. Our card was charged twice for the Team plan on 3 September, same invoice number. Please refund one.", "owner", "monthly", 2026, 9, 3, True),
    _c(2, "Owner here. We renewed our annual Business plan on 16 September by mistake. Can we get our money back?", "owner", "annual", 2026, 9, 16),
    _c(3, "I own the workspace. We bought the annual Team plan on 10 July and want to stop now. Refund for the rest of the year?", "owner", "annual", 2026, 7, 10),
    _c(4, "I'm a regular member of our workspace and I noticed we were charged twice this month. Can you refund the duplicate?", "member", "monthly", 2026, 9, 3, True),
    _c(5, "Billing admin here. Our monthly Team plan charged us on 18 September but we barely use it. Refund please, it was only 3 days ago.", "billing_admin", "monthly", 2026, 9, 18),
    _c(6, "As the workspace owner: our annual Business plan renewed on 7 September. We'd like to cancel and get a refund.", "owner", "annual", 2026, 9, 7),
    _c(7, "Workspace owner. Our annual plan renewed on 6 September and we want a refund.", "owner", "annual", 2026, 9, 6),
    _c(8, "I'm the billing admin. We started the annual Team plan on 30 August. It is not working for us, can we be refunded?", "billing_admin", "annual", 2026, 8, 30),
    _c(9, "I'm the owner. We first subscribed in September 2025 and the annual plan auto-renewed on 12 September. Please refund the renewal.", "owner", "annual", 2026, 9, 12),
    _c(10, "I'm a project manager in our workspace (not the owner). Our annual plan renewed 3 days ago and we want a refund.", "member", "annual", 2026, 9, 18),
    _c(11, "Owner here. When our annual plan renewed on 12 August the card was charged twice for the same invoice. I know it is past 14 days, but can we get the duplicate back?", "owner", "annual", 2026, 8, 12, True),
    _c(12, "Soy la propietaria del espacio de trabajo. Renovamos el plan anual hace 20 días. ¿Podemos cancelar y recibir el reembolso?", "owner", "annual", 2026, 9, 1),
    _c(13, "Ich bin Billing-Admin. Unser Monatsplan wurde am 3. September doppelt abgebucht (gleiche Rechnung). Bitte erstatten Sie die Doppelbuchung.", "billing_admin", "monthly", 2026, 9, 3, True),
    _c(14, "I'm the owner. We pay monthly for Business, charged on 1 September. We're cancelling today; please refund the unused part of the month.", "owner", "monthly", 2026, 9, 1),
    _c(15, "Owner of the workspace. I bought the annual Team plan last Monday and my team hates it. Can I get a refund?", "owner", "annual", 2026, 9, 14),
    _c(16, "I'm a workspace admin (I manage members and boards). Our annual plan renewed 5 days ago; please refund it.", "admin", "annual", 2026, 9, 16),
    _c(17, "Billing admin. Our annual Team plan renewed on 8 September. We want to cancel and get our money back.", "billing_admin", "annual", 2026, 9, 8),
    _c(18, "I'm the owner. I saw two charges of 144 USD on 3 September and thought it was a duplicate, but they are invoices for our two separate workspaces. Can I still get one refunded? We pay monthly.", "owner", "monthly", 2026, 9, 3),
    _c(19, "ワークスペースのオーナーです。10日前に年間プランが更新されました。返金してもらえますか？", "owner", "annual", 2026, 9, 11),
    _c(20, "Billing admin here. Our annual Team plan renewed on 1 September and nobody has logged in since. Refund?", "billing_admin", "annual", 2026, 9, 1),
]


# Prompts ---------------------------------------------------------------------

def policy_text() -> str:
    return get_article("billing-refunds").body


INSTRUCTION = (
    "You decide refund eligibility for Brightlane support, strictly by the policy below.\n"
    f"Today is {TODAY:%A %d %B %Y}.\n"
    "Possible decisions: refund_duplicate, refund_full, no_refund, needs_owner "
    "(needs_owner means the requester is not a workspace owner or billing admin).\n\n"
    "<policy>\n{policy}\n</policy>\n\n<ticket>\n{ticket}\n</ticket>\n"
)
DECISION_LINE = "End with one line exactly like: DECISION: <decision>"

DEMOS = (
    "Example ticket: I'm the billing admin. Our annual Business plan renewed on 10 September. Refund please.\n"
    "Reasoning: The requester is a billing admin, so they may request refunds. No duplicate charge is mentioned. "
    "The plan is annual and renewed 11 days before today, which is within 14 days, so a full refund applies.\n"
    "DECISION: refund_full\n\n"
    "Example ticket: I'm a member of the workspace. We were charged twice on 2 September.\n"
    "Reasoning: The requester is a regular member, not an owner or billing admin, so they cannot request a refund, "
    "even for a duplicate charge. The owner must ask.\n"
    "DECISION: needs_owner\n"
)

SCAFFOLD = (
    "Work through these steps in order and fill in every line:\n"
    "ROLE: <owner | billing_admin | admin | member>\n"
    "DUPLICATE: <yes | no>\n"
    "BILLING: <monthly | annual>\n"
    "CHARGED_ON: <YYYY-MM-DD>\n"
    "DAYS_SINCE: <whole days from CHARGED_ON to today>\n"
    "DECISION: <decision>  (apply the rules in this order: role, duplicate, annual within 14 days, otherwise no_refund)\n"
)


def _user(case: RefundCase, tail: str) -> list[dict]:
    body = INSTRUCTION.format(policy=policy_text(), ticket=case.text)
    return [{"role": "user", "content": body + "\n" + tail}]


def prompt_direct(case: RefundCase) -> list[dict]:
    return _user(case, "Answer with the decision only. " + DECISION_LINE)


def prompt_zero_shot_cot(case: RefundCase) -> list[dict]:
    return _user(case, "Think step by step before you answer. " + DECISION_LINE)


def prompt_demonstrated(case: RefundCase) -> list[dict]:
    return _user(case, DEMOS + "\nNow reason the same way about the ticket above. " + DECISION_LINE)


def prompt_scaffold(case: RefundCase) -> list[dict]:
    return _user(case, SCAFFOLD)


PROMPTS: dict[str, Callable[[RefundCase], list[dict]]] = {
    "direct": prompt_direct, "zero_shot_cot": prompt_zero_shot_cot,
    "demonstrated": prompt_demonstrated, "scaffold": prompt_scaffold,
}

# Parsing and scoring ------------------------------------------------------------

DECISION_RE = re.compile(r"^\s*\**DECISION\**\s*:\s*\**\s*([a-z_]+)", re.IGNORECASE | re.MULTILINE)


def parse_decision(text: str) -> str | None:
    """The last 'DECISION: <label>' line, if it names a known label."""
    found = DECISION_RE.findall(text or "")
    label = found[-1].lower() if found else None
    return label if label in LABELS else None


def wilson(k: int, n: int, z: float = 1.96) -> tuple[float, float]:
    """95 percent Wilson score interval for k successes out of n."""
    if n == 0:
        return (0.0, 1.0)
    p = k / n
    centre = (p + z * z / (2 * n)) / (1 + z * z / n)
    half = z * math.sqrt(p * (1 - p) / n + z * z / (4 * n * n)) / (1 + z * z / n)
    return (max(0.0, centre - half), min(1.0, centre + half))


@dataclass
class EvalResult:
    correct: int
    n: int
    unparsed: int
    usage: Usage
    wrong: list[tuple[str, str | None, str]]  # (case id, predicted, gold)

    def line(self, name: str) -> str:
        lo, hi = wilson(self.correct, self.n)
        return (f"{name:<14} {self.correct:>2}/{self.n}  95% CI {lo:.2f}-{hi:.2f}  unparsed {self.unparsed}  "
                f"tokens in/out {self.usage.input_tokens}/{self.usage.output_tokens}")


def evaluate(chat_fn: Callable, build: Callable[[RefundCase], list[dict]], cases: list[RefundCase] = CASES,
             **chat_kwargs) -> EvalResult:
    """Run one prompt style over the cases and score the parsed decision against the rule."""
    usage, correct, unparsed, wrong = Usage(), 0, 0, []
    for case in cases:
        result = chat_fn(build(case), **chat_kwargs)
        usage.input_tokens += result.usage.input_tokens
        usage.output_tokens += result.usage.output_tokens
        usage.reasoning_tokens += result.usage.reasoning_tokens
        label = parse_decision(result.text)
        unparsed += label is None
        if label == case.gold:
            correct += 1
        else:
            wrong.append((case.id, label, case.gold))
    return EvalResult(correct, len(cases), unparsed, usage, wrong)


def case_from_messages(messages: list[dict]) -> RefundCase:
    """Find which scenario a prompt is about (used by scripted stand-ins, never by real code)."""
    content = messages[-1]["content"]
    for case in CASES:
        if case.text in content:
            return case
    raise KeyError("no scenario text found in prompt")


# Classical baselines --------------------------------------------------------------

def shortcut_baseline(text: str) -> str:
    """One step, no intermediate facts: react to the most salient keyword."""
    t = text.lower()
    if re.search(r"twice|duplicate|doppel|duplicado", t):
        return "refund_duplicate"
    if re.search(r"annual|anual|年間", t):
        return "refund_full"
    return "no_refund"


MONTHS = {m: i for i, m in enumerate(["january", "february", "march", "april", "may", "june", "july",
                                      "august", "september", "october", "november", "december"], 1)}


def extract_facts_regex(text: str) -> Facts | None:
    """Pull the four facts out of English text with regular expressions; None if any is missing."""
    t = text.lower()
    if re.search(r"billing[- ]admin", t):
        role = "billing_admin"
    elif re.search(r"\bnot the owner\b|\bmember\b|project manager", t):
        role = "member"
    elif re.search(r"\bworkspace admin\b", t):
        role = "admin"
    elif re.search(r"\bowner\b|\bi own\b", t):
        role = "owner"
    else:
        return None
    billing = "annual" if "annual" in t else "monthly" if re.search(r"monthly|month", t) else None
    duplicate = bool(re.search(r"twice|duplicate", t)) and "separate" not in t
    charged_on = None
    if m := re.search(r"(\d{1,2}) (" + "|".join(MONTHS) + r")", t):
        charged_on = date(TODAY.year, MONTHS[m.group(2)], int(m.group(1)))
    elif m := re.search(r"(" + "|".join(MONTHS) + r") (\d{1,2})\b", t):
        charged_on = date(TODAY.year, MONTHS[m.group(1)], int(m.group(2)))
    elif m := re.search(r"(\d+) days ago", t):
        charged_on = date.fromordinal(TODAY.toordinal() - int(m.group(1)))
    if billing is None or charged_on is None:
        return None
    return Facts(role, billing, charged_on, duplicate)


# Ticket triage by keyword definitions read from a prompt -----------------------------

DEF_LINE = re.compile(r"^- ([a-z_]+): (.+)$", re.MULTILINE)
FALLBACK = re.compile(r"If no category fits, answer: ([a-z_]+)")


def keyword_reader(prompt: str, ticket_text: str) -> str:
    """A deterministic stand-in 'model' that follows a triage prompt's category definitions.

    It counts how many of each category's listed keywords appear in the ticket,
    picks the highest count (ties go to the category listed first), and uses the
    prompt's fallback when nothing matches. It is not a language model; it lets
    prompt search run for real without an API key.
    """
    text = ticket_text.lower()
    best, best_hits = None, 0
    for name, words in DEF_LINE.findall(prompt):
        hits = sum(1 for w in words.split(", ") if w.strip() and w.strip() in text)
        if hits > best_hits:
            best, best_hits = name, hits
    if best is None:
        m = FALLBACK.search(prompt)
        best = m.group(1) if m else "how_to"
    return best


TRIAGE_KEYWORDS: dict[str, list[str]] = {
    "billing": ["charge", "invoice", "price", "cost", "vat", "pay", "tax", "discount"],
    "cancellation": ["cancel", "refund", "money back"],
    "account_access": ["password", "locked", "2fa", "sign in", "saml", "gdpr", "region"],
    "bug": ["not loading", "stopped", "fails", "error", "bug", "down", "lost my changes"],
    "how_to": ["how do", "how many", "where", "export", "can i", "will the"],
    "feature_request": ["please add", "integration", "roadmap", "when will", "support"],
}


def triage_prompt(keywords: dict[str, list[str]], fallback: str = "how_to") -> str:
    """A triage prompt whose category definitions are keyword lists (order matters for ties)."""
    lines = "\n".join(f"- {name}: {', '.join(words)}" for name, words in keywords.items())
    return ("Classify the Brightlane support ticket into exactly one category.\n"
            f"Categories and the words that signal them:\n{lines}\n"
            f"If no category fits, answer: {fallback}\n"
            "Reply with the category name only.")


def triage_messages(prompt: str, ticket_text: str) -> list[dict]:
    return [{"role": "system", "content": prompt}, {"role": "user", "content": ticket_text}]


def keyword_responder(messages: list[dict], kwargs: dict) -> str:
    """ScriptedLLM responder: apply keyword_reader to the system prompt and the ticket."""
    return keyword_reader(messages[0]["content"], messages[-1]["content"])


def parse_category(text: str) -> str | None:
    word = (text or "").strip().strip(".`'\"").lower()
    return word if word in CATEGORIES else None

"""Part A: the refund task, its gold labels, and two classical baselines (one step vs several steps)."""
from collections import Counter

from examples.m05_helpers import CASES, TODAY, decide, extract_facts_regex, shortcut_baseline, wilson

print(f"{len(CASES)} scenarios judged as of {TODAY}; gold labels: {dict(Counter(c.gold for c in CASES))}\n")
print(f"{'id':<4} {'gold':<17} {'shortcut':<17} {'extract+rule':<17} days")
scores = {"shortcut": 0, "extract+rule": 0}
abstained = 0
for case in CASES:
    one_step = shortcut_baseline(case.text)
    facts = extract_facts_regex(case.text)
    stepwise = decide(facts) if facts else "(no facts)"
    abstained += facts is None
    scores["shortcut"] += one_step == case.gold
    scores["extract+rule"] += stepwise == case.gold
    days = (TODAY - case.facts.charged_on).days
    print(f"{case.id:<4} {case.gold:<17} {one_step:<17} {stepwise:<17} {days}")

print()
for name, k in scores.items():
    lo, hi = wilson(k, len(CASES))
    print(f"{name:<13} {k:>2}/{len(CASES)} correct  95% CI {lo:.2f}-{hi:.2f}")
print(f"extract+rule abstained on {abstained} and was wrong on "
      f"{len(CASES) - scores['extract+rule'] - abstained} of the cases it answered")

"""Part A: four ways to ask for the refund decision, their token cost, and a harness to score them.

Run without arguments to check the harness with scripted stand-ins (plumbing
only, no model quality). Run with --live to score a real model:
    LLM_PROVIDER=groq GROQ_API_KEY=... PYTHONPATH=. python examples/m05_cot_prompts.py --live
"""
import sys

from examples.m05_helpers import CASES, PROMPTS, TODAY, case_from_messages, decide, evaluate
from supportdesk.llm import Usage
from supportdesk.pricing import cost_usd
from supportdesk.stand_in import ScriptedLLM
from supportdesk.tokens import count_messages

# 1. Prompt size per style (input tokens are measured; output lengths are assumptions to replace).
ASSUMED_OUTPUT = {"direct": 8, "zero_shot_cot": 180, "demonstrated": 120, "scaffold": 60}
print(f"{'style':<14} {'input tok':>9} {'assumed out':>11} {'USD/1k gpt-oss-120b':>20} {'USD/1k gemini-3.5-flash':>24}")
for name, build in PROMPTS.items():
    tokens_in = count_messages(build(CASES[0]))
    usage = Usage(input_tokens=tokens_in, output_tokens=ASSUMED_OUTPUT[name])
    print(f"{name:<14} {tokens_in:>9} {ASSUMED_OUTPUT[name]:>11} "
          f"{1000 * cost_usd(usage, 'openai/gpt-oss-120b'):>20.3f} {1000 * cost_usd(usage, 'gemini-3.5-flash'):>24.3f}")


# 2. Plumbing check: a scripted responder that fills the scaffold from the true facts.
def oracle(messages, kwargs):
    f = case_from_messages(messages).facts
    return (f"ROLE: {f.role}\nDUPLICATE: {'yes' if f.duplicate else 'no'}\nBILLING: {f.billing}\n"
            f"CHARGED_ON: {f.charged_on}\nDAYS_SINCE: {(TODAY - f.charged_on).days}\nDECISION: {decide(f)}")


# 3. Format drift: the same answers, but a third of them written in a way the parser must reject.
def drifting(messages, kwargs):
    text = oracle(messages, kwargs)
    case = case_from_messages(messages)
    if int(case.id[1:]) % 3 == 0:
        label = text.rsplit("DECISION: ", 1)[1]
        return text.rsplit("\n", 1)[0] + f"\nDecision - {label.replace('_', ' ').title()}"
    if int(case.id[1:]) % 3 == 1:
        return text.replace("DECISION:", "**DECISION:**")
    return text


print()
for name, responder in (("oracle", oracle), ("drifting", drifting)):
    print(evaluate(ScriptedLLM(responder=responder), PROMPTS["scaffold"]).line(name))

# 4. The real measurement, when a key is set.
if "--live" in sys.argv:
    from supportdesk.llm import chat
    for name, build in PROMPTS.items():
        result = evaluate(chat, build, temperature=0.0, max_tokens=2000)
        print(result.line(name), " wrong:", result.wrong)

ROLE: member
DUPLICATE: yes
BILLING: monthly
CHARGED_ON: 2026-09-03
DAYS_SINCE: 18
DECISION: needs_owner