pf1e-simulator/src/pf1e_simulator/spells.py

"""PF1e spells: spell definitions, effect types, and JSON loader.

Spells are SpellSpec dataclasses with typed effects:
- DamageEffect: dice damage with optional save for half
- ConditionEffect: applies a condition (looked up in CONDITIONS_BY_NAME)
- HealEffect: restores HP (capped at hp_max)
- BuffEffect: grants StatModifiers as transient effects

Range categories: personal, touch, close, medium, long.
Save types: fort, ref, will, none.

Sources: PRPG Core Rulebook, Magic chapter.
"""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING, Literal, TypeGuard, get_args

from pf1e_simulator.dice import DiceExpr, dice_to_str, parse_dice
from pf1e_simulator.effects import BonusType, StatModifier, StatTarget

if TYPE_CHECKING:
    from pf1e_simulator.loaders.foundry import Json

SPELLS_DIR = Path(__file__).resolve().parent.parent.parent / "data" / "spells"

SpellRange = Literal["personal", "touch", "close", "medium", "long"]
SaveType = Literal["fort", "ref", "will", "none"]
# PF1e area templates (CRB, Magic - Aiming a Spell): burst (circle around a
# point), cone (90° quarter circle from the caster), line (5-ft ray from the
# caster to its length).
AreaShape = Literal["burst", "cone", "line"]
# PF1e casting times, normalized to budget codes (see combat._cast_time_slot).
CastTime = Literal["standard", "swift", "immediate", "free", "full_round", "one_round", "long"]


@dataclass(frozen=True)
class DamageDieScaling:
    """Per-caster-level extra damage dice (PF1e: "1d6 per caster level, max 10d6").

    - dice: the dice added per scaling step (e.g. 1d6)
    - per_levels: caster levels required per added die
    - max_dice: cap on the total dice added by the scaling
    """
    dice: DiceExpr
    per_levels: int
    max_dice: int


@dataclass(frozen=True)
class DamageFlatScaling:
    """Per-caster-level flat damage bonus (PF1e: "+1 per caster level, max +5").

    - amount: flat points added per scaling step
    - per_levels: caster levels required per added step
    - max_amount: cap on the total flat bonus from scaling (JSON key "max")
    """
    amount: int
    per_levels: int
    max_amount: int


@dataclass(frozen=True)
class ProjectileScaling:
    """Multi-projectile count scaling with caster level.

    PF1e (SRD spells/magic-missile.md): "one missile, plus one additional
    missile per two caster levels beyond 1st [...] maximum of five missiles";
    (SRD spells/scorching-ray.md): "one ray, plus one additional ray for every
    four levels beyond 3rd (to a maximum of three rays at 11th level)".

    - base: projectiles at the reference level (usually 1)
    - at_level: caster level from which extra projectiles begin counting
    - per_levels: caster levels required per additional projectile
    - max_projectiles: cap (JSON key "max")
    """
    base: int
    at_level: int
    per_levels: int
    max_projectiles: int

    def count_for(self, caster_level: int) -> int:
        """Projectiles fired by a caster of ``caster_level``.

        ``base + (caster_level - at_level) // per_levels``, the level delta
        floored at 0 (a caster below the reference level still fires the base
        projectiles), capped at ``max_projectiles``.
        """
        extra = max(0, caster_level - self.at_level) // self.per_levels
        return min(self.max_projectiles, self.base + extra)


@dataclass(frozen=True)
class RecurringDamageSpec:
    """Per-round recurring spell damage (PF1e Acid Arrow).

    - formula: dice rolled in each recurring round (2d4)
    - types: energy types (acid) — reduced by energy resistance/immunity,
      never by DR (it is spell damage)
    - base_rounds / per_levels / max_rounds: rounds of recurrence scale with
      caster level as ``base_rounds + caster_level // per_levels`` capped at
      ``max_rounds`` (SRD spells/acid-arrow.md: "1 round + 1 round per three
      levels [...] to a maximum of 6 additional rounds").
    """
    formula: DiceExpr
    types: tuple[str, ...]
    base_rounds: int
    per_levels: int
    max_rounds: int

    def rounds_for(self, caster_level: int) -> int:
        """Recurring rounds for a caster of ``caster_level``."""
        return min(self.max_rounds, self.base_rounds + caster_level // self.per_levels)


@dataclass(frozen=True)
class DamageEffect:
    formula: DiceExpr
    types: tuple[str, ...]
    # False = save negates (targeted); True = half damage on save (ray/line/cone/AoE)
    half_on_save: bool = False
    # Only consulted when half_on_save is False. True = a successful save deals
    # no damage (save negates, the targeted-spell default); False = the save has
    # no effect on this damage at all (SRD Sound Burst: "Fortitude partial"
    # negates only the stun — every creature takes the full dice regardless).
    save_negates: bool = True
    # PF1e caster-level scaling (CRB Magic — Caster Level; e.g. Fireball
    # "1d6 points of fire damage per caster level (maximum 10d6)"). Absent =
    # level-independent damage, the pre-S1b behaviour.
    scaling: DamageDieScaling | None = None
    scaling_flat: DamageFlatScaling | None = None
    # PF1e multi-projectile spells (Magic Missile, Scorching Ray): the effect
    # fires `count_for(caster_level)` projectiles, each rolling its own damage
    # and (for a ray) its own attack. Absent = a single projectile.
    projectiles: ProjectileScaling | None = None
    # PF1e recurring (damage-over-time) component (Acid Arrow): after the
    # initial hit the target takes this damage again at the start of each of
    # its turns for `rounds_for(caster_level)` rounds. Absent = no DoT.
    recurring: RecurringDamageSpec | None = None

    def effective_formula(self, caster_level: int) -> DiceExpr:
        """The damage formula resolved for a caster of ``caster_level``.

        PF1e (CRB, Magic — Caster Level): spells whose damage is "per caster
        level" grow with the caster and stop at the stated cap. Extra dice =
        min(max_dice, level // per_levels); flat bonus = min(max,
        (level // per_levels) * amount). The loader guarantees the scaling
        dice share the base formula's die type (or the base is flat).
        """
        count = self.formula.count
        sides = self.formula.sides
        bonus = self.formula.bonus
        if self.scaling is not None:
            extra = min(
                self.scaling.max_dice, caster_level // self.scaling.per_levels
            ) * self.scaling.dice.count
            if count == 0 and extra > 0:
                sides = self.scaling.dice.sides
            count += extra
        if self.scaling_flat is not None:
            bonus += min(
                self.scaling_flat.max_amount,
                (caster_level // self.scaling_flat.per_levels) * self.scaling_flat.amount,
            )
        return DiceExpr(count=count, sides=sides, bonus=bonus)


@dataclass(frozen=True)
class ConditionEffect:
    condition_name: str
    duration_rounds: int
    save_negates: bool = True


@dataclass(frozen=True)
class HealEffect:
    formula: DiceExpr


@dataclass(frozen=True)
class BuffEffect:
    modifiers: tuple[StatModifier, ...]
    duration_rounds: int


@dataclass(frozen=True)
class AreaEffect:
    """Area of effect spell (burst / cone / line).

    PF1e (CRB, Magic - Aiming a Spell): a burst/spread centers on a chosen
    point; a cone "shoots away from you in a quarter-circle in the direction
    you designate"; a line "shoots away from you in a line in the direction
    you designate" and "affects all creatures in squares through which the
    line passes". Cones and lines therefore originate at the CASTER and are
    aimed at the resolved target cell (or the U7b point-cast center).

    - shape: "burst" (circle), "cone" (90° quarter circle) or "line" (5 ft wide)
    - radius_ft: burst radius / cone length / line length, in feet
    - effects: the effects applied to each creature in the area
    - save_for_half: if True, successful save halves damage
    """
    shape: AreaShape = "burst"
    radius_ft: int = 20
    effects: tuple[SpellEffect, ...] = ()
    save_for_half: bool = True

@dataclass(frozen=True)
class SummonEffect:
    """Summoning spell effect.

    - creature: path to the creature JSON file (e.g., "data/monsters/celestial_wolf.json")
    - duration_rounds: how long the summoned creature stays
    """
    creature: str
    duration_rounds: int


SpellEffect = DamageEffect | ConditionEffect | HealEffect | BuffEffect | AreaEffect | SummonEffect


@dataclass(frozen=True)
class SpellSpec:
    name: str
    level: int
    school: str
    range: SpellRange = "medium"
    save: SaveType = "none"
    base_dc: int = 10
    # PF1e (CRB, Magic — Ranged Touch Attacks): the spell requires a ranged
    # touch attack to deliver (rays). Auto-hit spells (Magic Missile) and
    # area/save-only spells keep this False.
    requires_ranged_touch: bool = False
    cast_time: CastTime = "standard"
    effects: tuple[SpellEffect, ...] = ()
    description: str = ""


def spell_range_ft(spell: SpellSpec, caster_level: int) -> int:
    if spell.range == "personal":
        return 0
    if spell.range == "touch":
        return 5
    if spell.range == "close":
        return 25 + 5 * (caster_level // 2)
    if spell.range == "medium":
        return 100 + 10 * caster_level
    return 400 + 40 * caster_level


_FEAR_CONDITIONS = frozenset({"frightened", "panicked", "shaken", "cowering"})


def save_descriptors_for_effects(effects: tuple[SpellEffect, ...]) -> frozenset[str]:
    descriptors: set[str] = set()
    for effect in effects:
        if isinstance(effect, DamageEffect):
            descriptors.update(effect.types)
        elif isinstance(effect, ConditionEffect) and effect.condition_name in _FEAR_CONDITIONS:
            descriptors.add("fear")
        elif isinstance(effect, AreaEffect):
            descriptors.update(save_descriptors_for_effects(effect.effects))
    return frozenset(descriptors)


def spell_save_descriptors(spell: SpellSpec) -> frozenset[str]:
    """Spell save descriptors: effect-derived plus CRB school tags.

    Enchantment is mind_affecting, divination is divination.
    """
    descriptors = set(save_descriptors_for_effects(spell.effects))
    if spell.school == "enchantment":
        descriptors.add("mind_affecting")
    elif spell.school == "divination":
        descriptors.add("divination")
    return frozenset(descriptors)


def _as_dict(value: Json) -> dict[str, Json]:
    if isinstance(value, dict):
        return value
    msg = f"expected object, got {value!r}"
    raise ValueError(msg)


def _as_list(value: Json) -> list[Json]:
    if isinstance(value, list):
        return value
    msg = f"expected list, got {value!r}"
    raise ValueError(msg)


def _as_str(value: Json) -> str:
    if isinstance(value, str):
        return value
    msg = f"expected string, got {value!r}"
    raise ValueError(msg)


def _as_int(value: Json) -> int:
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        return int(value)
    if isinstance(value, str):
        return int(value)
    msg = f"expected integer, got {value!r}"
    raise ValueError(msg)


def _as_optional_str(value: Json) -> str | None:
    if value is None:
        return None
    return _as_str(value)


def _is_spell_range(value: str) -> TypeGuard[SpellRange]:
    return value in ("personal", "touch", "close", "medium", "long")


def _is_save_type(value: str) -> TypeGuard[SaveType]:
    return value in ("fort", "ref", "will", "none")


def _is_area_shape(value: str) -> TypeGuard[AreaShape]:
    return value in get_args(AreaShape)


def _is_cast_time(value: str) -> TypeGuard[CastTime]:
    return value in get_args(CastTime)


# RAW phrasings accepted by the JSON loader, mapped onto the budget codes.
_CAST_TIME_ALIASES: dict[str, CastTime] = {
    "1 round": "one_round",
    "one round": "one_round",
    "1 full-round action": "full_round",
    "full-round action": "full_round",
    "full-round": "full_round",
    "1 minute": "long",
    "3 rounds": "long",
    "10 minutes": "long",
    "1 hour": "long",
    "standard action": "standard",
    "swift action": "swift",
    "immediate action": "immediate",
    "free action": "free",
}


def _parse_cast_time(value: str) -> CastTime:
    """Casting time as a budget code, accepting the codes and RAW aliases."""
    code = value.strip().lower()
    if _is_cast_time(code):
        return code
    if code in _CAST_TIME_ALIASES:
        return _CAST_TIME_ALIASES[code]
    msg = f"unknown cast time: {value!r} (expected one of {get_args(CastTime)})"
    raise ValueError(msg)


def _is_stat_target(value: str) -> TypeGuard[StatTarget]:
    return value in get_args(StatTarget)


def _is_bonus_type(value: str) -> TypeGuard[BonusType]:
    return value in get_args(BonusType)


def _parse_modifier(value: Json) -> StatModifier:
    m = _as_dict(value)
    target = _as_str(m["target"])
    if not _is_stat_target(target):
        msg = f"unknown stat target: {target!r}"
        raise ValueError(msg)
    bonus = _as_str(m.get("bonus_type", ""))
    if not _is_bonus_type(bonus):
        msg = f"unknown bonus type: {bonus!r}"
        raise ValueError(msg)
    return StatModifier(
        target=target,
        value=_as_int(m["value"]),
        bonus_type=bonus,
        condition=_as_optional_str(m.get("condition")),
        weapon_filter=_as_optional_str(m.get("weapon_filter")),
    )


def _parse_die_scaling(value: Json) -> DamageDieScaling:
    data = _as_dict(value)
    dice = parse_dice(_as_str(data["dice"]))
    per_levels = _as_int(data.get("per_levels", 1))
    if per_levels < 1:
        msg = f"scaling per_levels must be >= 1, got {per_levels}"
        raise ValueError(msg)
    max_dice = _as_int(data["max_dice"])
    if max_dice < 0:
        msg = f"scaling max_dice must be >= 0, got {max_dice}"
        raise ValueError(msg)
    return DamageDieScaling(dice=dice, per_levels=per_levels, max_dice=max_dice)


def _parse_flat_scaling(value: Json) -> DamageFlatScaling:
    data = _as_dict(value)
    amount = _as_int(data.get("amount", 1))
    per_levels = _as_int(data.get("per_levels", 1))
    if per_levels < 1:
        msg = f"scaling_flat per_levels must be >= 1, got {per_levels}"
        raise ValueError(msg)
    max_amount = _as_int(data["max"])
    if max_amount < 0:
        msg = f"scaling_flat max must be >= 0, got {max_amount}"
        raise ValueError(msg)
    return DamageFlatScaling(
        amount=amount, per_levels=per_levels, max_amount=max_amount
    )


def _parse_projectiles(value: Json) -> ProjectileScaling:
    data = _as_dict(value)
    base = _as_int(data.get("base", 1))
    at_level = _as_int(data.get("at_level", 1))
    per_levels = _as_int(data["per_levels"])
    if base < 1:
        msg = f"projectiles base must be >= 1, got {base}"
        raise ValueError(msg)
    if per_levels < 1:
        msg = f"projectiles per_levels must be >= 1, got {per_levels}"
        raise ValueError(msg)
    if at_level < 1:
        msg = f"projectiles at_level must be >= 1, got {at_level}"
        raise ValueError(msg)
    max_projectiles = _as_int(data["max"])
    if max_projectiles < base:
        msg = f"projectiles max must be >= base, got {max_projectiles} < {base}"
        raise ValueError(msg)
    return ProjectileScaling(
        base=base,
        at_level=at_level,
        per_levels=per_levels,
        max_projectiles=max_projectiles,
    )


def _parse_recurring(value: Json) -> RecurringDamageSpec:
    data = _as_dict(value)
    formula = parse_dice(_as_str(data["formula"]))
    base_rounds = _as_int(data.get("base_rounds", 1))
    per_levels = _as_int(data["per_levels"])
    max_rounds = _as_int(data["max"])
    if base_rounds < 0:
        msg = f"recurring base_rounds must be >= 0, got {base_rounds}"
        raise ValueError(msg)
    if per_levels < 1:
        msg = f"recurring per_levels must be >= 1, got {per_levels}"
        raise ValueError(msg)
    if max_rounds < base_rounds:
        msg = f"recurring max must be >= base_rounds, got {max_rounds} < {base_rounds}"
        raise ValueError(msg)
    return RecurringDamageSpec(
        formula=formula,
        types=tuple(_as_str(t) for t in _as_list(data["types"])),
        base_rounds=base_rounds,
        per_levels=per_levels,
        max_rounds=max_rounds,
    )


def _parse_effect(data: Json) -> SpellEffect:
    effect = _as_dict(data)
    effect_type = _as_str(effect["type"])
    if effect_type == "damage":
        formula = parse_dice(_as_str(effect["formula"]))
        scaling_raw = effect.get("scaling")
        scaling = None if scaling_raw is None else _parse_die_scaling(scaling_raw)
        if (
            scaling is not None
            and formula.count > 0
            and scaling.dice.sides != formula.sides
        ):
            msg = (
                "damage scaling dice must match the base formula die type: "
                f"{dice_to_str(scaling.dice)} vs {dice_to_str(formula)}"
            )
            raise ValueError(msg)
        flat_raw = effect.get("scaling_flat")
        scaling_flat = None if flat_raw is None else _parse_flat_scaling(flat_raw)
        proj_raw = effect.get("projectiles")
        projectiles = None if proj_raw is None else _parse_projectiles(proj_raw)
        rec_raw = effect.get("recurring")
        recurring = None if rec_raw is None else _parse_recurring(rec_raw)
        return DamageEffect(
            formula=formula,
            types=tuple(_as_str(t) for t in _as_list(effect["types"])),
            half_on_save=bool(effect.get("half_on_save", False)),
            save_negates=bool(effect.get("save_negates", True)),
            scaling=scaling,
            scaling_flat=scaling_flat,
            projectiles=projectiles,
            recurring=recurring,
        )
    if effect_type == "condition":
        return ConditionEffect(
            condition_name=_as_str(effect["condition_name"]),
            duration_rounds=_as_int(effect["duration_rounds"]),
            save_negates=bool(effect.get("save_negates", True)),
        )
    if effect_type == "heal":
        return HealEffect(
            formula=parse_dice(_as_str(effect["formula"])),
        )
    if effect_type == "buff":
        return BuffEffect(
            modifiers=tuple(_parse_modifier(m) for m in _as_list(effect["modifiers"])),
            duration_rounds=_as_int(effect.get("duration_rounds", 0)),
        )
    if effect_type == "area":
        shape_raw = _as_str(effect.get("shape", "burst"))
        if not _is_area_shape(shape_raw):
            msg = f"unknown area shape: {shape_raw!r} (expected one of {get_args(AreaShape)})"
            raise ValueError(msg)
        return AreaEffect(
            shape=shape_raw,
            radius_ft=_as_int(effect.get("radius_ft", 20)),
            effects=tuple(_parse_effect(e) for e in _as_list(effect["effects"])),
            save_for_half=bool(effect.get("save_for_half", True)),
        )
    if effect_type == "summon":
        return SummonEffect(
            creature=_as_str(effect["creature"]),
            duration_rounds=_as_int(effect["duration_rounds"]),
        )
    msg = f"Unknown effect type: {effect_type}"
    raise ValueError(msg)


def load_spell(path: str | Path) -> SpellSpec:
    data = _as_dict(json.loads(Path(path).read_text()))
    effects = tuple(_parse_effect(e) for e in _as_list(data.get("effects", [])))
    range_raw = _as_str(data.get("range", "medium"))
    if not _is_spell_range(range_raw):
        msg = f"unknown spell range: {range_raw!r}"
        raise ValueError(msg)
    save_raw = _as_str(data.get("save", "none"))
    if not _is_save_type(save_raw):
        msg = f"unknown save type: {save_raw!r}"
        raise ValueError(msg)
    cast_time = _parse_cast_time(_as_str(data.get("cast_time", "standard")))
    return SpellSpec(
        name=_as_str(data["name"]),
        level=_as_int(data["level"]),
        school=_as_str(data["school"]),
        range=range_raw,
        save=save_raw,
        base_dc=_as_int(data.get("base_dc", 10)),
        requires_ranged_touch=bool(data.get("requires_ranged_touch", False)),
        cast_time=cast_time,
        effects=effects,
        description=_as_str(data.get("description", "")),
    )


def load_spell_registry(dir_path: str | Path) -> dict[str, SpellSpec]:
    registry: dict[str, SpellSpec] = {}
    for path in sorted(Path(dir_path).glob("*.json")):
        spell = load_spell(path)
        registry[spell.name] = spell
    return registry