#!/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] [seeds] [levels, e.g. 22,52] [potions]

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 main() -> None:
    mode = sys.argv[1] if len(sys.argv) > 1 else "solo"
    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()