#!/usr/bin/env python3
"""Add DRAGON_SWORD (593) and DRAGON_AXE (594) drops to 5 boss NPCs at
specific 1/N rates. Weights are calculated as total_w_old/(N-1) so the
resulting probability is exactly 1/N after the rows are added.
"""
import shutil, sys, os
sys.path.insert(0, '/Users/tomassimkus/VS/openrsc-develop')
os.chdir('/Users/tomassimkus/VS/openrsc-develop')
from update_drops import parse_update_sql, write_update_sql
from compare_drops import load_enums
from collections import defaultdict

name_to_npc_id, npc_id_to_name, _, item_id_to_name = load_enums()

DRAGON_SWORD = 593  # "Dragon Longsword"
DRAGON_AXE   = 594

# (npc_id, item_id, target_1/N)
TARGETS = [
    (477, DRAGON_AXE,   1500),   # KING_BLACK_DRAGON
    (291, DRAGON_SWORD, 5000),   # BLACK_DRAGON
    (291, DRAGON_AXE,   6000),
    (201, DRAGON_SWORD, 10000),  # RED_DRAGON
    (201, DRAGON_AXE,   12000),
    (290, DRAGON_SWORD, 10000),  # BLACK_DEMON
    (290, DRAGON_AXE,   12000),
    (344, DRAGON_SWORD, 11500),  # FIRE_GIANT
    (344, DRAGON_AXE,   13500),
]

shutil.copy("pk_npcdrops.sql", "pk_npcdrops_update.sql")
header, rows, footer = parse_update_sql()

# Strip any pre-existing DRAGON_SWORD / DRAGON_AXE rows on these NPCs
target_nids = {nid for nid, _, _ in TARGETS}
new_rows = [r for r in rows
            if not (r[0] in target_nids and r[2] in (DRAGON_SWORD, DRAGON_AXE))]
removed = len(rows) - len(new_rows)

# Per-NPC totals (re-computed from cleaned rows so the math is correct
# whether or not stale entries existed)
totals = defaultdict(int)
for r in new_rows:
    if r[3] > 0:
        totals[r[0]] += r[3]

added = []
for nid, iid, target_N in TARGETS:
    tw = totals[nid]
    w_new = max(1, round(tw / (target_N - 1)))
    # parse_update_sql gives tuples (npc, amount_str, item, weight, 0)
    new_rows.append((nid, "1", iid, w_new, 0))
    totals[nid] += w_new
    actual_N = round(totals[nid] / w_new)
    added.append((nid, iid, target_N, w_new, actual_N))

new_rows.sort(key=lambda r: (r[0], 0 if r[3] == 0 else 1, r[3], r[2]))
write_update_sql(header, new_rows, footer)
shutil.copy("pk_npcdrops_update.sql", "pk_npcdrops.sql")
shutil.copy("pk_npcdrops_update.sql", "pk_npcdrops_p2p_update.sql")

print(f"Removed {removed} pre-existing dragon-weapon rows on target NPCs.")
print(f"Added {len(added)} new dragon-weapon rows:\n")
print(f"  {'NPC':<22} {'Item':<16} {'Weight':>10} {'Target':>10} {'Actual':>10}")
print("  " + "-" * 72)
for nid, iid, target_N, w, actual_N in added:
    nname = npc_id_to_name.get(nid, f"#{nid}")
    iname = item_id_to_name.get(iid, f"#{iid}")
    print(f"  {nname:<22} {iname:<16} {w:>10,} {'1/'+f'{target_N:,}':>10} {'1/'+f'{actual_N:,}':>10}")