budget.py

from dataclasses import dataclass
from random import Random
from typing import Sequence

from primitives import Effect, PrimitiveId, School, primitive_allowed_for_school, render_effect

MAX_ENERGY_COST = 5
MIRACLE_BUDGET_CAP = 8


@dataclass(frozen=True)
class EffectPlan:
    primitive_id: PrimitiveId
    weight: int = 1


@dataclass(frozen=True)
class CardSpec:
    name: str
    cost: int
    school: School
    theme: str
    effect_plans: tuple[EffectPlan, ...]
    flavor: str = ""
    art_prompt: str = ""
    miracle: bool = False


@dataclass(frozen=True)
class Card:
    name: str
    cost: int
    school: School
    theme: str
    effects: tuple[Effect, ...]
    flavor: str = ""
    art_prompt: str = ""
    art_uri: str = ""

    # Return fixed, parseable rules text for the card.
    def rules_text(self) -> str:
        return " ".join(render_effect(effect) for effect in self.effects)


# Convert energy cost into point budget.
def budget_for_energy(cost: int) -> int:
    if cost < 1 or cost > MAX_ENERGY_COST:
        raise ValueError(f"card cost must be 1..{MAX_ENERGY_COST}")
    return cost


# Roll the bounded high-variance miracle budget.
def miracle_budget(cost: int, rng: Random) -> int:
    return min(MIRACLE_BUDGET_CAP, budget_for_energy(cost) * rng.randint(1, 3))


# Build a fully costed card from model-chosen effect shapes.
def cost_card(spec: CardSpec, rng: Random | None = None) -> Card:
    validate_card_spec(spec)
    point_budget = miracle_budget(spec.cost, rng or Random()) if spec.miracle else budget_for_energy(spec.cost)
    effects = cost_effects(point_budget, spec.effect_plans, spec.school)
    return Card(spec.name, spec.cost, spec.school, spec.theme, effects, spec.flavor, spec.art_prompt)


# Build fully costed effects from weighted primitive plans.
def cost_effects(point_budget: int, plans: Sequence[EffectPlan], school: School | None = None) -> tuple[Effect, ...]:
    usable_plans = tuple(plans[:point_budget])
    total_weight = sum(plan.weight for plan in usable_plans)
    taxed = len(usable_plans) > 1
    effects: list[Effect] = []
    spent = 0
    for index, plan in enumerate(usable_plans):
        remaining = point_budget - spent
        points = remaining if index == len(usable_plans) - 1 else max(1, point_budget * plan.weight // total_weight)
        spent += points
        effects.append(cost_effect(plan.primitive_id, points, taxed, school))
    return tuple(effects)


# Build one fully costed effect from an allocated point budget.
def cost_effect(primitive_id: PrimitiveId, points: int, taxed: bool = False, school: School | None = None) -> Effect:
    if points < 1:
        raise ValueError("effect points must be positive")
    match primitive_id:
        case "deal":
            return Effect("deal", amount=damage_amount(points, taxed))
        case "burn":
            return Effect("burn", amount=burn_amount(points, taxed), duration=2)
        case "bomb":
            return Effect("bomb", amount=damage_amount(points, taxed), delay=3)
        case "block":
            return Effect("block", amount=points * (2 if taxed else 3))
        case "ward":
            return Effect("ward", amount=points)
        case "weak":
            return Effect("weak", amount=damage_amount(points, taxed), duration=1)
        case "draw":
            return Effect("draw", amount=max(1, points * 2 // 3))
        case "energy":
            return Effect("energy", amount=max(1, points // 2))
        case "initiative":
            return Effect("initiative", duration=1)
        case "multi_hit":
            return Effect("multi_hit", amount=max(1, points), hits=2 if points < 4 else 3)
        case "vulnerable":
            return Effect("vulnerable", amount=max(1, points), duration=1 if points < 4 else 2)
        case "conditional":
            return Effect("conditional", amount=conditional_amount(points, taxed), condition="opponent_below_half")
        case "scaling":
            return Effect("scaling", amount=points, condition=scaling_condition(school))


# Return direct damage for a point allocation.
def damage_amount(points: int, taxed: bool) -> int:
    return max(1, points * 2 - (1 if taxed else 0))


# Return burn damage per tick for a point allocation.
def burn_amount(points: int, taxed: bool) -> int:
    return max(1, points - (1 if taxed and points > 1 else 0))


# Return maximum scaled conditional damage for a point allocation.
def conditional_amount(points: int, taxed: bool) -> int:
    return max(1, points - (1 if taxed and points > 1 else 0))


# Return the scaling counter used by a school.
def scaling_condition(school: School | None) -> str:
    return "shield_charge" if school == "earth" else "cards_played"


# Validate a model-authored card spec before costing it.
def validate_card_spec(spec: CardSpec) -> None:
    budget_for_energy(spec.cost)
    if not spec.effect_plans:
        raise ValueError("card must have at least one effect")
    for plan in spec.effect_plans:
        validate_effect_plan(plan, spec.school)


# Validate one model-chosen effect plan.
def validate_effect_plan(plan: EffectPlan, school: School) -> None:
    if plan.weight < 1:
        raise ValueError("effect weight must be positive")
    if not primitive_allowed_for_school(plan.primitive_id, school):
        raise ValueError(f"{plan.primitive_id} is not allowed for {school}")