#!/usr/bin/env python3
"""
Fix garbled OCR talent titles in the rules DB.
- Normalise mixed-case OCR titles to proper Title Case
- Set source/source_abbr/category correctly for all talent entries
- Remove duplicate sanitized entries where a clean 'talent_*' version exists
"""
import os, re, mysql.connector
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent

def load_env_file(path):
    if not path.exists():
        return
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith('#') or '=' not in line:
            continue
        key, value = line.split('=', 1)
        os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))

load_env_file(REPO / 'database' / '.env')
load_env_file(REPO / '.env')
if not os.environ.get('DB_PASSWORD'):
    raise RuntimeError('DB_PASSWORD must be set in database/.env or .env')

DB = dict(
    host=os.environ.get('DB_HOST', '192.168.1.113'),
    port=int(os.environ.get('DB_PORT', '3307')),
    user=os.environ.get('DB_USER', 'deathwatch'),
    password=os.environ['DB_PASSWORD'],
    database=os.environ.get('DB_NAME', 'deathwatch'),
)

def normalise_title(raw):
    """Convert OCR-garbled title to proper Title Case."""
    # Remove excess whitespace
    s = re.sub(r'\s+', ' ', raw).strip()
    # Lowercase everything first, then title-case
    words = s.lower().split()
    # Keep small words lowercase unless first word
    small = {'a','an','the','of','to','in','for','and','or','but','with','at','by','from'}
    result = []
    for i, w in enumerate(words):
        if i == 0 or w not in small:
            # Handle hyphenated words
            result.append('-'.join(p.capitalize() for p in w.split('-')))
        else:
            result.append(w)
    return ' '.join(result)

db = mysql.connector.connect(**DB)
c = db.cursor()

# 1. Get all sanitized talent entries (have "Prerequisites:" in content)
c.execute("""SELECT id, rule_id, title FROM rules
             WHERE source = 'sanitized' AND content LIKE 'Prerequisites:%'""")
sanitized_talents = c.fetchall()
print(f"Found {len(sanitized_talents)} sanitized talent entries")

# 2. Get existing clean talent rule_ids (from our earlier import)
c.execute("SELECT rule_id, title FROM rules WHERE source = 'Core Rulebook' AND category = 'talents'")
clean_talents = {row[0]: row[1] for row in c.fetchall()}
print(f"Found {len(clean_talents)} clean talent entries")

fixed = 0
deleted = 0

for (db_id, rule_id, title) in sanitized_talents:
    clean_title = normalise_title(title)
    expected_rid = 'talent_' + re.sub(r'[^a-z0-9]', '_', clean_title.lower())
    # Remove trailing underscores
    expected_rid = re.sub(r'_+$', '', expected_rid)
    expected_rid = re.sub(r'_+', '_', expected_rid)

    if expected_rid in clean_talents:
        # We already have a clean version - delete the sanitized duplicate
        c.execute("DELETE FROM rules WHERE id = %s", (db_id,))
        deleted += 1
    else:
        # No clean version exists - fix this entry's title/source/category
        c.execute("""UPDATE rules
                     SET title = %s, source = 'Core Rulebook', source_abbr = 'CR', category = 'talents'
                     WHERE id = %s""",
                  (clean_title, db_id))
        # Also add a proper rule_id if missing
        new_rid = expected_rid
        # Check if rule_id already conflicts
        c.execute("SELECT COUNT(*) FROM rules WHERE rule_id = %s AND id != %s", (new_rid, db_id))
        if c.fetchone()[0] == 0:
            c.execute("UPDATE rules SET rule_id = %s WHERE id = %s", (new_rid, db_id))
        fixed += 1

db.commit()
print(f"Fixed: {fixed} talent titles, Deleted: {deleted} duplicates")

# 3. Also fix remaining garbled-title talents (non-sanitized, left-over OCR artifacts)
c.execute("""SELECT id, title FROM rules
             WHERE category = 'talents' AND title REGEXP '[A-Z]{2,}[a-z]' """)
garbled = c.fetchall()
print(f"\nFixing {len(garbled)} remaining garbled-title talent entries...")
for (db_id, title) in garbled:
    clean = normalise_title(title)
    if clean != title:
        c.execute("UPDATE rules SET title = %s WHERE id = %s", (clean, db_id))
        print(f"  '{title}' → '{clean}'")

db.commit()

# 4. Final stats
c.execute("SELECT source, COUNT(*) FROM rules GROUP BY source ORDER BY COUNT(*) DESC")
print("\n=== Final rule counts by source ===")
total = 0
for src, cnt in c.fetchall():
    print(f"  {src}: {cnt}")
    total += cnt
print(f"  TOTAL: {total}")

c.execute("SELECT category, COUNT(*) FROM rules GROUP BY category ORDER BY COUNT(*) DESC")
print("\n=== By category ===")
for cat, cnt in c.fetchall():
    print(f"  {cat}: {cnt}")

c.close()
db.close()