"""Can one user's chat leak into another's? A small simulation of two shared caches in front of a shop's help bot.

Python 3 standard library only (plus lesson 23's code, imported, not copied). Run `python3 show_numbers.py` to print every
number the lesson quotes.

The toy, in five rules:
  1. A made-up shop has three made-up users. Each has two private facts, an order code and a street.
  2. The help bot answers four kinds of question. Three answers carry a private fact (an order code or a street) and
     one, opening hours, is the same for everyone.
  3. A response cache sits in front of the bot and keeps whole answers. Its key is either the question alone ("shared")
     or the user plus the question ("per user"). A request whose key is already stored is a hit and gets the stored answer.
  4. A hit is a leak when the stored answer is not the answer this user should get. In this toy every such answer carries
     another user's order code or street.
  5. A prompt cache (lesson 8's idea: a stored start of a prompt is reused whole or not at all) makes the wait before the
     first word shorter when it is hit. The waits come from lesson 23's best-case reading time plus a random queueing delay
     with a fixed seed. Real waits are noisier and real caches store per block; nothing here touches a real service.
"""
from __future__ import annotations

import random
import sys
from dataclasses import dataclass
from pathlib import Path

LESSONS = Path(__file__).resolve().parents[2]
sys.path.append(str(LESSONS / "why-ai-waits-before-the-first-word" / "code"))
import prefill_sim                                        # lesson 23: the best-case time of reading a prompt


# --- the shop, its users and its questions -------------------------------------------------------------------

@dataclass(frozen=True)
class User:
    name: str                   # "A", "B", "C": the widget and the lesson write "User A"
    order: str                  # a made-up order code
    street: str                 # a made-up street, no number, no town


USERS = [User("A", "A-1042", "Mill Lane"), User("B", "B-2210", "Quarry Road"), User("C", "C-3377", "Orchard Close")]
USER_BY_NAME = {u.name: u for u in USERS}

# (question as the user types it, the answer). {order} and {street} are filled from the user.
QUESTIONS = [
    ("Where is my order?", "Order {order} is out for delivery."),
    ("What is my delivery address?", "We deliver to {street}."),
    ("What are your opening hours?", "We open 9 to 5, Monday to Saturday."),
    ("Can you gift-wrap order {order}?", "Yes, order {order} can be gift-wrapped."),
]
PUBLIC = 2                      # the question whose answer is the same for everyone

# (user, question number): fixed, so every run gives the same table.
SEQUENCE = [("A", 0), ("B", 0), ("A", 2), ("C", 2), ("B", 1), ("A", 1), ("C", 0),
            ("A", 0), ("B", 3), ("B", 3), ("C", 1), ("B", 0), ("A", 2), ("C", 0)]


def fill(template: str, user: User) -> str:
    return template.replace("{order}", user.order).replace("{street}", user.street)


def question(user: User, q: int) -> str:
    """The words the user types; question 3 holds the user's own order code, so it differs from user to user."""
    return fill(QUESTIONS[q][0], user)


def answer(user: User, q: int) -> str:
    """The answer this user should get."""
    return fill(QUESTIONS[q][1], user)


# --- the response cache --------------------------------------------------------------------------------------

@dataclass(frozen=True)
class Step:
    n: int                      # request number, from 1
    user: str
    question: str
    key: str                    # what the cache looks up
    hit: bool
    written_by: int | None      # on a hit, the number of the request whose answer is stored
    served: str                 # the answer the user got
    right: str                  # the answer the user should have got
    leak: bool                  # served is not right


def cache_key(user: User, text: str, per_user: bool) -> str:
    return f"{user.name} | {text}" if per_user else text


def simulate(per_user: bool, sequence=SEQUENCE) -> list[Step]:
    """Run the requests through a response cache that never expires. A miss writes the right answer for the asker."""
    store: dict[str, tuple[str, int]] = {}
    steps = []
    for n, (name, q) in enumerate(sequence, 1):
        user, text, right = USER_BY_NAME[name], question(USER_BY_NAME[name], q), answer(USER_BY_NAME[name], q)
        key = cache_key(user, text, per_user)
        if key in store:
            served, by = store[key]
            steps.append(Step(n, name, text, key, True, by, served, right, served != right))
        else:
            store[key] = (right, n)
            steps.append(Step(n, name, text, key, False, None, right, right, False))
    return steps


def percent(part: int, whole: int) -> int:
    """Whole percent, rounded half up, in integers so that the widget's JavaScript prints the same digits."""
    return (200 * part + whole) // (2 * whole)


def totals(steps: list[Step]) -> dict:
    hits, leaks, n = sum(s.hit for s in steps), sum(s.leak for s in steps), len(steps)
    return {"requests": n, "hits": hits, "leaks": leaks, "rate": percent(hits, n)}


# --- the prompt cache and the wait before the first word ------------------------------------------------------

INSTRUCTION = 900               # tokens that start every prompt of the shop's bot, the same for everyone
NOTE = 100                      # tokens of the user's own note, right after them
TAIL = 20                       # tokens of the new question, last
TARGET = "Mill Lane"            # the street in User A's note
GUESSES = ["Elm Road", "Mill Lane", "Fern Close"]   # what User B tries; the middle one is right
RIGHT = GUESSES.index(TARGET)
QUEUE_MS = 60.0                 # mean of the random queueing delay added to every wait, in milliseconds (a toy figure)
TRIES = 25                      # waits measured per guess in one experiment
EXPERIMENTS = 2000
SEED = 2026


def prompt(street: str) -> list:
    """Token labels, equal where the text is equal: the shared instruction, the note about a street, and a fresh question."""
    return [("sys", i) for i in range(INSTRUCTION)] + [(street, i) for i in range(NOTE)] + [("question", i) for i in range(TAIL)]


class PromptCache:
    """Stored starts of prompts. A start is kept under a salt; a request reads a stored start only under the same salt,
    and only if its prompt begins with the whole of it (lesson 8's rule)."""

    def __init__(self):
        self.starts: list[tuple[str, list]] = []

    def store(self, salt: str, start: list) -> None:
        self.starts.append((salt, start))

    def reused(self, salt: str, prompt_tokens: list) -> int:
        return max((len(s) for t, s in self.starts if t == salt and prompt_tokens[:len(s)] == s), default=0)


def wait_ms(total: int, reused: int) -> float:
    """Best-case time to read the tokens that are not reused, from lesson 23's model at the chip's peak, in milliseconds."""
    return prefill_sim.prefill(total - reused).seconds * 1000


def ideal_waits(salted: bool) -> list[float]:
    """The wait of each of User B's three guesses, with no random delay, right after User A sent its note.
    Not salted: both users share one key space. Salted: each user's prompts are stored under that user's name."""
    cache = PromptCache()
    cache.store("A" if salted else "", prompt(TARGET)[:INSTRUCTION + NOTE])
    waits = []
    for street in GUESSES:
        tokens = prompt(street)
        waits.append(wait_ms(len(tokens), cache.reused("B" if salted else "", tokens)))
    return waits


def timing(salted: bool, seed: int = SEED, experiments: int = EXPERIMENTS, tries: int = TRIES) -> dict:
    """User B picks the guess with the shortest wait, either from one try per guess or from the average of `tries`.
    Each experiment draws fresh delays; the same seed gives the same draws in both modes, so only the key differs.
    The cache is as User A left it before every try (a real request would fill the cache itself)."""
    rng, ideal = random.Random(seed), ideal_waits(salted)
    one = many = 0
    sums = [0.0] * len(ideal)
    for _ in range(experiments):
        waits = [[w + rng.expovariate(1 / QUEUE_MS) for _ in range(tries)] for w in ideal]
        first = [w[0] for w in waits]
        means = [sum(w) / tries for w in waits]
        one += first.index(min(first)) == RIGHT
        many += means.index(min(means)) == RIGHT
        sums = [s + m for s, m in zip(sums, means)]
    return {"ideal": ideal, "average": [s / experiments for s in sums], "one": percent(one, experiments), "many": percent(many, experiments)}