Lab 12 — GatewayMeter under fire

"""GatewayMeter stage 12."""

"""GatewayMeter stage 12. See this lab's README and common/flink_job.py."""
from __future__ import annotations

import sys
sys.path.insert(0, "/labs")

from common.flink_job import build_stage, run_stage  # noqa: E402

def build_job(env):
    return build_stage(env, 12)

if __name__ == "__main__":
    run_stage(12)

import pathlib
import sys

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[2]))

import hashlib
import json
import subprocess

import pytest

from common.events import Event, REQUEST_COMPLETED, REQUEST_STARTED, TOKEN_GENERATED
from common.gatewaymeter import aggregate_costs, assemble, merge_partials, partial_costs
from common.invoice_diff import diff
from common.invoice_oracle import materialize
from common.produce import kafka_partition_for_key
from common.replay import schedule

ROOT = pathlib.Path(__file__).resolve().parents[2]
ORIGIN = 1_767_225_600_000

def test_full_invoice_matches_two_stage_aggregation():
    events = []
    for i in range(20):
        events += [
            Event(REQUEST_STARTED, f"r{i}", "acme" if i < 15 else "globex", "gpt-5.5", i, {"prompt_tokens": 100}),
            Event(REQUEST_COMPLETED, f"r{i}", "acme" if i < 15 else "globex", "gpt-5.5", 1000+i, {"completion_tokens": 20}),
        ]
    records, open_requests = assemble(events)
    assert open_requests == {}
    assert merge_partials(partial_costs(records)) == aggregate_costs(records)


def keyed(*events):
    return [f"{event.request_id}\t{event.to_json()}" for event in events]


def produce(*flags):
    out = subprocess.run(
        [sys.executable, str(ROOT / "common/produce.py"), "--seed", "7", "--start-ms", str(ORIGIN),
         "--rate", "20", *flags],
        capture_output=True, text=True, timeout=30, check=True,
    )
    return out.stdout.splitlines()


def test_oracle_prices_like_the_stage_12_job():
    lines = keyed(
        Event(REQUEST_STARTED, "r1", "acme", "gpt-5.5", 1, {"prompt_tokens": 100}),
        Event(REQUEST_COMPLETED, "r1", "acme", "gpt-5.5", 1001, {"completion_tokens": 20}),
        Event(REQUEST_STARTED, "r2", "globex", "gpt-5.5", 1, {"prompt_tokens": 100}),
        Event(REQUEST_COMPLETED, "r2", "globex", "gpt-5.5", 1001, {"completion_tokens": 20}),
    )
    # 1,920 micro-dollars at list price; acme's negotiated 10% discount comes from the same price
    # book the job's price lookup uses from stage 9 on.
    assert materialize(lines) == [
        {"tenant_id": "acme", "window_end": 60000, "micro_usd": 1728, "requests": 1},
        {"tenant_id": "globex", "window_end": 60000, "micro_usd": 1920, "requests": 1},
    ]


def test_oracle_bills_an_abandoned_request_when_its_timer_fires():
    lines = keyed(
        Event(REQUEST_STARTED, "r1", "globex", "gpt-5.5", 1_000, {"prompt_tokens": 100}),
        Event(TOKEN_GENERATED, "r1", "globex", "gpt-5.5", 1_010, {"n": 20}),
    )
    # The assembler's timer fires at start + 5 minutes and bills the tokens seen so far.
    assert materialize(lines) == [
        {"tenant_id": "globex", "window_end": 6 * 60_000, "micro_usd": 1920, "requests": 1}
    ]


def test_oracle_refuses_a_horizon_the_trace_cannot_close():
    # One minute of event time; a tenth of the requests arrive 30 s late.
    lines = produce("--requests", "1200", "--late-pct", "0.1")
    with pytest.raises(ValueError, match="cannot close"):
        materialize(lines, through_ms=ORIGIN + 60_000)
    # The minute before the origin holds the first late requests, and the trace can close it.
    closed = materialize(lines, through_ms=ORIGIN)
    assert closed and all(row["window_end"] == ORIGIN for row in closed)


def test_a_trace_is_reproducible_and_needs_no_clock():
    assert produce("--requests", "300") == produce("--requests", "300")


def test_a_waking_partition_delivers_its_held_requests_late_with_their_timestamps():
    plain = produce("--requests", "400")
    held = produce("--requests", "400", "--stall-partition", "5",
                   "--stall-after-seconds", "4", "--wake-after-seconds", "12")
    # Same requests, same events: only the delivery order changes, so the invoice cannot.
    assert sorted(held) == sorted(plain)
    assert held != plain
    assert materialize(held) == materialize(plain)

    starts = [json.loads(line.split("\t", 1)[1]) for line in held]
    starts = [event for event in starts if event["event_type"] == "RequestStarted"]
    seconds = [(event["ts"] - ORIGIN) / 1000 for event in starts]
    partition = [kafka_partition_for_key(event["request_id"]) for event in starts]
    wake = seconds.index(12.0)
    released = [i for i, (sec, p) in enumerate(zip(seconds, partition)) if p == 5 and 4 <= sec < 12]
    # Partition 5's 4-12 s requests arrive together at 12 s, after everything else before 12 s.
    assert released == list(range(wake - len(released), wake))
    assert all(i < released[0] for i, sec in enumerate(seconds) if sec < 12 and i not in released)


def test_replay_paces_by_the_newest_request_start():
    lines = keyed(
        Event(REQUEST_STARTED, f"r-{ORIGIN}-0000001", "acme", "gpt-5.5", ORIGIN, {}),
        Event(REQUEST_COMPLETED, f"r-{ORIGIN}-0000001", "acme", "gpt-5.5", ORIGIN + 3_000, {}),
        Event(REQUEST_STARTED, f"r-{ORIGIN}-0000002", "acme", "gpt-5.5", ORIGIN + 1_000, {}),
        # Late, or held back by a stalled partition: old timestamps do not move the pace.
        Event(REQUEST_STARTED, f"r-{ORIGIN}-0000003", "acme", "gpt-5.5", ORIGIN - 30_000, {}),
        Event(REQUEST_STARTED, f"r-{ORIGIN}-0000004", "acme", "gpt-5.5", ORIGIN + 2_500, {}),
    )
    assert [at for at, _ in schedule(lines)] == [0.0, 0.0, 1.0, 1.0, 2.5]


def row(tenant, minute, requests, micro):
    return json.dumps({"tenant_id": tenant, "window_end": ORIGIN + minute * 60_000,
                       "requests": requests, "micro_usd": micro})


def test_diff_materializes_corrections_whatever_order_kafka_returns_them_in():
    oracle = [row("acme", 1, 10, 900), row("globex", 1, 5, 500)]
    output = [
        row("acme", 1, 10, 900),       # the correction, read first from another partition
        row("acme", 1, 8, 700),        # the first firing
        row("globex", 1, 5, 500),
        row("globex", 2, 1, 100),      # the trace's tail, past the oracle's horizon
    ]
    report = diff(output, oracle)
    assert report["ok"] and report["matched"] == 2
    assert report["emissions"] == {"total": 3, "corrected_minutes": 1, "repeated_values": 0}


def test_diff_reports_what_does_not_match():
    oracle = [row("acme", 1, 10, 900), row("globex", 1, 5, 500)]
    output = [row("acme", 1, 10, 1000), row("acme", 1, 10, 900), row("initech", 1, 1, 1)]
    report = diff(output, oracle)
    assert not report["ok"]
    assert [c["tenant_id"] for c in report["conflicts"]] == ["acme"]
    assert [m["tenant_id"] for m in report["missing"]] == ["globex"]
    assert [u["tenant_id"] for u in report["unexpected"]] == ["initech"]


def test_the_capstone_trace_is_the_one_the_readme_describes():
    readme = (ROOT / "12-capstone-gatewaymeter/README.md").read_text(encoding="utf-8")
    command = readme.split("python3 common/produce.py ", 1)[1].split(" > evidence/trace.jsonl", 1)[0]
    flags = command.replace("\\\n", " ").split()
    out = subprocess.run([sys.executable, str(ROOT / "common/produce.py"), *flags],
                         capture_output=True, check=True, timeout=30).stdout
    assert hashlib.sha256(out).hexdigest() in readme

"""The event envelope every lab shares.

The gateway emits one JSON object per thing that happens to a request. `ts` is **event time** —
the millisecond the gateway says the thing happened — and it is the only clock the job may bill by.
Kafka's append time is not it, and neither is the moment an operator sees the record.
"""
from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any, Dict

# The six things the gateway reports. Frozen: a lab that needs a new fact adds a payload field,
# never a seventh type.
REQUEST_STARTED = "RequestStarted"          # payload: prompt_tokens
TOKEN_GENERATED = "TokenGenerated"          # payload: n
TOOL_CALLED = "ToolCalled"                  # payload: tool, latency_ms
AGENT_STEP_COMPLETED = "AgentStepCompleted"  # payload: step, tokens_in, tokens_out
MODEL_ERROR = "ModelError"                  # payload: code
REQUEST_COMPLETED = "RequestCompleted"      # payload: status, completion_tokens, latency_ms

EVENT_TYPES = (
    REQUEST_STARTED, TOKEN_GENERATED, TOOL_CALLED,
    AGENT_STEP_COMPLETED, MODEL_ERROR, REQUEST_COMPLETED,
)


@dataclass
class Event:
    event_type: str
    request_id: str
    tenant_id: str
    model: str
    ts: int                                   # event time, ms since epoch
    payload: Dict[str, Any] = field(default_factory=dict)

    def to_json(self) -> str:
        return json.dumps({
            "event_type": self.event_type, "request_id": self.request_id,
            "tenant_id": self.tenant_id, "model": self.model, "ts": self.ts,
            "payload": self.payload,
        })


def parse_event(raw: str) -> Event:
    """JSON line -> Event. A plain map; no state, no side effects, no clock.

    Malformed records raise. Module 2 is where the job decides what to do about that; until then
    a bad record should stop you, not vanish.
    """
    d = json.loads(raw)
    return Event(
        event_type=d["event_type"], request_id=d["request_id"], tenant_id=d["tenant_id"],
        model=d["model"], ts=int(d["ts"]), payload=d.get("payload") or {},
    )

"""Pure GatewayMeter business semantics shared by the labs and their fast tests.

Nothing in this module imports Flink.  A lab can therefore distinguish two questions:

* is the invoice calculation correct for this event trace? and
* does Flink preserve that calculation under partitioning, time, and failure?

The first question belongs in a fast unit test.  The second belongs in the compose cluster.
Money is represented as integer micro-dollars so replay and partial aggregation are deterministic.
"""
from __future__ import annotations

from dataclasses import dataclass, field, replace
from typing import Iterable, Mapping, Sequence
import zlib

from common.events import REQUEST_COMPLETED, REQUEST_STARTED, TOKEN_GENERATED, Event

MINUTE_MS = 60_000
PRICE_PER_1K = {
    "gpt-5.5": (12_000, 36_000),
    "gpt-5.4-mini": (1_000, 4_000),
    "claude-opus-5": (15_000, 75_000),
    "gemini-3.6-flash": (500, 2_000),
}


@dataclass
class OpenRequest:
    request_id: str
    tenant_id: str
    model: str
    started_ts: int
    prompt_tokens: int
    completion_tokens: int = 0


@dataclass(frozen=True)
class RequestRecord:
    request_id: str
    tenant_id: str
    model: str
    started_ts: int
    completed_ts: int
    latency_ms: int
    prompt_tokens: int
    completion_tokens: int
    status: str
    price_multiplier_ppm: int = 1_000_000
    price_source: str = "list"


@dataclass(frozen=True)
class TenantCostMinute:
    tenant_id: str
    window_end: int
    micro_usd: int
    requests: int

    def key(self) -> tuple[str, int]:
        return self.tenant_id, self.window_end

    def to_dict(self) -> dict[str, int | str]:
        return {
            "tenant_id": self.tenant_id,
            "window_end": self.window_end,
            "micro_usd": self.micro_usd,
            "requests": self.requests,
        }


@dataclass(frozen=True)
class PartialCost:
    tenant_id: str
    salt: int
    window_end: int
    micro_usd: int
    requests: int


@dataclass
class RequestAssemblerModel:
    """A small deterministic oracle for the keyed RequestAssembler operator."""

    open_requests: dict[str, OpenRequest] = field(default_factory=dict)

    def apply(self, event: Event) -> RequestRecord | None:
        if event.event_type == REQUEST_STARTED:
            self.open_requests[event.request_id] = OpenRequest(
                event.request_id,
                event.tenant_id,
                event.model,
                event.ts,
                int(event.payload.get("prompt_tokens", 0)),
            )
            return None
        current = self.open_requests.get(event.request_id)
        if current is None:
            return None
        if event.event_type == TOKEN_GENERATED:
            current.completion_tokens += int(event.payload.get("n", 0))
            return None
        if event.event_type != REQUEST_COMPLETED:
            return None
        self.open_requests.pop(event.request_id, None)
        return RequestRecord(
            request_id=current.request_id,
            tenant_id=current.tenant_id,
            model=current.model,
            started_ts=current.started_ts,
            completed_ts=event.ts,
            latency_ms=int(event.payload.get("latency_ms", event.ts - current.started_ts)),
            prompt_tokens=current.prompt_tokens,
            completion_tokens=int(event.payload.get("completion_tokens", current.completion_tokens)),
            status=str(event.payload.get("status", "ok")),
        )


# From stage 6 the assembler sets an event-time timer at start + 5 minutes; a request still open
# when it fires is billed as abandoned, for the tokens seen so far.
ABANDON_AFTER_MS = 5 * MINUTE_MS


def abandoned_record(request: OpenRequest, fired_at: int) -> RequestRecord:
    """The record the assembler's timer emits for a request that never completed."""
    return RequestRecord(
        request_id=request.request_id,
        tenant_id=request.tenant_id,
        model=request.model,
        started_ts=request.started_ts,
        completed_ts=fired_at,
        latency_ms=fired_at - request.started_ts,
        prompt_tokens=request.prompt_tokens,
        completion_tokens=request.completion_tokens,
        status="abandoned",
    )


def assemble(events: Iterable[Event]) -> tuple[list[RequestRecord], dict[str, OpenRequest]]:
    model = RequestAssemblerModel()
    records: list[RequestRecord] = []
    for event in events:
        record = model.apply(event)
        if record is not None:
            records.append(record)
    return records, model.open_requests


def request_cost_micro_usd(
    record: RequestRecord,
    prices: Mapping[str, tuple[int, int]] = PRICE_PER_1K,
) -> int:
    """Return deterministic micro-dollars, rounded once at the request boundary."""
    input_per_1k, output_per_1k = prices[record.model]
    list_micro = (
        record.prompt_tokens * input_per_1k
        + record.completion_tokens * output_per_1k
        + 500
    ) // 1000
    return (list_micro * record.price_multiplier_ppm + 500_000) // 1_000_000


def window_end(timestamp_ms: int, size_ms: int = MINUTE_MS) -> int:
    return (timestamp_ms // size_ms + 1) * size_ms


def aggregate_costs(
    records: Iterable[RequestRecord],
    *,
    timestamps: Mapping[str, int] | None = None,
) -> list[TenantCostMinute]:
    """Aggregate by tenant-minute; timestamps can substitute processing/arrival time in Lab 4."""
    totals: dict[tuple[str, int], tuple[int, int]] = {}
    for record in records:
        ts = timestamps[record.request_id] if timestamps is not None else record.completed_ts
        key = record.tenant_id, window_end(ts)
        micro, count = totals.get(key, (0, 0))
        totals[key] = micro + request_cost_micro_usd(record), count + 1
    return [TenantCostMinute(t, end, micro, count)
            for (t, end), (micro, count) in sorted(totals.items())]


def salt_for(request_id: str, buckets: int = 8) -> int:
    return zlib.crc32(request_id.encode("utf-8")) % buckets


def partial_costs(records: Iterable[RequestRecord], buckets: int = 8) -> list[PartialCost]:
    totals: dict[tuple[str, int, int], tuple[int, int]] = {}
    for record in records:
        key = record.tenant_id, salt_for(record.request_id, buckets), window_end(record.completed_ts)
        micro, count = totals.get(key, (0, 0))
        totals[key] = micro + request_cost_micro_usd(record), count + 1
    return [PartialCost(t, salt, end, micro, count)
            for (t, salt, end), (micro, count) in sorted(totals.items())]


def merge_partials(updates: Sequence[PartialCost]) -> list[TenantCostMinute]:
    """Replace a salt's prior value; a late window re-fire is a correction, not another add."""
    latest: dict[tuple[str, int, int], PartialCost] = {}
    for update in updates:
        latest[update.tenant_id, update.window_end, update.salt] = update
    totals: dict[tuple[str, int], tuple[int, int]] = {}
    for update in latest.values():
        key = update.tenant_id, update.window_end
        micro, count = totals.get(key, (0, 0))
        totals[key] = micro + update.micro_usd, count + update.requests
    return [TenantCostMinute(t, end, micro, count)
            for (t, end), (micro, count) in sorted(totals.items())]


def upsert_minutes(emissions: Iterable[TenantCostMinute]) -> dict[tuple[str, int], TenantCostMinute]:
    """The billing-store view: recovery duplicates and late corrections share one natural key."""
    result: dict[tuple[str, int], TenantCostMinute] = {}
    for emission in emissions:
        result[emission.key()] = emission
    return result


def with_negotiated_price(record: RequestRecord, multiplier_ppm: int) -> RequestRecord:
    return replace(record, price_multiplier_ppm=multiplier_ppm, price_source="negotiated")


# The billing service's price book. The job's price lookup (stage 9 on) and the invoice oracle both
# read it, so they cannot disagree about what a tenant pays.
NEGOTIATED_PRICE_PPM = {"acme": 900_000}


def negotiated_price(record: RequestRecord) -> RequestRecord:
    """The price lookup's answer: the tenant's negotiated multiplier, or the list price."""
    multiplier = NEGOTIATED_PRICE_PPM.get(record.tenant_id)
    return record if multiplier is None else with_negotiated_price(record, multiplier)


def operator_watermark(partition_watermarks: Mapping[int, int], idle: Iterable[int] = ()) -> int:
    """The minimum active input watermark; useful as the deterministic oracle in Lab 5."""
    idle_set = set(idle)
    active = [value for partition, value in partition_watermarks.items() if partition not in idle_set]
    if not active:
        raise ValueError("an operator watermark needs at least one active input")
    return min(active)


def replay_from_checkpoint(
    records: Sequence[RequestRecord], checkpoint_index: int
) -> list[TenantCostMinute]:
    """Model restore-and-replay: the final materialized invoice must equal a failure-free run."""
    before = aggregate_costs(records[:checkpoint_index])
    replayed = aggregate_costs(records[checkpoint_index:])
    # In a real window operator the checkpoint also contains open accumulators.  Re-aggregating the
    # full logical input is the business oracle against which the cluster output is compared.
    del before, replayed
    return aggregate_costs(records)

"""The one PyFlink GatewayMeter job, parameterised by the course stage.

Each lab's ``gatewaymeter/job.py`` calls :func:`run_stage` with its module number.  Keeping the
operators here makes changes cumulative and keeps frozen UIDs spelled once.  The lab README tells
the learner which section to inspect and which environment switch deliberately breaks it.
"""
from __future__ import annotations

import asyncio
import json
import os

from common.events import REQUEST_COMPLETED, REQUEST_STARTED, TOKEN_GENERATED, parse_event
from common.gatewaymeter import (
    ABANDON_AFTER_MS,
    MINUTE_MS,
    OpenRequest,
    PartialCost,
    RequestRecord,
    TenantCostMinute,
    abandoned_record,
    negotiated_price,
    request_cost_micro_usd,
    salt_for,
)

BOOTSTRAP = "kafka:19092"
INPUT_TOPIC = "gateway-events"
INPUT_PARTITIONS = 8                 # gateway-events' partitions; from stage 4, one source subtask each
LATENESS_MS = 2 * MINUTE_MS


def _types():
    from pyflink.common import Types
    return Types.PICKLED_BYTE_ARRAY()


def _source(env, stage: int):
    from pyflink.common import Duration, WatermarkStrategy
    from pyflink.common.serialization import SimpleStringSchema
    from pyflink.datastream.connectors.kafka import KafkaOffsetsInitializer, KafkaSource
    from pyflink.common.watermark_strategy import TimestampAssigner

    source = (
        KafkaSource.builder()
        .set_bootstrap_servers(BOOTSTRAP)
        .set_topics(INPUT_TOPIC)
        .set_group_id(f"gatewaymeter-m{stage:02d}")
        .set_starting_offsets(KafkaOffsetsInitializer.earliest())
        .set_value_only_deserializer(SimpleStringSchema())
        .build()
    )
    if stage < 4:
        raw = env.from_source(source, WatermarkStrategy.no_watermarks(), "gateway-events")
        return raw.uid("gateway-events").map(parse_event, output_type=_types()).name("parse").uid("parse")

    class EventTime(TimestampAssigner):
        def extract_timestamp(self, value, record_timestamp):
            return value.ts

    # Event time is Event.ts. PyFlink's from_source hands the source only the Java half of a
    # WatermarkStrategy and drops a Python TimestampAssigner, which would leave event time at the
    # moment the producer sent each record. So the source carries no watermarks, and
    # assign_timestamps_and_watermarks, which does run a Python assigner, generates them right after
    # it. Watermarks are then per source subtask; with one subtask per partition that is per
    # partition, which is what a silent partition (module 5) needs. It takes all eight slots.
    strategy = WatermarkStrategy.for_bounded_out_of_orderness(Duration.of_seconds(20))
    if stage >= 5 and os.getenv("GM_DISABLE_IDLENESS") != "1":
        strategy = strategy.with_idleness(Duration.of_seconds(15))
    strategy = strategy.with_timestamp_assigner(EventTime())     # last: with_idleness drops it

    parallelism = env.get_parallelism()
    env.set_parallelism(int(os.getenv("GM_SOURCE_PARALLELISM", str(INPUT_PARTITIONS))))
    raw = env.from_source(source, WatermarkStrategy.no_watermarks(), "gateway-events")
    events = raw.uid("gateway-events").map(parse_event, output_type=_types()).name("parse").uid("parse")
    timed = events.assign_timestamps_and_watermarks(strategy).name("event-time").uid("event-time")
    env.set_parallelism(parallelism)
    return timed


def _records(events, stage: int):
    from pyflink.common import Time, Types
    from pyflink.datastream.functions import KeyedProcessFunction
    from pyflink.datastream.state import StateTtlConfig, ValueStateDescriptor

    class RequestAssembler(KeyedProcessFunction):
        def open(self, runtime_context):
            descriptor = ValueStateDescriptor("open", _types())
            if stage >= 3:
                ttl = (
                    StateTtlConfig.new_builder(Time.minutes(10))
                    .set_update_type(StateTtlConfig.UpdateType.OnCreateAndWrite)
                    .set_state_visibility(StateTtlConfig.StateVisibility.NeverReturnExpired)
                    .build()
                )
                descriptor.enable_time_to_live(ttl)
            self.open_request = runtime_context.get_state(descriptor)

        def process_element(self, event, ctx):
            current = self.open_request.value()
            if event.event_type == REQUEST_STARTED:
                current = OpenRequest(
                    event.request_id,
                    event.tenant_id,
                    event.model,
                    event.ts,
                    int(event.payload.get("prompt_tokens", 0)),
                )
                self.open_request.update(current)
                if stage >= 6:
                    ctx.timer_service().register_event_time_timer(event.ts + ABANDON_AFTER_MS)
                return
            if current is None:
                return
            if event.event_type == TOKEN_GENERATED:
                current.completion_tokens += int(event.payload.get("n", 0))
                self.open_request.update(current)
                return
            if event.event_type != REQUEST_COMPLETED:
                return
            self.open_request.clear()
            yield RequestRecord(
                current.request_id,
                current.tenant_id,
                current.model,
                current.started_ts,
                event.ts,
                int(event.payload.get("latency_ms", event.ts - current.started_ts)),
                current.prompt_tokens,
                int(event.payload.get("completion_tokens", current.completion_tokens)),
                str(event.payload.get("status", "ok")),
            )

        def on_timer(self, timestamp, ctx):
            current = self.open_request.value()
            if current is None:
                return
            self.open_request.clear()
            yield abandoned_record(current, timestamp)

    keyed = events.key_by(
        (lambda event: event.tenant_id)
        if os.getenv("GM_BREAK_KEY_BY") == "1"
        else (lambda event: event.request_id),
        key_type=Types.STRING(),
    )
    return keyed.process(RequestAssembler(), output_type=_types()).name("assemble").uid("assemble")


def _enrich(records, stage: int):
    if stage < 9:
        return records
    from pyflink.common import Time
    from pyflink.datastream.async_data_stream import AsyncDataStream
    from pyflink.datastream.functions import AsyncFunction

    class PriceLookup(AsyncFunction):
        def __init__(self, delay_seconds):
            self.delay_seconds = delay_seconds

        async def async_invoke(self, record):
            await asyncio.sleep(self.delay_seconds)
            # The same price book the invoice oracle uses: acme's negotiated 10% discount.
            return [negotiated_price(record)]

    return (
        AsyncDataStream.unordered_wait(
            records,
            PriceLookup(float(os.getenv("GM_PRICE_DELAY_SECONDS", "0.02"))),
            Time.seconds(5),
            capacity=int(os.getenv("GM_ASYNC_CAPACITY", "100")),
            output_type=_types(),
        )
        .name("enrich-price")
        .uid("enrich-price")
    )


def _cost_windows(records, stage: int):
    from pyflink.common import Time, Types
    from pyflink.datastream.functions import AggregateFunction, ProcessWindowFunction
    from pyflink.datastream.output_tag import OutputTag
    from pyflink.datastream.window import TumblingEventTimeWindows, TumblingProcessingTimeWindows

    class CostAggregate(AggregateFunction):
        def create_accumulator(self):
            return 0, 0

        def add(self, record, accumulator):
            return accumulator[0] + request_cost_micro_usd(record), accumulator[1] + 1

        def get_result(self, accumulator):
            return accumulator

        def merge(self, left, right):
            return left[0] + right[0], left[1] + right[1]

    class EmitCost(ProcessWindowFunction):
        def process(self, key, context, elements):
            micro, requests = next(iter(elements))
            yield TenantCostMinute(key, context.window().end, micro, requests)

    class EmitPartial(ProcessWindowFunction):
        def process(self, key, context, elements):
            tenant, salt = key
            micro, requests = next(iter(elements))
            yield PartialCost(tenant, salt, context.window().end, micro, requests)

    if stage >= 11:
        keyed = records.key_by(
            lambda record: (record.tenant_id, salt_for(record.request_id)),
            key_type=Types.TUPLE([Types.STRING(), Types.INT()]),
        )
        windowed = keyed.window(TumblingEventTimeWindows.of(Time.minutes(1)))
        if stage >= 6:
            windowed = windowed.allowed_lateness(LATENESS_MS)
        partials = (
            windowed.aggregate(CostAggregate(), EmitPartial(),
                               accumulator_type=Types.TUPLE([Types.LONG(), Types.LONG()]),
                               output_type=_types())
            .name("tenant-cost-partial")
            .uid("tenant-cost-partial")
        )
        return _merge_partial_stream(partials), None

    keyed = records.key_by(lambda record: record.tenant_id, key_type=Types.STRING())
    assigner = (
        TumblingProcessingTimeWindows.of(Time.minutes(1))
        if stage == 4 and os.getenv("GM_CLOCK", "event") == "processing"
        else TumblingEventTimeWindows.of(Time.minutes(1))
    )
    windowed = keyed.window(assigner)
    late_tag = OutputTag("tenant-cost-late", _types())
    if stage >= 6:
        windowed = windowed.allowed_lateness(LATENESS_MS).side_output_late_data(late_tag)
    result = (
        windowed.aggregate(CostAggregate(), EmitCost(),
                           accumulator_type=Types.TUPLE([Types.LONG(), Types.LONG()]),
                           output_type=_types())
        .name("tenant-cost")
        .uid("tenant-cost")
    )
    return result, result.get_side_output(late_tag) if stage >= 6 else None


def _merge_partial_stream(partials):
    from pyflink.common import Types
    from pyflink.datastream.functions import KeyedProcessFunction
    from pyflink.datastream.state import MapStateDescriptor

    class MergePartials(KeyedProcessFunction):
        def open(self, runtime_context):
            self.parts = runtime_context.get_map_state(
                MapStateDescriptor("partials", Types.STRING(), _types())
            )

        def _total(self, tenant_id, end):
            values = [value for value in self.parts.values() if value.window_end == end]
            return TenantCostMinute(
                tenant_id,
                end,
                sum(value.micro_usd for value in values),
                sum(value.requests for value in values),
            )

        def process_element(self, partial, ctx):
            self.parts.put(f"{partial.window_end}:{partial.salt}", partial)
            fire_at = partial.window_end - 1
            cleanup_at = partial.window_end + LATENESS_MS
            if ctx.timer_service().current_watermark() >= fire_at:
                yield self._total(partial.tenant_id, partial.window_end)
            else:
                ctx.timer_service().register_event_time_timer(fire_at)
            ctx.timer_service().register_event_time_timer(cleanup_at)

        def on_timer(self, timestamp, ctx):
            if timestamp % MINUTE_MS == MINUTE_MS - 1:
                yield self._total(ctx.get_current_key(), timestamp + 1)
                return
            end = timestamp - LATENESS_MS
            for key in list(self.parts.keys()):
                if key.startswith(f"{end}:"):
                    self.parts.remove(key)

    return (
        partials.key_by(lambda partial: partial.tenant_id, key_type=Types.STRING())
        .process(MergePartials(), output_type=_types())
        .name("tenant-cost")
        .uid("tenant-cost")
    )


def _configure_recovery(env, stage: int):
    if stage < 7:
        return
    from pyflink.common import Configuration
    from pyflink.datastream import CheckpointingMode, ExternalizedCheckpointRetention

    env.enable_checkpointing(10_000, CheckpointingMode.EXACTLY_ONCE)
    checkpoint = env.get_checkpoint_config()
    checkpoint.set_checkpoint_timeout(60_000)
    checkpoint.set_min_pause_between_checkpoints(2_000)
    checkpoint.set_max_concurrent_checkpoints(1)
    checkpoint.set_externalized_checkpoint_retention(
        ExternalizedCheckpointRetention.RETAIN_ON_CANCELLATION
    )
    if stage >= 9 and os.getenv("GM_UNALIGNED") == "1":
        checkpoint.enable_unaligned_checkpoints()
    config = Configuration()
    config.set_string("state.backend.type", os.getenv("GM_STATE_BACKEND", "rocksdb"))
    config.set_string("execution.checkpointing.storage", "filesystem")
    config.set_string("execution.checkpointing.dir", "file:///checkpoints")
    config.set_string("restart-strategy.type", "fixed-delay")
    config.set_string("restart-strategy.fixed-delay.attempts", os.getenv("GM_RESTART_ATTEMPTS", "3"))
    config.set_string("restart-strategy.fixed-delay.delay", "10 s")
    env.configure(config)


def _sink(costs, stage: int):
    from pyflink.common import Types

    if stage < 8:
        costs.map(lambda value: json.dumps(value.to_dict(), sort_keys=True),
                  output_type=Types.STRING()).print()
        return
    from pyflink.common.serialization import SimpleStringSchema
    from pyflink.datastream.connectors.base import DeliveryGuarantee
    from pyflink.datastream.connectors.kafka import KafkaRecordSerializationSchema, KafkaSink

    guarantee_name = os.getenv("GM_SINK_GUARANTEE", "exactly_once")
    guarantee = (
        DeliveryGuarantee.AT_LEAST_ONCE
        if guarantee_name == "at_least_once"
        else DeliveryGuarantee.EXACTLY_ONCE
    )
    serializer = (
        KafkaRecordSerializationSchema.builder()
        .set_topic("tenant-cost")
        .set_value_serialization_schema(SimpleStringSchema())
        .build()
    )
    builder = (
        KafkaSink.builder()
        .set_bootstrap_servers(BOOTSTRAP)
        .set_record_serializer(serializer)
        .set_delivery_guarantee(guarantee)
        .set_property("transaction.timeout.ms", "900000")
    )
    if guarantee == DeliveryGuarantee.EXACTLY_ONCE:
        builder = builder.set_transactional_id_prefix(
            os.getenv("GM_TRANSACTIONAL_ID_PREFIX", f"gatewaymeter-m{stage:02d}-tenant-cost")
        )
    strings = costs.map(lambda value: json.dumps(value.to_dict(), sort_keys=True),
                        output_type=Types.STRING())
    strings.sink_to(builder.build()).name("sink-tenant-cost").uid("sink-tenant-cost")


def build_stage(env, stage: int):
    from pyflink.common import Types

    _configure_recovery(env, stage)
    env.set_parallelism(int(os.getenv("GM_PARALLELISM", "4")))
    if stage >= 10:
        env.set_max_parallelism(int(os.getenv("GM_MAX_PARALLELISM", "128")))
    env.add_python_file("/labs/common")
    events = _source(env, stage)
    records = _records(events, stage)
    if stage <= 3:
        records.print()
        return records
    records = _enrich(records, stage)
    costs, late = _cost_windows(records, stage)
    if late is not None:
        late.map(lambda value: f"LATE {value.request_id}",
                 output_type=Types.STRING()).print()
    _sink(costs, stage)
    return costs


def run_stage(stage: int) -> None:
    from pyflink.datastream import StreamExecutionEnvironment

    env = StreamExecutionEnvironment.get_execution_environment()
    build_stage(env, stage)
    env.execute(f"GatewayMeter - module {stage}")

#!/usr/bin/env python3
"""Reset the local ``tenant-cost`` topic for one isolated billing experiment.

Exactly-once recovery applies within one submitted job. It does not deduplicate
records produced by several independent GatewayMeter jobs. This tool creates a
clean observation boundary while protecting a running lab from accidental data
loss: it refuses to proceed while the local Flink cluster reports an active job.

``--with-input`` also empties ``gateway-events``. Every stage reads that topic from
the earliest offset, so without it the capstone would re-read the traffic of every
earlier lab alongside the trace its oracle describes. Flink cannot see a producer
left running from an earlier lab, so the tool then watches the empty topic for a few
seconds and fails if anything is still writing to it.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path
import subprocess
import time
import urllib.request


ROOT = Path(__file__).resolve().parents[1]
FLINK_JOBS_URL = "http://127.0.0.1:8081/jobs/overview"
TOPIC = "tenant-cost"
TOPIC_PARTITIONS = {"tenant-cost": 4, "gateway-events": 8}
# tenant-cost keeps every record: the sink stamps each result with its event time, and a seeded
# trace's minutes are older than Kafka's 7-day retention, which would delete them within minutes.
TOPIC_CONFIG = {"tenant-cost": ["--config", "retention.ms=-1"]}
QUIET_SECONDS = 5


def active_job_ids(payload: dict) -> list[str]:
    terminal = {"FINISHED", "FAILED", "CANCELED", "SUSPENDED"}
    return [
        str(job["jid"])
        for job in payload.get("jobs", [])
        if str(job.get("state", "")).upper() not in terminal
    ]


def flink_active_jobs() -> list[str]:
    try:
        with urllib.request.urlopen(FLINK_JOBS_URL, timeout=3) as response:
            return active_job_ids(json.load(response))
    except Exception as exc:
        raise SystemExit(
            "Cannot read the local Flink job list at http://127.0.0.1:8081. "
            "Start the compose cluster before resetting the topic."
        ) from exc


def kafka_topics(*args: str) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        [
            "docker", "compose", "exec", "-T", "kafka",
            "/opt/kafka/bin/kafka-topics.sh",
            "--bootstrap-server", "localhost:19092",
            *args,
        ],
        cwd=ROOT,
        check=True,
        text=True,
        capture_output=True,
    )


def records_in(topic: str) -> int:
    """Sum of the end offsets over the topic's partitions."""
    out = subprocess.run(
        [
            "docker", "compose", "exec", "-T", "kafka",
            "/opt/kafka/bin/kafka-get-offsets.sh",
            "--bootstrap-server", "localhost:19092", "--topic", topic,
        ],
        cwd=ROOT, check=True, text=True, capture_output=True,
    ).stdout
    return sum(int(line.rsplit(":", 1)[1]) for line in out.splitlines() if line.startswith(f"{topic}:"))


def reset_topic(topic: str) -> None:
    kafka_topics("--delete", "--if-exists", "--topic", topic)
    for _ in range(30):
        listed = kafka_topics("--list").stdout.splitlines()
        if topic not in listed:
            break
        time.sleep(1)
    else:
        raise SystemExit(f"Timed out waiting for {topic} deletion")

    partitions = TOPIC_PARTITIONS[topic]
    kafka_topics(
        "--create", "--topic", topic, "--partitions", str(partitions),
        "--replication-factor", "1", *TOPIC_CONFIG.get(topic, []),
    )
    print(f"Reset {topic}: empty topic with {partitions} partitions")


def main(argv: list[str] | None = None) -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--with-input", action="store_true",
                        help="also empty gateway-events (module 12's clean trace)")
    args = parser.parse_args(argv)
    topics = [TOPIC, "gateway-events"] if args.with_input else [TOPIC]

    active = flink_active_jobs()
    if active:
        joined = ", ".join(active)
        raise SystemExit(
            f"Refusing to reset {' and '.join(topics)}: active Flink job(s): {joined}. "
            "Cancel them in the Web UI, then run this command again."
        )
    for topic in topics:
        reset_topic(topic)
    if args.with_input:
        time.sleep(QUIET_SECONDS)
        written = records_in("gateway-events")
        if written:
            raise SystemExit(
                f"gateway-events received {written} records in the {QUIET_SECONDS} s after the reset: "
                "a producer from an earlier lab is still running. Stop it, then reset again."
            )


if __name__ == "__main__":
    main()