from typing import List

from . import collection
from .plants import POOL, VARIABLE_TYPE_UNITS, make_leveled, plant_rarity
from .types import ALL_TYPES, TYPE_BEATS
from .unit import ROLE_ATTACKER, ROLE_SUPPORT, ROLE_TANK, TEAM_ENEMY, TEAM_PLAYER, Unit


DIFFICULTY_RARITIES = {
    1: ("rare",),
    2: ("rare", "super_rare"),
    3: ("rare", "super_rare"),
    4: ("super_rare", "super_super_rare"),
    5: ("super_super_rare",),
}

# Low-rarity plants that join floor enemies at higher floors (deterministic pick).
FLOOR_RARE_PLANTS = ["Marigold", "Daffodil", "Peony", "Primrose",
                     "Hibiscus", "Dandelion", "Venus Flytrap"]
# Low-rarity tank plants used as enemy tanks on later floors.
FLOOR_TANK_PLANTS = ["Marigold", "Daffodil", "Venus Flytrap"]

# Floor enemies scale up faster and hit harder / are beefier than player units
# of the same level (a flat boost on top of the leveled stats).
FLOOR_ENEMY_LEVEL_STEP = 4  # +1 enemy level every 4 floors (was every 5)
FLOOR_ENEMY_ATK_MULT = 1.20  # +20% attack (more damage)
FLOOR_ENEMY_HP_MULT = 1.25   # +25% max health (tankier)


def floor_enemy_count(floor: int) -> int:
    """Number of enemies on a battle floor: 1 on floors 1-3, 2 on 4-9, 3 on
    10-15, then a full 4-enemy team from floor 16 on."""
    if floor <= 3:
        return 1
    if floor <= 9:
        return 2
    if floor <= 15:
        return 3
    return 4


def floor_enemy_names(floor: int, player_type: str) -> List[str]:
    """Deterministic enemy names for a battle floor. Floors 1-3 are a single
    pinwheel; the team grows from there — 3+ enemies from floor 10 and always
    4 from floor 16. Teams mix pinwheels, butterflies and toadstools, and
    duplicates are allowed (e.g. two pinwheels). Tank plants join from floor
    40 and low-rarity plants from floor 45, replacing critter slots so a team
    never exceeds 4 enemies. Boss floors field the boss alone."""
    from . import bosses
    boss = bosses.boss_for_floor(floor)
    if boss:
        return [boss]
    count = floor_enemy_count(floor)
    # Lead with an attacker; the rest cycle through the trio so each floor gets
    # a deterministic mix of roles (with repeats) built from the critters.
    names = ["Pinwheel"]
    cycle = ("Butterfly", "Toadstool", "Pinwheel", "Toadstool",
             "Butterfly", "Pinwheel", "Butterfly", "Toadstool")
    for i in range(1, count):
        names.append(cycle[(floor + i) % len(cycle)])
    # High floors mix in plants, replacing critter slots (4-enemy cap kept).
    if floor >= 40 and count >= 3:
        names[-1] = FLOOR_TANK_PLANTS[(floor // 2) % len(FLOOR_TANK_PLANTS)]
    if floor >= 45 and count >= 4:
        names[-2] = FLOOR_RARE_PLANTS[(floor * 7) % len(FLOOR_RARE_PLANTS)]
    return names


def roll_enemy_types(names: List[str], player_type: str, floor: int) -> List[str]:
    """Types for a floor's enemies (parallel to `names`). The floor 1-3 pinwheel is
    always the type the player is strong against; other critters get a fixed type
    derived from the floor (no rerolling — the same floor always fields the same
    team, and duplicate types within a team are allowed); plants keep their own
    fixed type."""
    from . import bosses
    types = []
    for i, n in enumerate(names):
        if bosses.is_boss(n):
            types.append(bosses.boss_type(n))
        elif n in VARIABLE_TYPE_UNITS:
            if floor <= 3:
                types.append(TYPE_BEATS.get(player_type, "garden"))
            else:
                types.append(ALL_TYPES[(floor * 7 + i * i) % len(ALL_TYPES)])
        else:
            types.append(POOL[n]["type"])
    return types


def floor_enemy_level(floor: int) -> int:
    """Enemy level for a floor (ramps one level every FLOOR_ENEMY_LEVEL_STEP
    floors, steeper than the old every-5 curve)."""
    return 1 + (floor - 1) // FLOOR_ENEMY_LEVEL_STEP


def make_floor_enemies(floor: int, player_type: str) -> List[Unit]:
    """Build the enemy unit team for a floor, with level/stage scaling and a
    fresh roll of critter types."""
    names = floor_enemy_names(floor, player_type)
    types = roll_enemy_types(names, player_type, floor)
    return build_floor_enemies(names, types, floor)


def build_floor_enemies(names: List[str], types: List[str], floor: int) -> List[Unit]:
    """Build enemy units from an already-rolled name/type list (so the picker and
    the actual battle always show the same enemies). Floor enemies also get a
    flat attack/health boost so they out-stat player units of the same level.
    Boss floors build the boss instead."""
    from . import bosses
    if names and bosses.is_boss(names[0]):
        return [bosses.make_boss(names[0], floor)]
    level = floor_enemy_level(floor)
    stage = max(1, min(5, 1 + (floor - 1) // 20))
    units = [make_leveled(n, TEAM_ENEMY, level, stage, ptype=t)
             for n, t in zip(names, types)]
    for u in units:
        u.attack = int(u.attack * FLOOR_ENEMY_ATK_MULT)
        u.max_hp = int(u.max_hp * FLOOR_ENEMY_HP_MULT)
        u.hp = u.max_hp
    return units


def difficulty_pool(difficulty: int) -> List[str]:
    rarities = DIFFICULTY_RARITIES.get(difficulty, DIFFICULTY_RARITIES[1])
    return [n for n in POOL if plant_rarity(n) in rarities]


TOWER_STAGES = 20
TOWER_INJECT_FLOOR = 10
TOWER_INJECT_CHANCE_MIN = 0.20
TOWER_INJECT_CHANCE_MAX = 0.60


def tower_inject_chance(stage: int) -> float:
    """Chance (0..1) that a tower enemy at `stage` carries an injected ability."""
    if stage < TOWER_INJECT_FLOOR:
        return 0.0
    frac = min(1.0, (stage - TOWER_INJECT_FLOOR) / (TOWER_STAGES - TOWER_INJECT_FLOOR))
    return TOWER_INJECT_CHANCE_MIN + frac * (TOWER_INJECT_CHANCE_MAX - TOWER_INJECT_CHANCE_MIN)


def roll_enemy_injection(name: str) -> dict:
    """Roll an injected ability for an enemy plant, or None. It comes from a
    same-role plant's extractable skills; legendary plants only pull from other
    legendary plants so they stay strong. Never replaces a slot with the same
    ability it already has."""
    import random

    from .skills import skill_to_dict
    my_rarity = plant_rarity(name)
    my_role = POOL[name]["role"]
    base = POOL[name]["skills"]
    candidates = []
    for other, spec in POOL.items():
        if other == name:
            continue
        if spec["role"] != my_role:
            continue
        if my_rarity == "super_super_rare" and plant_rarity(other) != "super_super_rare":
            continue
        for _, sk in collection.extractable_skills(other):
            candidates.append(sk)
    if not candidates:
        return None
    pairs = []
    for sk in candidates:
        for slot in (1, 2):
            if slot < len(base) and base[slot].name != sk.name:
                pairs.append((sk, slot))
    if not pairs:
        return None
    sk, slot = random.choice(pairs)
    return {"index": slot, "source": name, "skill": skill_to_dict(sk)}


def make_tower_enemies(names: List[str], level: int, stage: int) -> List[Unit]:
    """Build a tower enemy team; each enemy independently rolls a chance to
    carry one injected ability (floors 10+) and gains runes at higher stages."""
    import random

    from . import runes
    from .skills import skill_from_dict
    units = []
    estage = tower_enemy_stage(stage)
    for n in names:
        u = make_leveled(n, TEAM_ENEMY, level, estage)
        if random.random() < tower_inject_chance(stage):
            inj = roll_enemy_injection(n)
            if inj and inj["index"] < len(u.skills):
                u.skills[inj["index"]] = skill_from_dict(inj["skill"])
                u.injected_index = inj["index"]
                u.injected_name = inj["skill"]["name"]
        u.runes = runes.roll_enemy_runes(stage)
        units.append(u)
    return units


def roll_tower_enemy_team(stage: int, level: int, exclude=()) -> tuple:
    """Roll a full tower enemy team once, returning (names, injections, runes).
    The recipe is cached per stage so a retry rebuilds the exact same team
    instead of rerolling it."""
    import random

    from . import runes
    names = roll_tower_enemies(stage, exclude=exclude)
    injections = []
    for n in names:
        inj = None
        if random.random() < tower_inject_chance(stage):
            inj = roll_enemy_injection(n)
        injections.append(inj)
    enemy_runes = [runes.roll_enemy_runes(stage) for _ in names]
    return names, injections, enemy_runes


def build_tower_enemy_team(names: List[str], injections, enemy_runes,
                           level: int, stage: int) -> List[Unit]:
    """Rebuild tower enemy units from a previously rolled team recipe."""
    from .skills import skill_from_dict
    units = []
    estage = tower_enemy_stage(stage)
    for n, inj, rns in zip(names, injections, enemy_runes or []):
        u = make_leveled(n, TEAM_ENEMY, level, estage)
        if inj and inj["index"] < len(u.skills):
            u.skills[inj["index"]] = skill_from_dict(inj["skill"])
            u.injected_index = inj["index"]
            u.injected_name = inj["skill"]["name"]
        u.runes = list(rns) if rns else []
        units.append(u)
    return units


def tower_pool(stage: int) -> tuple:
    """Rarities available at a tower stage (stage 1 = rares only, ramping up)."""
    if stage <= 1:
        return ("rare",)
    if stage <= 3:
        return ("rare", "super_rare")
    if stage <= 5:
        return ("rare", "super_rare")
    if stage <= 8:
        return ("super_rare",)
    if stage <= 12:
        return ("super_rare", "super_super_rare")
    return ("super_super_rare",)


def tower_enemy_stage(stage: int) -> int:
    """Evolution stage for tower enemies: they evolve as the tower ramps, which
    multiplies their base stats (1.0 / 1.25 / 1.55)."""
    if stage >= 11:
        return 3
    if stage >= 5:
        return 2
    return 1


def tower_enemy_level(player_level: int, stage: int) -> int:
    return max(1, min(33, player_level + (stage - 1) * 3 // 4))


def roll_tower_enemies(stage: int, exclude=()) -> List[str]:
    """Roll the tower enemy team for a stage. Count grows 1 -> 4; rarity ramps.
    Tanks are scarce at high stages (there are no super-super-rare tanks), so the
    tank slot is filled from the super-rare pool when the stage pool has none."""
    import random
    pool = [n for n in POOL
            if plant_rarity(n) in tower_pool(stage) and n not in exclude]
    count = 1 if stage == 1 else min(4, (stage + 3) // 2)
    if count == 1:
        return [random.choice(pool)]

    def _by(role, source):
        return [n for n in source if POOL[n]["role"] == role]

    tank_pool = _by(ROLE_TANK, pool)
    if not tank_pool:
        tank_pool = [n for n in POOL
                     if POOL[n]["role"] == ROLE_TANK and plant_rarity(n) == "super_rare"
                     and n not in exclude]

    chosen = []
    for role, source in ((ROLE_TANK, tank_pool), (ROLE_SUPPORT, pool),
                         (ROLE_ATTACKER, pool), (ROLE_ATTACKER, pool)):
        if len(chosen) >= count:
            break
        cands = [n for n in _by(role, source) if n not in chosen]
        if cands:
            chosen.append(random.choice(cands))
    remaining = [n for n in pool if n not in chosen]
    random.shuffle(remaining)
    return (chosen + remaining)[:count]


def tower_reward(stage: int):
    """Reward for defeating a tower stage: (water, tower_packs)."""
    if stage <= 5:
        return 1, 0
    return 0, 1 + (stage - 6) // 5


def roll_enemies(difficulty: int, exclude=()) -> List[str]:
    """Roll a 4-plant enemy team for a difficulty stage, aiming for
    1 tank + 1 support + 2 attackers. Falls back gracefully when a role
    isn't present in the pool (e.g. no legendary tanks on stage 5)."""
    import random
    pool = [n for n in difficulty_pool(difficulty) if n not in exclude]
    if len(pool) < 4:  # safety fallback: never fewer than a full team
        pool = [n for n in POOL if n not in exclude]

    def _by_role(role):
        return [n for n in pool if POOL[n]["role"] == role]

    chosen = []
    for role in (ROLE_TANK, ROLE_SUPPORT, ROLE_ATTACKER, ROLE_ATTACKER):
        candidates = [n for n in _by_role(role) if n not in chosen]
        if candidates:
            chosen.append(random.choice(candidates))
    remaining = [n for n in pool if n not in chosen]
    random.shuffle(remaining)
    return (chosen + remaining)[:4]


def pick_team(names: List[str], team: int) -> List[Unit]:
    units = []
    for n in names:
        prog = collection.get_progress(n)
        units.append(make_leveled(n, team, prog["level"], prog["stage"],
                                  ptype=collection.unit_type(n)))
    return units


def make_enemy_team(names: List[str], level: int = 1, stage: int = 1) -> List[Unit]:
    return [make_leveled(n, TEAM_ENEMY, level, stage) for n in names]


def player_team_level(names: List[str]) -> int:
    levels = [collection.get_progress(n)["level"] for n in names]
    return max(1, round(sum(levels) / max(1, len(levels))))


def player_team_stage(names: List[str]) -> int:
    stages = [collection.get_progress(n)["stage"] for n in names]
    return max(1, round(sum(stages) / max(1, len(stages))))


DEFAULT_PLAYERS = ["Lavender", "Marigold", "Sunflower", "Zinnia"]
DEFAULT_ENEMIES = ["Rose", "Tulip", "Camelia", "Hellebore"]


def make_roster():
    """Deterministic level-1 rosters (used by tests; independent of save data)."""
    return ([make_leveled(n, TEAM_PLAYER, 1, 1) for n in DEFAULT_PLAYERS],
            [make_leveled(n, TEAM_ENEMY, 1, 1) for n in DEFAULT_ENEMIES])


def all_names() -> List[str]:
    return list(POOL.keys())