scrypt/warden/author.py

"""The encounter author: the Warden re-composes the next fight — in bands.

The cardinal rule holds. Composition is deterministic: a small set of
mutations of the YAML-authored base script, each REJECTED at build time
unless it stays inside threat-budget bands (total cost and a cumulative
front-load cap). The LLM's entire authority is choosing WHICH pre-proven
variant will sting this player hardest, informed by its memory shards.
A bot smoke-sim guards the winner; any failure anywhere returns the base
script exactly as authored.
"""

from __future__ import annotations

import random
from dataclasses import dataclass

from scrypt.engine.bots import simulate
from scrypt.engine.cards import Card
from scrypt.engine.combat import LANES, EncounterScript, Result, ScriptedPlay

from .context import build_messages
from .harness import Harness, Tool

BASE_LABEL = "as scheduled"

# Threat bands: a variant may shave a fifth or add 15% — never more.
TOTAL_LOW, TOTAL_HIGH = 0.80, 1.15
FRONTLOAD_SLACK = 7  # cumulative threat may lead the base by at most this
GATE_DROP = 0.25     # bot winrate may fall at most this far below base
GATE_SEEDS = 24


def card_threat(card: Card) -> int:
    return card.power * 2 + card.health + 2 * len(card.sigils)


def script_cost(script: EncounterScript) -> int:
    return sum(card_threat(p.card) for turn in script for p in turn)


def _cumulative(script: EncounterScript) -> list[int]:
    total, out = 0, []
    for turn in script:
        total += sum(card_threat(p.card) for p in turn)
        out.append(total)
    return out


def within_bands(candidate: EncounterScript, base: EncounterScript) -> bool:
    if len(candidate) != len(base):
        return False
    if any(len(turn) > 2 for turn in candidate):
        return False
    if any(not (0 <= p.lane < LANES) for turn in candidate for p in turn):
        return False
    base_cost = script_cost(base)
    cost = script_cost(candidate)
    if not (base_cost * TOTAL_LOW <= cost <= base_cost * TOTAL_HIGH):
        return False
    for cand_cum, base_cum in zip(_cumulative(candidate), _cumulative(base)):
        if cand_cum > base_cum + FRONTLOAD_SLACK:
            return False
    return True


# ----------------------------------------------------------- mutations

def _copy(script: EncounterScript) -> list[list[ScriptedPlay]]:
    return [list(turn) for turn in script]


def _flat(script) -> list[tuple[int, int, ScriptedPlay]]:
    return [(t, i, p) for t, turn in enumerate(script) for i, p in enumerate(turn)]


def _free_lane(turn: list[ScriptedPlay], rng: random.Random) -> int | None:
    free = sorted(set(range(LANES)) - {p.lane for p in turn})
    return rng.choice(free) if free else None


def _swarm(script, content, rng):
    """Trade the heaviest scheduled play for cheap early pressure."""
    out = _copy(script)
    plays = _flat(out)
    if not plays:
        return None
    t, i, _ = max(plays, key=lambda x: card_threat(x[2].card))
    del out[t][i]
    reaper = content.card("reaper")
    placed = 0
    for turn_idx in range(1, len(out)):
        if placed == 2:
            break
        if len(out[turn_idx]) < 2:
            lane = _free_lane(out[turn_idx], rng)
            if lane is not None:
                out[turn_idx].append(ScriptedPlay(lane=lane, card=reaper))
                placed += 1
    return out if placed == 2 else None


def _dead_air(script, content, rng):
    """The ground forces stand down; everything comes in over the wall."""
    out = _copy(script)
    kill = content.card("kill-signal")
    grounded = [
        (t, i) for t, i, p in _flat(out)
        if not p.card.has("tunneling") and card_threat(p.card) >= card_threat(kill)
    ]
    if len(grounded) < 2:
        return None
    for t, i in rng.sample(grounded, 2):
        out[t][i] = ScriptedPlay(lane=out[t][i].lane, card=kill)
    return out


def _panes(script, content, rng):
    """Cheap bodies become multiplexers: adjacency stops being safe."""
    out = _copy(script)
    mux = content.card("multiplexer")
    cheap = sorted(
        (
            (t, i) for t, i, p in _flat(out)
            if p.card.id != "multiplexer" and card_threat(p.card) <= card_threat(mux)
        ),
        key=lambda x: card_threat(out[x[0]][x[1]].card),
    )
    if len(cheap) < 2:
        return None
    for t, i in cheap[:2]:
        out[t][i] = ScriptedPlay(lane=out[t][i].lane, card=mux)
    return out


def _wall(script, content, rng):
    """Two small plays are recalled; something large is scheduled instead."""
    out = _copy(script)
    plays = sorted(_flat(out), key=lambda x: card_threat(x[2].card))
    if len(plays) < 3:
        return None
    for t, i, _ in sorted(plays[:2], key=lambda x: (x[0], -x[1])):
        del out[t][i]
    mid = len(out) // 2
    for turn_idx in (mid, mid + 1, mid - 1):
        if 0 <= turn_idx < len(out) and len(out[turn_idx]) < 2:
            lane = _free_lane(out[turn_idx], rng)
            if lane is not None:
                out[turn_idx].append(
                    ScriptedPlay(lane=lane, card=content.card("cron-golem"))
                )
                return out
    return None


MUTATIONS = [
    ("the swarm", "many cheap processes, early — the board fills before they breathe", _swarm),
    ("dead air", "ground forces swapped for tunneling kill-signals; blockers stop mattering", _dead_air),
    ("panes within panes", "small bodies become multiplexers; standing next to anything hurts", _panes),
    ("the wall", "small plays recalled; a cron-golem is scheduled in their place", _wall),
]


@dataclass(frozen=True)
class Variant:
    label: str
    description: str
    script: EncounterScript


def variants(content, base: EncounterScript, rng: random.Random) -> list[Variant]:
    """Every pre-proven option, base first. Mutations that fail to build
    or fall outside the bands simply don't make the menu."""
    out = [Variant(BASE_LABEL, "the encounter exactly as scheduled", base)]
    for label, desc, fn in MUTATIONS:
        candidate = fn(base, content, rng)
        if candidate is not None and within_bands(candidate, base):
            out.append(Variant(label, desc, candidate))
    return out


# ----------------------------------------------------------- the choice

def choice_frame(encounter_name: str, shards: str, options: list[Variant]) -> str:
    menu = "\n".join(f"- {v.label}: {v.description}" for v in options)
    memory = shards.strip() or "- (no observations on file yet)"
    return (
        f"You are re-composing the encounter '{encounter_name}' for this "
        "specific player.\nWhat you remember of them:\n"
        f"{memory}\n\nPre-approved variants:\n{menu}\n\n"
        "Choose the variant that punishes their habits hardest by calling the tool."
    )


async def _choose(backend, encounter_name: str, shards: str,
                  options: list[Variant], rng: random.Random) -> str:
    labels = [v.label for v in options]
    if backend is None:
        return rng.choice(labels)
    picked: list[str] = []
    tool = Tool(
        name="compose",
        description="pick one pre-approved encounter variant",
        schema={"properties": {"variant": {"enum": labels}}, "required": ["variant"]},
        handler=lambda args: picked.append(args["variant"]) or "composed",
    )
    harness = Harness(backend, [tool], max_steps=2, max_tokens=120)
    await harness.run(build_messages(choice_frame(encounter_name, shards, options)))
    return picked[0] if picked else rng.choice(labels)


def smoke_gate(content, candidate: EncounterScript, base: EncounterScript,
               deck: list[Card], seeds: int = GATE_SEEDS) -> bool:
    """The floor bot plays the player's own deck against both scripts.
    A variant that craters the floor winrate is too hot to ship."""
    side = [content.card("bit")] * 20

    def rate(script) -> float:
        wins = sum(
            simulate(list(deck), side, script, seed).result is Result.PLAYER_WIN
            for seed in range(seeds)
        )
        return wins / seeds

    return rate(candidate) >= rate(base) - GATE_DROP


async def compose(
    backend,
    content,
    encounter_id: str,
    deck: list[Card],
    shards: str,
    rng: random.Random,
) -> tuple[EncounterScript, str | None]:
    """(script, variant_label-or-None). None means: exactly as authored."""
    base = content.encounters[encounter_id]["script"]
    options = variants(content, base, rng)
    if len(options) <= 1:
        return base, None
    try:
        label = await _choose(
            backend, content.encounters[encounter_id]["name"], shards, options, rng
        )
    except Exception:
        return base, None
    if label == BASE_LABEL:
        return base, None
    chosen = next(v for v in options if v.label == label)
    if not smoke_gate(content, chosen.script, base, deck):
        return base, None
    return chosen.script, chosen.label