generator.py

from dataclasses import dataclass
import json
import random
import subprocess
import tempfile
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any, Protocol, Sequence

from budget import Card, CardSpec, EffectPlan, cost_card
from local_llm import ChatCompleter
from primitives import PrimitiveId, School, school_bias, school_primitives


class CardModelClient(Protocol):
    # Return one raw model-authored card payload.
    def create_card(self, payload: dict[str, Any]) -> dict[str, Any]:  # pragma: no cover
        ...


class CardPackClient(Protocol):
    # Return one raw model-authored pack payload.
    def create_pack(self, payload: dict[str, Any]) -> dict[str, Any]:  # pragma: no cover
        ...


@dataclass(frozen=True)
class DeckCardContext:
    name: str
    cost: int
    rules_text: str


@dataclass(frozen=True)
class CodexCardClient:
    command: tuple[str, ...] = ("codex", "exec", "--ephemeral", "--skip-git-repo-check", "-")
    timeout_seconds: int = 120
    cwd: str = "."

    # Return one raw card payload from a Codex invocation.
    def create_card(self, payload: dict[str, Any]) -> dict[str, Any]:
        pack_payload = dict(payload)
        pack_payload.setdefault("cost", 1)
        pack_payload["pack_size"] = 1
        return self.create_pack(pack_payload)["cards"][0]

    # Return one raw card pack payload from a Codex invocation.
    def create_pack(self, payload: dict[str, Any]) -> dict[str, Any]:
        with tempfile.TemporaryDirectory() as tmpdir:
            output_path = Path(tmpdir) / "codex-card-pack.json"
            command = self.command[:-1] + ("--output-last-message", str(output_path), self.command[-1])
            result = run_codex_command(command, codex_pack_prompt(payload), self.cwd, self.timeout_seconds)
            if result.returncode != 0:
                raise RuntimeError(result.stderr.strip() or "Codex card generation failed")
            return json.loads(extract_json_object(output_path.read_text()))


@dataclass(frozen=True)
class MiniCPMCardClient:
    chat: ChatCompleter

    # Return one raw card payload from a local MiniCPM endpoint.
    def create_card(self, payload: dict[str, Any]) -> dict[str, Any]:
        pack_payload = dict(payload)
        pack_payload.setdefault("cost", 1)
        pack_payload["pack_size"] = 1
        return self.create_pack(pack_payload)["cards"][0]

    # Return one raw card pack payload from a local MiniCPM endpoint.
    def create_pack(self, payload: dict[str, Any]) -> dict[str, Any]:
        text = self.chat.complete(card_system_prompt(), codex_pack_prompt(payload))
        return json.loads(extract_json_object(text))


@dataclass(frozen=True)
class LlamaCppCardClient:
    chat: ChatCompleter

    # Return one raw card payload from a llama.cpp MiniCPM endpoint.
    def create_card(self, payload: dict[str, Any]) -> dict[str, Any]:
        return generate_llamacpp_card(self.chat, payload, ())

    # Return one raw card pack, generating the cards concurrently across
    # llama-server slots, then de-duplicating names so no two picks collide.
    # A single card's failure falls back only that card, never the whole pack.
    def create_pack(self, payload: dict[str, Any]) -> dict[str, Any]:
        size = int(payload.get("pack_size", 3))
        focuses = pack_focuses(payload, size)
        with ThreadPoolExecutor(max_workers=size) as pool:
            cards = list(pool.map(lambda index: self.one_card(payload, focuses[index]), range(size)))
        return {"cards": dedup_pack_names(cards, str(payload.get("school", "fire")))}

    # Generate one card centered on its slot's primitive, degrading to a single-card
    # fallback on any model failure so one slow slot never collapses the pack.
    def one_card(self, payload: dict[str, Any], focus: str = "") -> dict[str, Any]:
        try:
            return generate_llamacpp_card(self.chat, payload, (), focus=focus)
        except Exception:
            need = f"explore primitive_id {focus}" if focus else llama_candidate_need(payload, {})
            return repair_llamacpp_card(fallback_raw_for_need(need, payload), payload, [])


TOOL_SCHEMA: dict[str, Any] = {
    "name": "create_card",
    "parameters": {
        "type": "object",
        "required": ["name", "flavor", "art_prompt", "effects"],
        "properties": {
            "name": {"type": "string"},
            "flavor": {"type": "string"},
            "art_prompt": {"type": "string"},
            "effects": {
                "type": "array",
                "items": {
                    "type": "object",
                    "required": ["primitive_id"],
                    "properties": {
                        "primitive_id": {"type": "string"},
                        "weight": {"type": "integer", "minimum": 1},
                    },
                },
            },
        },
    },
}

PACK_SCHEMA: dict[str, Any] = {
    "type": "object",
    "required": ["cards"],
    "properties": {
        "cards": {
            "type": "array",
            "minItems": 3,
            "maxItems": 3,
            "items": TOOL_SCHEMA["parameters"],
        }
    },
}


# Generate one engine-costed card through a mockable model boundary.
def generate_card(
    client: CardModelClient,
    school: School,
    theme: str,
    current_deck: Sequence[Card],
    cost: int,
    draft_anchors: Sequence[Card] = (),
) -> Card:
    raw = client.create_card(generator_payload(school, theme, current_deck, draft_anchors))
    spec = parse_card_payload(raw, school, theme, cost)
    return cost_card(spec)


# Generate one engine-costed draft pack through a mockable model boundary.
def generate_pack(
    client: CardPackClient,
    school: School,
    theme: str,
    current_deck: Sequence[Card],
    cost: int,
    pack_size: int = 3,
    draft_anchors: Sequence[Card] = (),
    quick: bool = False,
) -> tuple[Card, ...]:
    payload = pack_payload(school, theme, current_deck, cost, pack_size, draft_anchors)
    payload["quick"] = quick
    raw = client.create_pack(payload)
    pack = parse_pack_payload(raw, school, theme, cost)
    if len(pack) < pack_size:
        raise ValueError("generated pack had the wrong size")
    return pack[:pack_size]


# Generate several engine-costed cards through one model call.
def generate_cards(
    client: CardPackClient,
    school: School,
    theme: str,
    current_deck: Sequence[Card],
    costs: Sequence[int],
    draft_anchors: Sequence[Card] = (),
) -> tuple[Card, ...]:
    raw = client.create_pack(cards_payload(school, theme, current_deck, costs, draft_anchors))
    if len(raw["cards"]) != len(costs):
        raise ValueError("generated card batch had the wrong size")
    return tuple(cost_card(parse_card_payload(card, school, theme, cost)) for card, cost in zip(raw["cards"], costs))


# Build the deck-aware model payload for card generation.
def generator_payload(school: School, theme: str, current_deck: Sequence[Card], draft_anchors: Sequence[Card] = ()) -> dict[str, Any]:
    return {
        "school": school,
        "theme": theme,
        "school_identity": school_identity(school),
        "school_bias": school_bias(school),
        "primitive_guidance": primitive_guidance(school),
        "allowed_primitives": school_primitives(school),
        "current_deck": [card_context(card) for card in current_deck],
        "current_deck_summary": deck_summary(current_deck),
        "draft_anchors": [card_context(card) for card in draft_anchors],
        "draft_anchor_summary": deck_summary(draft_anchors),
        "draft_direction": draft_direction(school, draft_anchors),
        "draft_need": draft_need(school, current_deck),
        "tool_schema": TOOL_SCHEMA,
        "instruction": "Choose effect shapes only; numeric magnitudes are assigned by the engine.",
    }


# Build the deck-aware model payload for draft pack generation.
def pack_payload(
    school: School,
    theme: str,
    current_deck: Sequence[Card],
    cost: int,
    pack_size: int,
    draft_anchors: Sequence[Card] = (),
) -> dict[str, Any]:
    payload = generator_payload(school, theme, current_deck, draft_anchors)
    payload["cost"] = cost
    payload["pack_size"] = pack_size
    payload["pack_schema"] = PACK_SCHEMA
    return payload


# Build the deck-aware model payload for multi-card generation.
def cards_payload(
    school: School,
    theme: str,
    current_deck: Sequence[Card],
    costs: Sequence[int],
    draft_anchors: Sequence[Card] = (),
) -> dict[str, Any]:
    payload = generator_payload(school, theme, current_deck, draft_anchors)
    payload["costs"] = tuple(costs)
    payload["pack_size"] = len(costs)
    payload["pack_schema"] = PACK_SCHEMA
    return payload


# Build the prompt sent to Codex for one draft pack.
def codex_pack_prompt(payload: dict[str, Any]) -> str:
    cost_line = cost_instruction(payload)
    return "\n".join(
        (
            "Generate a Tabras draft pack as strict JSON only.",
            "Return exactly this shape: {\"cards\": [{\"name\": str, \"flavor\": str, \"art_prompt\": str, \"effects\": [{\"primitive_id\": str, \"weight\": int}]}]}",
            "Do not include markdown, comments, or numeric magnitudes like damage/block/heal values.",
            "The engine owns all final numbers. You only choose primitive ids and weights.",
            "Make cards in the same response meaningfully different: vary primitive mixes, names, and tactical roles.",
            "The flavor and art_prompt must describe the same concrete scene.",
            "The art_prompt is for image generation: describe subject, setting, mood, and magic; never mention cards, frames, text, UI, or rules.",
            f"School: {payload['school']}",
            f"Theme: {payload['theme']}",
            f"Class identity: {payload['school_identity']}",
            f"School bias: {payload['school_bias']}",
            f"Primitive guidance: {payload['primitive_guidance']}",
            f"Deck primitive counts: {json.dumps(payload['current_deck_summary'])}",
            f"Draft anchor cards: {json.dumps(payload['draft_anchors'])}",
            f"Draft anchor primitive counts: {json.dumps(payload['draft_anchor_summary'])}",
            f"Draft direction: {payload['draft_direction']}",
            f"Draft need: {payload['draft_need']}",
            cost_line,
            f"Pack size: {payload['pack_size']}",
            f"Return exactly {payload['pack_size']} cards: no fewer and no extras.",
            f"Allowed primitives: {json.dumps(payload['allowed_primitives'])}",
            f"Current deck context: {json.dumps(payload['current_deck'])}",
        )
    )


# Return the system prompt for card generation.
def card_system_prompt() -> str:
    return "You are Tabras card authoring logic. Output strict JSON only. No thinking. No markdown."


NAME_BAN = (
    "deal",
    "damage",
    "bomb",
    "burn",
    "block",
    "draw",
    "scaling",
    "multi",
    "hit",
    "initiative",
    "vulnerable",
    "weak",
    "ward",
    "energy",
    "pressure",
    "clock",
    "plan",
    "tempo",
    "payoff",
    "fast",
    "immediate",
    "card",
    "school",
    "anime",
    "fantasy",
    "spell",
    "primitive",
)
NAME_STEM_BAN = ("burn", "fast", "flam", "pressur", "primitiv")
# Deep per-school name pools so a duplicate is replaced with a fresh evocative
# name (drawn in order against the run's used-name set) instead of a number.
FALLBACK_NAMES: dict[str, tuple[str, ...]] = {
    "fire": (
        "Cinder Warrant", "Ashen Oath", "Ember Brand", "Sable Pyre", "Charcoal Saint", "Red Sigil",
        "Smolder Pact", "Soot Crown", "Furnace Hymn", "Molten Verdict", "Scorch Litany", "Kindled Wrath",
        "Pyre Mantle", "Coalheart", "Searing Edict", "Brimstone Vow", "Cinderfall", "Ashen Reverie",
        "Ember Requiem", "Slag Reliquary", "Firstflame", "Cauldron Mark", "Emberwake", "Hearthscar",
    ),
    "ice": (
        "Glass Edict", "Rime Mirror", "Pale Aurora", "Silver Stillness", "Needle Choir", "Blue Omen",
        "Frostbound Vow", "Hoarfrost Litany", "Winter Verdict", "Crystal Reverie", "Sleet Sigil", "Glacial Hymn",
        "Pale Reliquary", "Shivermark", "Frozen Requiem", "Snowfall Oath", "Cold Aurora", "Brittle Crown",
        "Wintermute", "Frostwake", "Gelid Pact", "Mirrorfrost", "Still Omen", "Icebound Mantle",
    ),
    "earth": (
        "Stone Covenant", "Cairn Writ", "Rootbound Seal", "Granite Vow", "Mossbound Relic", "Deep Warden",
        "Loam Litany", "Bedrock Hymn", "Quarry Verdict", "Tor Reliquary", "Sediment Sigil", "Boulder Oath",
        "Furrow Mark", "Earthen Requiem", "Gravel Ledger", "Standing Stone", "Mirebound", "Cragheart",
        "Hollow Crown", "Slatewake", "Claybound Pact", "Old Cairn", "Verdant Writ", "Stonebound Vow",
    ),
}
SCHOOL_FORBIDDEN_IMAGERY: dict[str, tuple[str, ...]] = {
    "fire": ("glacier", "frost", "frozen", "ice", "snow", "winter", "blizzard", "earth"),
    "ice": ("cinder", "ember", "flame", "inferno", "pyre", "ashen", "volcano", "fire", "earth"),
    "earth": ("glacier", "frost", "inferno", "flame", "ember", "blizzard", "fire", "ice"),
}


# Return whether a card name is empty or leans on mechanic words.
def generic_card_name(name: str) -> bool:
    words = [word.strip(".,:;!?'\"").lower() for word in name.split()]
    return not words or any(word in NAME_BAN or word.isdigit() or any(word.startswith(stem) for stem in NAME_STEM_BAN) for word in words)


# Return the shared naming and flavor rules for card prompts.
def naming_rules() -> tuple[str, ...]:
    return (
        "Name the card with one to three evocative English words of concrete imagery from the world's fantasy.",
        f"Never use these words in the name: {', '.join(NAME_BAN)}.",
        "Bad names repeat mechanics, like Fire Deal or Pressure Clock.",
        "Write flavor as one short English in-world sentence with no mechanic or plan words.",
    )


# Return school-specific imagery constraints for prompts and validation.
def school_imagery_rule(school: str) -> str:
    forbidden = ", ".join(SCHOOL_FORBIDDEN_IMAGERY.get(school, ()))
    if not forbidden:
        return "Keep name, flavor, and art imagery aligned with the school."
    return f"Do not use off-school imagery for {school}; avoid these words in name, flavor, and art_prompt: {forbidden}."


# Build a compact one-card prompt for llama.cpp MiniCPM.
def llamacpp_card_prompt(payload: dict[str, Any], pack_cards: Sequence[dict[str, Any]], focus: str = "") -> str:
    used_names = [card.get("name", "") for card in pack_cards]
    pack_counts = pack_summary(pack_cards)
    candidate_need = llama_candidate_need(payload, pack_counts)
    focus_line = (
        (f"Center this card on primitive_id {focus}, but invent its fantasy and tactical purpose freely.",)
        if focus
        else ()
    )
    return "\n".join(
        (
            "Create exactly one Tabras card as JSON.",
            # effects + name come first so the card survives if the model is cut off
            # mid-output; keep flavor and art_prompt short so the JSON always closes.
            "Shape: {\"effects\": [{\"primitive_id\": string, \"weight\": 1}], \"name\": string, \"flavor\": string, \"art_prompt\": string}",
            "Use one or two effects. Do not write numeric damage, block, heal, or rules numbers.",
            "Keep flavor under 12 words. Keep art_prompt under 18 words. Output only one compact JSON object.",
            "Fit the school identity, but invent the card's fantasy and tactical purpose freely.",
            "Do not describe Fire as a people or species; use fast damage plan or pressure plan.",
            "Prefer a new tactical purpose over repeating the current deck or pack.",
            *focus_line,
            "Make art_prompt match the flavor sentence as a concrete image scene.",
            "The art_prompt must describe subject, setting, mood, and magic; never mention cards, frames, text, UI, or rules.",
            *naming_rules(),
            school_imagery_rule(str(payload["school"])),
            f"School: {payload['school']}",
            f"Theme: {payload['theme']}",
            f"Class identity: {payload['school_identity']}",
            f"Bias: {payload['school_bias']}",
            f"Primitive guidance: {payload['primitive_guidance']}",
            f"Deck primitive counts: {json.dumps(payload['current_deck_summary'])}",
            f"Draft anchor cards: {json.dumps(payload['draft_anchors'])}",
            f"Draft anchor primitive counts: {json.dumps(payload['draft_anchor_summary'])}",
            f"Draft direction: {payload['draft_direction']}",
            f"Current draft need: {payload['draft_need']}",
            f"This candidate should: {candidate_need}",
            f"Allowed primitive_id values: {json.dumps(payload['allowed_primitives'])}",
            f"Existing deck: {json.dumps(payload['current_deck'])}",
            f"Already in this pack: {json.dumps(used_names)}",
            f"Pack primitive counts so far: {json.dumps(pack_counts)}",
        )
    )


# Generate one llama.cpp card with one corrective retry; quick mode takes the first card.
def generate_llamacpp_card(chat: ChatCompleter, payload: dict[str, Any], pack_cards: Sequence[dict[str, Any]], focus: str = "") -> dict[str, Any]:
    prompt = llamacpp_card_prompt(payload, pack_cards, focus)
    card = parse_llamacpp_card_text(chat.complete(card_system_prompt(), prompt), payload, pack_cards, discard_bad_text=False)
    need = llama_candidate_need(payload, pack_summary(pack_cards))
    if payload.get("quick"):
        return repair_llamacpp_card(card, payload, pack_cards)
    if card_acceptable(card, need, payload):
        return card
    retry = parse_llamacpp_card_text(chat.complete(card_system_prompt(), llamacpp_retry_prompt(payload, need)), payload, pack_cards, discard_bad_text=False)
    if card_acceptable(retry, need, payload):
        return retry
    # Fall back to this slot's focus primitive (not always deal) so a pack whose
    # names get rejected still shows varied effects instead of three identical cards.
    fallback_need = f"explore primitive_id {focus}" if focus else need
    return repair_llamacpp_card(fallback_raw_for_need(fallback_need, payload), payload, pack_cards)


# Return whether a generated card meets its need with a non-generic name.
def card_acceptable(card: dict[str, Any], need: str, payload: dict[str, Any]) -> bool:
    return (
        card_satisfies_candidate_need(card, need)
        and not generic_card_name(str(card.get("name", "")))
        and not theme_name_leak(str(card.get("name", "")), str(payload["theme"]))
        and not school_flavor_mismatch(str(payload["school"]), card)
    )


# Build a corrective prompt for a missed candidate need.
def llamacpp_retry_prompt(payload: dict[str, Any], need: str) -> str:
    return "\n".join(
        (
            "Create exactly one corrected Tabras card as JSON.",
            "Shape: {\"effects\": [{\"primitive_id\": string, \"weight\": 1}], \"name\": string, \"flavor\": string, \"art_prompt\": string}",
            "Keep flavor under 12 words and art_prompt under 18 words; output one compact JSON object.",
            "Make flavor and art_prompt describe the same concrete scene.",
            "The art_prompt must describe an image scene only; no cards, text, frames, UI, or rules.",
            *naming_rules(),
            school_imagery_rule(str(payload["school"])),
            f"School: {payload['school']}",
            f"Theme: {payload['theme']}",
            f"Allowed primitive_id values: {json.dumps(payload['allowed_primitives'])}",
            f"Correction requirement: {need}",
            "If the requirement mentions primitive_id block, include an effect with primitive_id exactly \"block\".",
            "If the requirement mentions primitive_id scaling, include an effect with primitive_id exactly \"scaling\".",
            "If the requirement mentions primitive_id multi_hit, include an effect with primitive_id exactly \"multi_hit\".",
            "If the requirement mentions primitive_id initiative, include an effect with primitive_id exactly \"initiative\".",
            "If the requirement mentions primitive_id burn, include an effect with primitive_id exactly \"burn\".",
            "If the requirement mentions primitive_id bomb, include an effect with primitive_id exactly \"bomb\".",
            "If the requirement mentions primitive_id draw, include an effect with primitive_id exactly \"draw\".",
            "JSON only.",
        )
    )


# Return the cost instruction for a Codex generation payload.
def cost_instruction(payload: dict[str, Any]) -> str:
    if "costs" in payload:
        return f"Card costs by index: {json.dumps(payload['costs'])}"
    return f"Card cost for every candidate: {payload['cost']}"


# Run Codex for one card-generation prompt.
def run_codex_command(command: tuple[str, ...], prompt: str, cwd: str, timeout_seconds: int) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        command,
        input=prompt,
        text=True,
        cwd=cwd,
        check=False,
        capture_output=True,
        timeout=timeout_seconds,
    )


# Convert a costed card into compact generator context.
def card_context(card: Card) -> dict[str, Any]:
    return DeckCardContext(card.name, card.cost, card.rules_text()).__dict__


# Return the design identity for one school.
def school_identity(school: School) -> str:
    identities = {
        "fire": "Fire races with direct pressure, burn clocks, delayed bombs, and occasional scaling finishers.",
        "ice": "Ice wins through tempo, forced sequencing, vulnerability windows, and precise multi-hit pressure.",
        "earth": "Earth absorbs pressure with block, banks blocked damage as shield charge, then uses scaling to convert charge into burst; plain deal is secondary.",
    }
    return identities[school]


# Return school-specific primitive design guidance.
def primitive_guidance(school: School) -> str:
    guidance = {
        "fire": "Use deal for immediate pressure, burn for damage over turns, bomb for delayed burst, and scaling for finishers.",
        "ice": "Use initiative for turn order, vulnerable before multi_hit or deal, and conditional for precise finishers.",
        "earth": "Use block to create shield charge, scaling to spend shield charge as burst, ward for rare protection, and weak/draw as support.",
    }
    return guidance[school]


# Return a concise direction inferred from anchor cards.
def draft_direction(school: School, draft_anchors: Sequence[Card]) -> str:
    counts = deck_summary(draft_anchors)
    if not counts:
        return "No anchor cards yet; generate a strong candidate that could define the deck's direction."
    if school == "fire":
        return fire_anchor_direction(counts)
    if school == "ice":
        return ice_anchor_direction(counts)
    return earth_anchor_direction(counts)


# Return Fire direction from anchor primitives.
def fire_anchor_direction(counts: dict[str, int]) -> str:
    if counts.get("bomb", 0) or counts.get("burn", 0):
        return "Lean into a pressure clock: pair delayed damage with immediate deal, draw, or scaling finishers."
    return "Lean into fast pressure: add burn, bomb, draw, or finishers that support the anchor cards."


# Return Ice direction from anchor primitives.
def ice_anchor_direction(counts: dict[str, int]) -> str:
    if counts.get("vulnerable", 0) and not counts.get("multi_hit", 0):
        return "Exploit vulnerability windows with multi_hit, initiative, and precise finishers."
    if counts.get("multi_hit", 0) and not counts.get("vulnerable", 0):
        return "Support multi-hit pressure with vulnerable, initiative, draw, or conditional finishers."
    return "Build a tempo chain around initiative, vulnerable, multi_hit, and conditional payoff."


# Return Earth direction from anchor primitives.
def earth_anchor_direction(counts: dict[str, int]) -> str:
    if counts.get("block", 0) and not counts.get("scaling", 0):
        return "Use block as the shield-charge engine and add scaling payoff, weak, draw, or conversion."
    if counts.get("scaling", 0):
        return "Support shield-charge burst with block, survival, draw, weak, and occasional ward."
    return "Build the absorb-then-retaliate loop with block and shield-charge payoff."


# Return a per-candidate need for one-card llama.cpp generation.
def llama_candidate_need(payload: dict[str, Any], pack_counts: dict[str, int]) -> str:
    draft_need_text = str(payload["draft_need"]).lower()
    school = str(payload["school"])
    if school == "fire" and pack_counts.get("deal", 0) > 0 and pack_counts.get("burn", 0) == 0:
        return "explore primitive_id burn so this pack is not another direct damage card."
    if school == "fire" and pack_counts.get("deal", 0) > 0 and pack_counts.get("bomb", 0) == 0:
        return "explore primitive_id bomb so this pack has delayed pressure."
    if school == "fire" and pack_counts.get("deal", 0) > 0 and pack_counts.get("draw", 0) == 0:
        return "explore primitive_id draw as support instead of another direct damage card."
    if "add block" in draft_need_text and pack_counts.get("block", 0) == 0:
        return "explore primitive_id block as the shield-charge engine while inventing the card fantasy freely."
    if "scaling" in draft_need_text and pack_counts.get("scaling", 0) == 0:
        return "explore primitive_id scaling as the shield-charge payoff while inventing the card fantasy freely."
    if "multi_hit" in draft_need_text and pack_counts.get("multi_hit", 0) == 0:
        return "explore primitive_id multi_hit as the payoff for vulnerability while inventing the card fantasy freely."
    if "initiative" in draft_need_text and pack_counts.get("initiative", 0) == 0:
        return "explore primitive_id initiative to control turn order while inventing the card fantasy freely."
    if "plain damage is covered" in draft_need_text and pack_counts.get("deal", 0) > 0:
        return "avoid another primary deal card; explore block, scaling, weak, draw, or support."
    return "fit the draft need with a distinct name, fantasy, and tactical purpose."


# Return whether a card satisfies a per-candidate need.
def card_satisfies_candidate_need(card: dict[str, Any], need: str) -> bool:
    effects = card.get("effects", [])
    primitives = [effect.get("primitive_id") for effect in effects if isinstance(effect, dict)] if isinstance(effects, list) else []
    if "primitive_id scaling" in need:
        return "scaling" in primitives
    if "primitive_id block" in need:
        return "block" in primitives
    if "primitive_id multi_hit" in need:
        return "multi_hit" in primitives
    if "primitive_id initiative" in need:
        return "initiative" in primitives
    if "primitive_id burn" in need:
        return "burn" in primitives
    if "primitive_id bomb" in need:
        return "bomb" in primitives
    if "primitive_id draw" in need:
        return "draw" in primitives
    if "avoid another primary deal" in need:
        return not primitives or primitives[0] != "deal"
    return True


# Build a minimal raw response that satisfies a missed primitive need.
def fallback_raw_for_need(need: str, payload: dict[str, Any]) -> dict[str, Any]:
    allowed = tuple(payload["allowed_primitives"])
    primitive = next((candidate for candidate in allowed if f"primitive_id {candidate}" in need), allowed[0])
    return {"effects": [{"primitive_id": primitive, "weight": 1}]}


# Count primitive ids represented in a costed deck.
def deck_summary(deck: Sequence[Card]) -> dict[str, int]:
    counts: Counter[str] = Counter()
    for card in deck:
        counts.update(effect.primitive_id for effect in card.effects)
    return ordered_counts(counts)


# Count primitive ids represented in raw generated pack cards.
def pack_summary(cards: Sequence[dict[str, Any]]) -> dict[str, int]:
    counts: Counter[str] = Counter()
    for card in cards:
        effects = card.get("effects", [])
        if isinstance(effects, list):
            counts.update(effect.get("primitive_id") for effect in effects if isinstance(effect, dict))
    return ordered_counts(counts)


# Return a compact deck-aware drafting need.
def draft_need(school: School, deck: Sequence[Card]) -> str:
    counts = deck_summary(deck)
    if school == "earth":
        return earth_draft_need(counts)
    if school == "fire":
        return fire_draft_need(counts)
    return ice_draft_need(counts)


# Return an Earth-specific drafting need.
def earth_draft_need(counts: dict[str, int]) -> str:
    protection = counts.get("block", 0) + counts.get("ward", 0)
    payoff = counts.get("scaling", 0) + counts.get("deal", 0) + counts.get("conditional", 0)
    if counts.get("scaling", 0) > 0 and counts.get("block", 0) <= counts.get("scaling", 0):
        return "Scaling payoff exists; add block to bank shield charge before adding more plain deal."
    if counts.get("block", 0) > 0 and counts.get("scaling", 0) == 0:
        return "Add a shield-charge payoff using scaling; avoid making plain deal the main Earth payoff."
    if counts.get("deal", 0) >= protection + 2:
        return "Plain damage is covered; prefer block, scaling shield-charge payoff, weak, draw, or conversion."
    if protection >= payoff + 2:
        return "Protection is covered; prefer shield-charge payoff, pressure, weak, draw, or conversion instead of more ward."
    if counts.get("block", 0) == 0:
        return "Add block so Earth can bank shield charge before spending it."
    return "Improve the absorb-then-retaliate loop with a distinct support or payoff card."


# Return a Fire-specific drafting need.
def fire_draft_need(counts: dict[str, int]) -> str:
    delayed = counts.get("burn", 0) + counts.get("bomb", 0)
    pressure = counts.get("deal", 0) + counts.get("scaling", 0)
    if delayed >= pressure + 2:
        return "Delayed damage is covered; prefer immediate deal, scaling finishers, draw, or energy."
    if counts.get("deal", 0) == 0:
        return "Add immediate deal so Fire can race before delayed damage resolves."
    return "Add a distinct pressure card that advances Fire's fast damage plan."


# Return an Ice-specific drafting need.
def ice_draft_need(counts: dict[str, int]) -> str:
    setup = counts.get("initiative", 0) + counts.get("vulnerable", 0)
    payoff = counts.get("multi_hit", 0) + counts.get("conditional", 0) + counts.get("deal", 0)
    if counts.get("vulnerable", 0) > 0 and counts.get("multi_hit", 0) == 0:
        return "Vulnerability setup exists; add a multi_hit payoff instead of more plain deal."
    if counts.get("initiative", 0) == 0 and counts.get("vulnerable", 0) >= 2:
        return "Add initiative to control turn order around the vulnerability window."
    if counts.get("deal", 0) >= setup + 3:
        return "Plain damage is covered; prefer initiative, vulnerable, multi_hit, conditional, draw, or tempo."
    if setup >= payoff + 2:
        return "Tempo setup is covered; prefer payoff damage or multi-hit pressure."
    return "Add a tempo card that creates or exploits a vulnerability window."


# Return counts with empty keys removed.
def ordered_counts(counts: Counter[str]) -> dict[str, int]:
    return {key: counts[key] for key in sorted(counts) if key}


# Parse model payload into a card spec without trusting model-authored numbers.
def parse_card_payload(raw: dict[str, Any], school: School, theme: str, cost: int) -> CardSpec:
    effects = tuple(parse_effect_plan(effect) for effect in raw["effects"])
    return CardSpec(
        name=str(raw["name"]),
        cost=cost,
        school=school,
        theme=theme,
        effect_plans=effects,
        flavor=str(raw.get("flavor", "")),
        art_prompt=str(raw.get("art_prompt", "")),
    )


# Parse model payload into costed draft cards.
def parse_pack_payload(raw: dict[str, Any], school: School, theme: str, cost: int) -> tuple[Card, ...]:
    return tuple(cost_card(parse_card_payload(card, school, theme, cost)) for card in raw["cards"])


# Parse one model-authored effect plan.
def parse_effect_plan(raw: dict[str, Any]) -> EffectPlan:
    primitive_id = raw["primitive_id"]
    weight = max(1, int(raw.get("weight", 1)))
    return EffectPlan(primitive_id=primitive_id, weight=weight)  # type: ignore[arg-type]


# Normalize one-card or pack-shaped model JSON into one card object.
def normalize_card_response(raw: dict[str, Any]) -> dict[str, Any]:
    if "cards" in raw and isinstance(raw["cards"], list) and raw["cards"]:
        return normalize_card_response(raw["cards"][0])
    if "shape" in raw and isinstance(raw["shape"], dict):
        return normalize_card_response(raw["shape"])
    if "effects" not in raw and len(raw) == 1:
        value = next(iter(raw.values()))
        if isinstance(value, dict):
            return normalize_card_response(value)
    return raw


# Parse one llama.cpp card response, falling back on malformed JSON.
def parse_llamacpp_card_text(
    text: str,
    payload: dict[str, Any],
    pack_cards: Sequence[dict[str, Any]],
    discard_bad_text: bool = True,
) -> dict[str, Any]:
    try:
        raw = normalize_card_response(json.loads(extract_json_object(text)))
    except (TypeError, ValueError, json.JSONDecodeError):
        raw = {}
    # Small models (MiniCPM) often emit malformed JSON or nest the effects, which
    # would collapse to a single default "deal". When the structured effects yield
    # no allowed primitive, salvage the primitives the model named in its raw text.
    allowed = tuple(payload["allowed_primitives"])
    if not valid_effect_ids(raw.get("effects"), allowed):
        salvaged = salvage_effects_from_text(text, allowed)
        if salvaged:
            raw["effects"] = salvaged
    return repair_llamacpp_card(raw, payload, pack_cards, discard_bad_text)


# Return the allowed primitive ids present in a well-formed effects list.
def valid_effect_ids(raw: Any, allowed: Sequence[str]) -> list[str]:
    if not isinstance(raw, list):
        return []
    return [e["primitive_id"] for e in raw if isinstance(e, dict) and e.get("primitive_id") in allowed]


# Recover effects by scanning raw model text for allowed primitive ids in the
# order they appear (deduped, capped at two), so malformed JSON still yields
# the mechanic the model intended instead of defaulting to plain deal.
def salvage_effects_from_text(text: str, allowed: Sequence[str]) -> list[dict[str, Any]]:
    low = text.lower()
    positions: list[tuple[int, str]] = []
    for pid in allowed:
        candidates = [low.find(pid), low.find(pid.replace("_", " "))]
        hits = [pos for pos in candidates if pos != -1]
        if hits:
            positions.append((min(hits), pid))
    positions.sort()
    seen: list[str] = []
    for _, pid in positions:
        if pid not in seen:
            seen.append(pid)
    return [{"primitive_id": pid, "weight": 1} for pid in seen[:2]]


# Repair cheap llama.cpp JSON into the strict card schema.
def repair_llamacpp_card(
    raw: dict[str, Any],
    payload: dict[str, Any],
    pack_cards: Sequence[dict[str, Any]],
    discard_bad_text: bool = True,
) -> dict[str, Any]:
    allowed = tuple(payload["allowed_primitives"])
    effects = repair_effects(raw.get("effects"), allowed)
    if discard_bad_text and unusable_raw_card(raw, payload):
        raw = {}
    school = str(payload["school"])
    taken = {str(card.get("name", "")) for card in pack_cards}
    name = distinct_name(str(raw.get("name") or fallback_card_name(school, effects, pack_cards)), taken, school)
    return {
        "name": name,
        "flavor": str(raw.get("flavor", "")),
        "art_prompt": str(raw.get("art_prompt") or fallback_art_prompt(payload)),
        "effects": effects,
    }


# Return whether generated text leans into another school's imagery.
def school_flavor_mismatch(school: str, raw: dict[str, Any]) -> bool:
    text = " ".join(str(raw.get(key, "")) for key in ("name", "flavor", "art_prompt")).lower()
    words = {word.strip(".,:;!?'\"").lower() for word in text.split()}
    return any(word in words for word in SCHOOL_FORBIDDEN_IMAGERY.get(school, ()))


# Return whether raw model text must be discarded before repair.
def unusable_raw_card(raw: dict[str, Any], payload: dict[str, Any]) -> bool:
    name = str(raw.get("name", ""))
    return (
        generic_card_name(name)
        or theme_name_leak(name, str(payload["theme"]))
        or school_flavor_mismatch(str(payload["school"]), raw)
    )


# Return whether the name leaks meta theme labels instead of in-world imagery.
def theme_name_leak(name: str, theme: str) -> bool:
    name_words = {word.strip(".,:;!?'\"").lower() for word in name.split()}
    theme_words = {word.strip(".,:;!?'\"").lower() for word in theme.split()}
    return any(word in name_words for word in theme_words if word in {"anime", "fantasy", "wuxia", "school", "world"})


# Return `preferred` if it is free, else the first unused name from the school's
# pool, so a collision draws a fresh evocative name rather than a numeric suffix.
def distinct_name(preferred: str, taken: set[str], school: str) -> str:
    preferred = preferred.strip()
    if preferred and preferred not in taken:
        return preferred
    pool = FALLBACK_NAMES.get(school, FALLBACK_NAMES["fire"])
    for name in pool:
        if name not in taken:
            return name
    index = 2  # reached only past ~24 distinct names in one run; numeric is the last resort
    base = preferred or pool[0]
    while f"{base} {index}" in taken:
        index += 1
    return f"{base} {index}"


# Return a distinct primitive focus per pack slot so concurrently-built cards
# span different effects (a soft prompt hint only, so it adds no extra calls).
def pack_focuses(payload: dict[str, Any], size: int) -> list[str]:
    allowed = list(payload.get("allowed_primitives", []))
    if not allowed:
        return [""] * size
    # Rotate the focus window by a random offset each pack so packs differ from
    # one another (not always deal/burn/bomb), while staying distinct within a pack.
    start = random.randrange(len(allowed))
    return [str(allowed[(start + index) % len(allowed)]) for index in range(size)]


# Give every pick in a parallel-built pack a distinct name from the school pool.
def dedup_pack_names(cards: list[dict[str, Any]], school: str) -> list[dict[str, Any]]:
    taken: set[str] = set()
    for card in cards:
        name = distinct_name(str(card.get("name", "")), taken, school)
        card["name"] = name
        taken.add(name)
    return cards




# Repair a model-authored effect list to allowed primitive ids.
def repair_effects(raw: Any, allowed: Sequence[str]) -> list[dict[str, Any]]:
    effects = [effect for effect in raw if isinstance(effect, dict)] if isinstance(raw, list) else []
    repaired = [repair_effect(effect, allowed) for effect in effects if effect.get("primitive_id") in allowed]
    return repaired[:2] or [{"primitive_id": allowed[0], "weight": 1}]


# Repair one model-authored effect to a minimal engine-owned plan.
def repair_effect(raw: dict[str, Any], allowed: Sequence[str]) -> dict[str, Any]:
    primitive_id = raw["primitive_id"] if raw.get("primitive_id") in allowed else allowed[0]
    return {"primitive_id": primitive_id, "weight": max(1, int(raw.get("weight", 1)))}


# Build a fallback card name that avoids the names already in the pack.
def fallback_card_name(school: str, effects: Sequence[dict[str, Any]], pack_cards: Sequence[dict[str, Any]]) -> str:
    return distinct_name("", {str(card.get("name", "")) for card in pack_cards}, school)


# Build fallback art text when the model omits it.
def fallback_art_prompt(payload: dict[str, Any]) -> str:
    return f"{payload['theme']} {payload['school']} card art"


# Extract the first JSON object from model text.
def extract_json_object(text: str) -> str:
    start = text.find("{")
    if start == -1:
        raise ValueError("Codex response did not contain a JSON object")
    end = json_object_end(text, start)
    candidate = balance_json_closers(text[start:end])
    json.loads(candidate)
    return candidate


# Return the end offset of the first balanced JSON object.
def json_object_end(text: str, start: int) -> int:
    stack: list[str] = []
    in_string = False
    escaped = False
    for index, char in enumerate(text[start:], start):
        if escaped:
            escaped = False
            continue
        if char == "\\" and in_string:
            escaped = True
            continue
        if char == "\"":
            in_string = not in_string
            continue
        if in_string:
            continue
        if char == "{":
            stack.append("}")
        elif char == "[":
            stack.append("]")
        elif stack and char == stack[-1]:
            stack.pop()
            if not stack:
                return index + 1
    return len(text)


# Append missing JSON object or array closers.
def balance_json_closers(text: str) -> str:
    stack: list[str] = []
    in_string = False
    escaped = False
    for char in text:
        if escaped:
            escaped = False
            continue
        if char == "\\" and in_string:
            escaped = True
            continue
        if char == "\"":
            in_string = not in_string
            continue
        if in_string:
            continue
        if char == "{":
            stack.append("}")
        elif char == "[":
            stack.append("]")
        elif stack and char == stack[-1]:
            stack.pop()
    return text + "".join(reversed(stack))