#!/usr/bin/env python3
"""Balance report: an hour of fighting for each trade with its own weapon, at each tier.

    python3 tools/balance.py [solo|party|bosses] [seeds] [levels, e.g. 22,52] [potions]
    python3 tools/balance.py dungeon [seeds] [level] [perfect]

`dungeon` plays each pull of every dungeon for a group of four and a pair, everyone with the
tier's gear, food and health potions. `perfect` stands in for active play: real cooldowns, no
waiting on auto's rules, quick interrupts. Bosses should beat an all-auto group and fall to one
that plays.

Each row is one encounter: wins and losses in the hour, XP an hour, food eaten an hour,
and the share of a kill's time spent fighting (the rest is resting between fights).
"""

from __future__ import annotations

import random
import statistics
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import game  # noqa: E402

T0 = 1_800_000_000.0
WEIGHT = {"miner": "heavy", "builder": "heavy", "homesteader": "light", "ranger": "medium"}
# (combat level, the tier of gear worn, the encounters to try)
CHECKPOINTS = [
    (10, 0, ["fight_wolves", "fight_goblins", "fight_scouts"]),
    (22, 1, ["fight_clay_golems", "fight_cave_trolls", "fight_restless_dead", "fight_harpies", "fight_ruins"]),
    (37, 2, ["fight_stone_guardians", "fight_lizardfolk", "fight_bog_wraiths", "fight_giant_spiders", "fight_barrow_king"]),
    (52, 3, ["fight_iron_golems", "fight_orcs", "fight_mine_ghosts", "fight_wyverns", "fight_green_dragon"]),
    (67, 4, ["fight_ice_elementals", "fight_frost_trolls", "fight_banshees", "fight_griffins", "fight_blue_dragon"]),
    (82, 5, ["fight_magma_constructs", "fight_dragonkin", "fight_shades", "fight_rocs", "fight_red_dragon"]),
]
FOODS = ["cooked_shrimp", "cooked_trout", "cooked_salmon", "smoked_eel", "cooked_lobster", "cooked_swordfish"]


def gear(cls: str, tier: int) -> dict:
    """The trade's own weapon and its usual armour weight, all of one tier (and a shield for a sword)."""
    style = game.CLASS_STYLE[cls]
    equip = {}
    for item, spec in game.ITEMS.items():
        if spec.get("tier") != tier or spec.get("req") != ["combat", game.TIER_LEVELS[tier]]:
            continue
        if spec.get("style") == style or (style == "sword" and spec.get("slot") == "shield") or \
                spec.get("weight") == WEIGHT[cls] and spec["slot"] in ("helmet", "body", "legs"):
            equip[spec["slot"]] = item
    return equip


POTIONS = False


def fighter(world: dict, uid: str, cls: str, level: int, tier: int) -> dict:
    player = game.new_player(uid, uid.title(), cls, T0)
    world["players"][uid] = player
    player["xp"]["combat"] = game.XP_TABLE[level]
    player["equip"].update(gear(cls, tier))
    player["hp"] = game.max_hp(player)
    player["bank"] = {FOODS[tier]: 50_000, "arrows": 10_000}
    if POTIONS:
        player["bank"][f"health_potion_{tier + 1}"] = 10_000
    return player


def hour(classes: list, level: int, tier: int, fight: str, seed: int) -> dict:
    world = game.new_world(T0)
    players = [fighter(world, f"p{i}", cls, level, tier) for i, cls in enumerate(classes)]
    game.start_action(world, players[0], fight, T0)
    for other in players[1:]:
        game.join_fight(world, other, players[0]["id"], T0)
    rng = random.Random(seed)
    end = T0 + 3600
    for now in range(int(T0) + 60, int(end) + 1, 60):
        for player in players:
            player["last_active"] = now
        game.catch_up(world, now, rng)
    away = players[0].get("away") or game.empty_summary()
    eaten = sum(away["eaten"].values())
    return {"wins": away["wins"], "losses": away["losses"], "xp": away["xp"].get("combat", 0), "eaten": eaten,
            "stopped": away["stopped"]}


def report(groups: dict, seeds: int, levels: set | None = None) -> None:
    print(f"{'who':24} {'lvl':>3} {'encounter':22} {'wins':>5} {'lost':>4} {'xp/h':>6} {'food/h':>6} {'stop':>4}")
    for name, classes in groups.items():
        for level, tier, fights in CHECKPOINTS:
            if levels and level not in levels:
                continue
            for fight in fights:
                rows = [hour(classes, level, tier, fight, seed) for seed in range(seeds)]
                mean = lambda key: statistics.mean(row[key] for row in rows)
                stops = sum(1 for row in rows if row["stopped"])
                label = game.ACTION_BY_ID[fight]["name"][:22]
                print(f"{name:24} {level:>3} {label:22} {mean('wins'):>5.0f} {mean('losses'):>4.1f} "
                      f"{mean('xp'):>6.0f} {mean('eaten'):>6.0f} {stops:>4}")
        print()


BOSSES = ["fight_ruins", "fight_barrow_king", "fight_green_dragon", "fight_blue_dragon", "fight_red_dragon"]


def dungeon(seeds: int, level: int | None) -> None:
    for key, spec in game.DUNGEONS.items():
        tier = spec["tier"]
        at = level or spec["level"] + 3
        action = game.ACTION_BY_ID[key]
        for group in (["builder", "homesteader", "miner", "ranger"], ["builder", "homesteader"]):
            global POTIONS
            POTIONS = True
            print(f"{spec['name']}: {'+'.join(group)} at {at}")
            print(f"{'pull':18} {'win':>5} {'secs':>5} {'downs':>5} {'potions':>7}")
            for index, (name, _pack) in enumerate(spec["pulls"]):
                wins, secs, downs, drunk = 0, [], [], []
                for seed in range(seeds):
                    world = game.new_world(T0)
                    players = [fighter(world, f"p{i}", cls, at, tier) for i, cls in enumerate(group)]
                    fight = game.draw_pull(world, players, action, {"pull": index, "wipes": 0}, random.Random(seed), T0)
                    if fight["result"] == "won":
                        wins += 1
                        secs.append(fight["duration"] - fight["rest"])
                    downs.append(sum(1 for m in fight["members"] if m["down"]))
                    drunk.append(sum(1 for m in fight["members"] if m["drank"]))
                mean = lambda xs: statistics.mean(xs) if xs else 0
                print(f"{name:18} {wins / seeds:>5.0%} {mean(secs):>5.0f} {mean(downs):>5.1f} {mean(drunk):>7.1f}")
            print()


def main() -> None:
    mode = sys.argv[1] if len(sys.argv) > 1 else "solo"
    if mode == "dungeon":
        if "perfect" in sys.argv:
            game.combat.AUTO_CD, game.combat.AUTO_DELAY, game.combat.AUTO_REACT = 1.0, 0.0, 0.25
            print("-- perfect hands --")
        dungeon(int(sys.argv[2]) if len(sys.argv) > 2 else 20,
                int(sys.argv[3]) if len(sys.argv) > 3 and sys.argv[3].isdigit() else None)
        return
    seeds = int(sys.argv[2]) if len(sys.argv) > 2 else 3
    levels = {int(level) for level in sys.argv[3].split(",")} if len(sys.argv) > 3 and sys.argv[3] != "all" else None
    global POTIONS
    POTIONS = "potions" in sys.argv[4:]
    if mode == "bosses":
        global CHECKPOINTS
        CHECKPOINTS = [(level, tier, [f for f in fights if f in BOSSES]) for level, tier, fights in CHECKPOINTS]
        report({"miner+ranger": ["miner", "ranger"], "builder+homesteader": ["builder", "homesteader"],
                "homesteader+ranger": ["homesteader", "ranger"], "all four": ["builder", "homesteader", "miner", "ranger"]},
               seeds, levels)
    elif mode == "party":
        report({"builder+homesteader": ["builder", "homesteader"],
                "all four": ["builder", "homesteader", "miner", "ranger"]}, seeds, levels)
    else:
        report({cls: [cls] for cls in game.CLASS_STYLE}, seeds, levels)


if __name__ == "__main__":
    main()