from __future__ import annotations

import random
from typing import Callable, Iterable, List, Optional

from . import runes
from .skills import (
    S_ATTACK_DOWN, S_ATTACK_UP, S_CRIT_DOWN, S_DEFENSE_DOWN, S_DEFENSE_UP,
    S_HEAL_BOOST, S_HOT, S_PREVENT_FIRST_DAMAGE, S_SHIELD, S_SNOWFLAKES, S_SPEED_DOWN,
    S_SPEED_UP, S_VULNERABLE, SHIELD_DURATION,
)


class Relic:
    """A roguelike blessing collected during a Tower run. `apply` runs once at
    the start of each battle for the rest of the run. Rune-reward relics
    (`rune_reward` set) have no battle effect and instead grant runes when the
    run ends; they may be picked more than once."""

    __slots__ = ("key", "name", "desc", "rarity", "apply", "rune_reward")

    def __init__(self, key: str, name: str, desc: str, rarity: str,
                 apply: Callable[["Battle"], None], rune_reward: str = "") -> None:
        self.key = key
        self.name = name
        self.desc = desc
        self.rarity = rarity
        self.apply = apply
        self.rune_reward = rune_reward

    def __repr__(self) -> str:
        return self.name


# --- apply helpers ------------------------------------------------------

def _players(battle):
    return [u for u in battle.players if u.alive]


def _pct_stat(attr, pct):
    def apply(battle):
        for u in _players(battle):
            setattr(u, attr, int(getattr(u, attr) * (1 + pct)))
    return apply


def _add_stat(attr, amount):
    def apply(battle):
        for u in _players(battle):
            setattr(u, attr, getattr(u, attr) + amount)
    return apply


def _pct_hp(pct):
    def apply(battle):
        for u in _players(battle):
            u.max_hp = int(u.max_hp * (1 + pct))
            u.hp = u.max_hp
    return apply


def _crit(pct):
    def apply(battle):
        for u in _players(battle):
            u.crit = min(1.0, u.crit + pct)
    return apply


def _shield(pct):
    def apply(battle):
        for u in _players(battle):
            u.status.setdefault(S_SHIELD, []).append((int(u.max_hp * pct), SHIELD_DURATION))
    return apply


def _buff(key, value, duration):
    def apply(battle):
        for u in _players(battle):
            u.status[key] = (value, duration)
    return apply


def _enemy_debuff(key, value, duration):
    def apply(battle):
        for u in battle.enemies:
            if u.alive:
                u.status[key] = (value, duration)
    return apply


def _heal_boost(pct):
    def apply(battle):
        for u in _players(battle):
            u.status[S_HEAL_BOOST] = pct
    return apply


def _hot(amount, turns):
    def apply(battle):
        for u in _players(battle):
            u.status[S_HOT] = [(amount, turns)]
    return apply


def _snowflakes(value):
    def apply(battle):
        for u in _players(battle):
            u.status[S_SNOWFLAKES] = u.status.get(S_SNOWFLAKES, 0) + value
    return apply


def _combine(*apps):
    def apply(battle):
        for fn in apps:
            fn(battle)
    return apply


def _prevent_first_damage():
    def apply(battle):
        for u in _players(battle):
            u.status[S_PREVENT_FIRST_DAMAGE] = 1
    return apply


def _enemy_vulnerable(hits):
    def apply(battle):
        for u in battle.enemies:
            if u.alive:
                u.status[S_VULNERABLE] = u.status.get(S_VULNERABLE, 0) + hits
    return apply


# --- the relic pool -----------------------------------------------------

RELIC_POOL: List[Relic] = [
    # common
    Relic("iron_bark", "Iron Bark",
          "Your team starts the battle with a shield for 10% of their max health.",
          "common", _shield(0.10)),
    Relic("vitality", "Vitality",
          "+12% max health for your team.",
          "common", _pct_hp(0.12)),
    Relic("granite", "Granite Roots",
          "+10% defense for your team.",
          "common", _pct_stat("defense", 0.10)),
    Relic("sunlit", "Sunlit Growth",
          "+8% attack for your team.",
          "common", _pct_stat("attack", 0.08)),
    Relic("swift", "Swift Petals",
          "+5 speed for your team.",
          "common", _add_stat("speed", 5)),
    Relic("sharp_eye", "Sharp Vision",
          "+5% critical chance for your team.",
          "common", _crit(0.05)),
    Relic("morning_dew", "Morning Dew",
          "Your team starts the battle with Attack Up for 1 turn.",
          "common", _buff(S_ATTACK_UP, 10, 1)),
    Relic("dawn_aura", "Dawn Aura",
          "Your team starts the battle with Defense Up for 1 turn.",
          "common", _buff(S_DEFENSE_UP, 8, 1)),
    Relic("regrowth", "Regrowth",
          "Your team starts the battle with a heal-over-time (20 HP/turn for 3 turns).",
          "common", _hot(20, 3)),
    # rare
    Relic("living_spring", "Living Spring",
          "Your team's heals are 15% stronger.",
          "rare", _heal_boost(15)),
    Relic("chilling_gust", "Chilling Gust",
          "Enemies start the battle with Speed Down for 1 turn.",
          "rare", _enemy_debuff(S_SPEED_DOWN, 8, 1)),
    Relic("weakening_fog", "Weakening Fog",
          "Enemies start the battle with Attack Down for 1 turn.",
          "rare", _enemy_debuff(S_ATTACK_DOWN, 10, 1)),
    Relic("frost_blades", "Frosted Blades",
          "Your team starts the battle with +5 bonus hit damage.",
          "rare", _snowflakes(5)),
    Relic("gale_wind", "Gale Wind",
          "+10 speed for your team.",
          "rare", _add_stat("speed", 10)),
    Relic("rune_cache", "Rune Cache",
          "No battle effect. When the tower run ends, gain a random Tier 1 rune. "
          "Can be picked more than once.",
          "rare", lambda battle: None, rune_reward="tier1"),
    # epic
    Relic("thicket", "Thick Thicket",
          "Your team starts the battle with a shield for 25% of their max health.",
          "epic", _shield(0.25)),
    Relic("rusting_rain", "Rusting Rain",
          "Enemies start the battle with Defense Down for 1 turn.",
          "epic", _enemy_debuff(S_DEFENSE_DOWN, 20, 1)),
    Relic("overpower", "Overpower",
          "+15% attack for your team.",
          "epic", _pct_stat("attack", 0.15)),
    Relic("apex", "Apex Growth",
          "+20% max health for your team.",
          "epic", _pct_hp(0.20)),
    Relic("fearless", "Fearless",
          "+10% critical chance for your team.",
          "epic", _crit(0.10)),
    Relic("expose", "Expose",
          "Enemies start the battle Vulnerable for 2 hits.",
          "epic", _enemy_vulnerable(2)),
    Relic("rune_crypt", "Rune Crypt",
          "No battle effect. When the tower run ends, gain a Tier 2 rune (brown). "
          "Can be picked more than once.",
          "epic", lambda battle: None, rune_reward="brown"),
    # legendary
    Relic("protective_aura", "Protective Aura",
          "Each ally takes 80% less damage from the first hit they receive in a battle.",
          "legendary", _prevent_first_damage()),
    Relic("aegis", "Aegis",
          "Allies start the battle with a shield for 40% of their max health.",
          "legendary", _shield(0.40)),
    Relic("overgrowth", "Overgrowth",
          "Allies start the battle with +20% max health and +10% attack.",
          "legendary", _combine(_pct_hp(0.20), _pct_stat("attack", 0.10))),
    Relic("second_wind", "Second Wind",
          "The first time an ally would fall in a battle, they survive at 50% HP instead.",
          "legendary", lambda battle: None),
]

RELIC_BY_KEY = {r.key: r for r in RELIC_POOL}

# Relative draw weight per rarity tier; legendary relics should show up rarely.
RELIC_RARITY_WEIGHT = {"common": 100, "rare": 60, "epic": 35, "legendary": 12}


def _tier_weights(stage):
    """Draw weights for a tower stage. Better tiers get rarer at low stages and
    more common at high stages, but are never impossible."""
    frac = 0.0
    if stage:
        frac = min(1.0, max(0.0, (stage - 1) / 19.0))  # 0 at stage 1, 1 at stage 20
    common = 100
    rare = 60 * (0.9 + 0.2 * frac)
    epic = 35 * (0.7 + 0.6 * frac)
    legendary = 12 * (0.5 + 1.0 * frac)
    return {"common": common, "rare": rare, "epic": epic, "legendary": legendary}


def by_key(key: str) -> Optional[Relic]:
    return RELIC_BY_KEY.get(key)


def rune_reward_count(relic_keys: Iterable[str]) -> int:
    """How many runes the picked rune-reward relics will grant when the run ends."""
    return sum(1 for k in relic_keys if getattr(by_key(k), "rune_reward", ""))


def grant_rune_rewards(relic_keys: Iterable[str]) -> list:
    """Grant the runes promised by any rune-reward relics collected during a
    tower run (called when the run ends). Returns the granted runes."""
    from . import collection  # lazy: collection doesn't import relics
    granted = []
    for key in relic_keys:
        relic = by_key(key)
        reward = getattr(relic, "rune_reward", "")
        if reward == "tier1":
            granted.append(collection.add_rune(random.choice(runes.TIER1_COLORS)))
        elif reward == "brown":
            granted.append(collection.add_rune("brown"))
    return granted


def roll_relics(n: int = 3, exclude: Iterable[str] = (), stage: Optional[int] = None) -> List[Relic]:
    """Roll `n` distinct relics, excluding already-owned keys. Draws are weighted
    by rarity, scaled by the tower stage so better tiers appear less in the early
    stages (but are never impossible). Rune-reward relics are never excluded, so
    they can be picked more than once in a run."""
    ex = set(exclude)
    pool = [r for r in RELIC_POOL if (r.key not in ex) or r.rune_reward]
    if not pool:
        return []
    weights_by_tier = _tier_weights(stage)
    weights = [weights_by_tier.get(r.rarity, 50) for r in pool]
    chosen = []
    while len(chosen) < n and pool:
        pick = random.choices(pool, weights=weights, k=1)[0]
        chosen.append(pick)
        idx = pool.index(pick)
        pool.pop(idx)
        weights.pop(idx)
    return chosen