from __future__ import annotations

import random
from collections import Counter

from .plants import POOL, plant_rarity
from .types import ALL_TYPES

PACK_TYPES = ("lesser", "tower", "general") + ALL_TYPES  # lesser, tower, general + one per type
REWARD_PACK_TYPES = ("general",) + ALL_TYPES  # good packs only (win rewards)
RARITY_WEIGHTS = {"common": 50, "rare": 35, "super_rare": 5, "super_super_rare": 2, "ultra_rare": 1}
LESSER_RARITY_WEIGHTS = {"common": 70, "rare": 25, "super_rare": 2, "super_super_rare": 0, "ultra_rare": 0}
TOWER_RARITY_WEIGHTS = {"common": 2, "rare": 15, "super_rare": 25, "super_super_rare": 15, "ultra_rare": 5}


def _weighted(names, weights_map=None) -> str:
    weights_map = weights_map or RARITY_WEIGHTS
    weights = [weights_map.get(plant_rarity(n), 20) for n in names]
    return random.choices(names, weights=weights, k=1)[0]


def open_pack(pack_type: str) -> str:
    """Return a random plant name for the given pack type."""
    if pack_type == "lesser":
        return _weighted(list(POOL.keys()), LESSER_RARITY_WEIGHTS)
    if pack_type == "tower":
        return _weighted(list(POOL.keys()), TOWER_RARITY_WEIGHTS)
    if pack_type == "general":
        return _weighted(list(POOL.keys()))
    if pack_type in ALL_TYPES:
        # type pack: roll the rarity first (with the general ladder), then pick a
        # plant of that rarity from the type; fall back to the whole pool only if
        # that rarity isn't present in the type.
        dist = _rarity_odds(list(POOL.keys()), RARITY_WEIGHTS)
        rarity = random.choices(list(dist.keys()),
                                weights=[dist[r] for r in dist], k=1)[0]
        candidates = [n for n in POOL
                      if POOL[n]["type"] == pack_type and plant_rarity(n) == rarity]
        if not candidates:
            candidates = [n for n in POOL if plant_rarity(n) == rarity]
        return random.choice(candidates)
    return _weighted(list(POOL.keys()))


def _rarity_odds(pool, weights_map) -> dict:
    """Chance (0..1) of each rarity when rolling from `pool` with `weights_map`."""
    counts = Counter(plant_rarity(n) for n in pool)
    total = sum(counts.get(r, 0) * w for r, w in weights_map.items())
    if total <= 0:
        return {r: 0.0 for r in weights_map}
    return {r: counts.get(r, 0) * weights_map[r] / total for r in weights_map}


def pack_odds(pack_type: str) -> dict:
    """Chance (0..1) of each rarity for the given pack type."""
    whole = list(POOL.keys())
    if pack_type == "lesser":
        return _rarity_odds(whole, LESSER_RARITY_WEIGHTS)
    if pack_type == "tower":
        return _rarity_odds(whole, TOWER_RARITY_WEIGHTS)
    if pack_type == "general":
        return _rarity_odds(whole, RARITY_WEIGHTS)
    # type packs roll the rarity first with the general ladder, so their rarity
    # odds match the general pack (the plant then comes from the type when possible).
    return _rarity_odds(whole, RARITY_WEIGHTS)