#!/usr/bin/env python3
"""Fix DRAGON_SWORD (593) and DRAGON_AXE (594) weights on the 9 bosses so:
- Sword is ~1.5x more common than that NPC's DRAGON_MED_HELMET (combined rate)
- Axe is ~1.2x more common than DRAGON_MED_HELMET (slightly rarer than sword)
- For bosses with no DRAGON_MED_HELMET (RED/BLUE/Lessers/Greater), use the
  BLACK_DEMON/FIRE_GIANT d_med rate as the baseline.

Removes existing sword/axe entries and adds single new entries with updated weights.
Right half is NOT touched (already matched to left half).
"""
import 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

_, npc_id_to_name, _, _ = load_enums()

# Target 1/x rates per NPC for sword and axe
# (derived from d_med combined rate × 1.5 for sword, × 1.2 for axe)
TARGETS = {
    477: {              "axe": 200},   # KBD                 — apex; axe only, sword removed
    291: {"sword": 300,  "axe": 400},  # Black Dragon        — hard (anti-drag)
    290: {"sword": 350,  "axe": 450},  # Black Demon         — hard
    201: {"sword": 350,  "axe": 450},  # Red Dragon          — hard (anti-drag)
    344: {"sword": 500,  "axe": 650},  # Fire Giant          — mid
}

# NPCs whose existing dragon sword/axe rows should be stripped if present, even
# though they're no longer in TARGETS. Lets the script idempotently "demote"
# previously-eligible bosses.
LEGACY_DROP_NPCS = {
    202,  # Blue Dragon            — removed by user decision
    184,  # Greater Demon          — removed by user decision
    22,   # Lesser Demon           — removed by user decision
    181,  # Lesser Demon W/Maze    — removed by user decision
}

DRAGON_SWORD = 593
DRAGON_AXE = 594

header, rows, footer = parse_update_sql()

# Step 1: remove existing sword/axe entries for bosses in TARGETS or LEGACY list
strip_scope = set(TARGETS.keys()) | LEGACY_DROP_NPCS
removed = 0
new_rows = []
for r in rows:
    nid, amt, iid, w, idx = r
    if nid in strip_scope and iid in (DRAGON_SWORD, DRAGON_AXE):
        removed += 1
        continue
    new_rows.append(r)

# Step 2: compute new weights based on each NPC's CURRENT total (excluding the
# soon-to-be-added sword/axe weights themselves) — we need to solve:
#   sword_weight / (existing_total + sword_weight + axe_weight) = 1 / sword_odds
# This becomes coupled. Use iterative approximation: assume the contribution
# of new sword+axe is small, compute weights from existing_total. Re-check rate
# after adding. Two passes converges to <1% deviation.

# Per-NPC: get current total (without the removed sword/axe)
existing_total = defaultdict(int)
for nid, amt, iid, w, _ in new_rows:
    if nid in TARGETS:
        existing_total[nid] += w

# First-pass weights (sword/axe optional per-NPC; None means skip that weapon)
def compute_weights(nid, base_total):
    targets = TARGETS[nid]
    sword_w = max(1, round(base_total / targets["sword"])) if "sword" in targets else None
    axe_w   = max(1, round(base_total / targets["axe"]))   if "axe"   in targets else None
    return sword_w, axe_w

# Two-pass refinement
weights = {}
for nid in TARGETS:
    base = existing_total[nid]
    sword_w, axe_w = compute_weights(nid, base)
    full_total = base + (sword_w or 0) + (axe_w or 0)
    if sword_w is not None:
        sword_w = max(1, round(full_total / TARGETS[nid]["sword"]))
    if axe_w is not None:
        axe_w = max(1, round(full_total / TARGETS[nid]["axe"]))
    weights[nid] = (sword_w, axe_w)

# Step 3: add new rows (skip weapons that are None for this NPC)
added = []
for nid, (sword_w, axe_w) in weights.items():
    if sword_w is not None:
        new_rows.append((nid, '1', DRAGON_SWORD, sword_w, 0))
    if axe_w is not None:
        new_rows.append((nid, '1', DRAGON_AXE,   axe_w,   0))
    added.append((nid, sword_w, axe_w))

# Sort & write
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)

print(f"Removed {removed} old sword/axe entries.")
print(f"\nNew weights and resulting rates:")
print(f"{'NPC':<24} {'sword_w':>12} {'sword_rate':>14} {'axe_w':>12} {'axe_rate':>14}")
print("-" * 80)

# Recompute final rates
final_total = defaultdict(int)
for nid, amt, iid, w, _ in new_rows:
    if nid in TARGETS:
        final_total[nid] += w

for nid, sword_w, axe_w in sorted(added, key=lambda x: -(x[1] or 0)):
    name = npc_id_to_name.get(nid, f"#{nid}")
    tot = final_total[nid]
    s_str = f"{sword_w:>12,}" if sword_w is not None else f"{'—':>12}"
    a_str = f"{axe_w:>12,}"   if axe_w   is not None else f"{'—':>12}"
    s_rate = f"1/{round(tot/sword_w):,}" if sword_w else "—"
    a_rate = f"1/{round(tot/axe_w):,}"   if axe_w   else "—"
    print(f"{name:<24} {s_str} {s_rate:>14} {a_str} {a_rate:>14}")