#!/usr/bin/env python3
"""Build the Brandling roster from the web-design client export.

  python tools/build_roster.py --csv ../clickup_web_design_export.csv [--limit N] [--jobs 16]

Pipeline (every network step is cached in art_src/creatures/cache/, so reruns only do new work):

  1. tools/roster_csv.py reads the CSV: company name, domain, staging-site URL, services and key
     highlights, colour words of the logo note. Nothing else leaves the CSV. Non-companies are
     dropped and duplicates (same normalised name) merged.
  2. Logos: Company_Sprite_Gen/logo_fetch.py looks on the domain's homepage, then on the
     staging site, then site icons and favicon services. cache/logos/<domain>.png + .json.
     A logo shared by 2+ companies is a template placeholder (e.g. the agency's own logo on an
     unfinished staging site) and is ignored.
  3. Types: Company_Sprite_Gen/business_pokemon_type.py (OpenRouter jev classifier) on
     "name + services + highlights". cache/types.json keeps choice + probabilities per slug.
  4. Sprites: Company_Sprite_Gen/generate_sprite.generate(..., write_files=False) with the
     logo's colours, else the colours named in the logo note, else a per-creature shade of
     the type's colour (type_variant_colors).
     cache/render.json skips sprites whose inputs didn't change (--rerender redoes all).
  5. Writes assets/creatures/roster.json, sprites/<slug>.png, contact_sheet.png and
     art_src/creatures/report.json.

Species ids are permanent (players' saves store them; the game requires id == index):
the first run (no roster.json yet) numbers species in slug order. Every later run keeps each
slug's id, appends new slugs after the last id (in slug order) and keeps companies that left
the CSV in the roster with "retired": true, their sprite kept, so no id is ever reused; a
retired company that comes back gets its old id. A --limit run leaves the species it didn't
sample untouched instead of retiring them. Every entry has "retired" (false for live ones).

OPENROUTER_API_KEY comes from the environment or ../.env (next to this repo).
"""
import argparse
import colorsys
import hashlib
import json
import os
import random
import re
import sys
import threading
import time
from collections import Counter
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from pathlib import Path

TOOLS = Path(__file__).resolve().parent
ROOT = TOOLS.parent
CSG = Path(os.environ.get("COMPANY_SPRITE_GEN", ROOT.parent / "Company_Sprite_Gen")).resolve()
ENV_FILE = ROOT.parent / ".env"
SRC = ROOT / "art_src" / "creatures"
CACHE = SRC / "cache"
LOGOS = CACHE / "logos"
TYPES_CACHE = CACHE / "types.json"
RENDER_CACHE = CACHE / "render.json"
REPORT = SRC / "report.json"
OUT = ROOT / "assets" / "creatures"
SPRITES = OUT / "sprites"

LIBRARY = "pixellab"
SPRITE_SIZE = 64
SHARED_LOGO_LIMIT = 2          # the same logo image on this many companies = template placeholder
TYPE2_MIN = 0.25
RARITY_CUTS = (("common", 0.60), ("uncommon", 0.87), ("rare", 0.97), ("legendary", 1.0))
RARITY_RANK = {r: i for i, (r, _) in enumerate(RARITY_CUTS)}
RATES = {"common": (0.55, 0.08), "uncommon": (0.40, 0.12), "rare": (0.25, 0.18), "legendary": (0.12, 0.25)}
BOTH_HABITAT_TYPES = {"ice", "flying", "poison"}
BOTH_HABITAT_CHANCE = 0.4
CONTACT_COUNT = 300


def log(msg):
    print(msg, flush=True)


def setup_paths():
    """Make Company_Sprite_Gen importable. Its part registries resolve their art folders
    relative to the working directory, so the process works from there (all our paths are
    absolute)."""
    for p in (str(TOOLS), str(CSG)):
        if p not in sys.path:
            sys.path.insert(0, p)
    os.chdir(CSG)


def load_env():
    if os.environ.get("OPENROUTER_API_KEY") or not ENV_FILE.exists():
        return
    for line in ENV_FILE.read_text().splitlines():
        key, sep, value = line.strip().partition("=")
        if sep and key.strip() == "OPENROUTER_API_KEY":
            os.environ["OPENROUTER_API_KEY"] = value.strip().strip("'\"")


def unit(slug, salt):
    """Deterministic number in [0, 1) from the slug."""
    return int(hashlib.sha256(f"{salt}:{slug}".encode()).hexdigest()[:13], 16) / 16 ** 13


def read_json(path, default):
    try:
        return json.loads(Path(path).read_text())
    except (OSError, ValueError):
        return default


def write_json(path, data, indent=1):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(json.dumps(data, indent=indent, sort_keys=True) + "\n")
    tmp.replace(path)


# ============================================================ logos

def fetch_logos(companies, jobs):
    """slug -> (Path or None, method or None) for every company."""
    import logo_fetch
    out = {c.slug: (None, None) for c in companies}
    todo = [c for c in companies if c.domain or c.staging]
    t0 = time.time()
    with ThreadPoolExecutor(jobs) as ex:
        futs = {ex.submit(logo_fetch.fetch_logo_info, c.domain, LOGOS, [c.staging] if c.staging else [],
                          None if c.domain else f"name-{c.slug}"): c for c in todo}
        for i, f in enumerate(as_completed(futs), 1):
            c = futs[f]
            try:
                out[c.slug] = f.result()
            except Exception as e:           # one broken site must not stop the run
                log(f"  ! logo {c.slug}: {type(e).__name__}: {e}")
            if i % 100 == 0 or i == len(todo):
                log(f"  logos {i}/{len(todo)} ({time.time() - t0:.0f}s)")
    return out


def logo_digest(path):
    """Fingerprint that survives re-encoding / resizing (catches one template on many sites)."""
    from PIL import Image
    img = Image.open(path).convert("RGBA").resize((16, 16), Image.BILINEAR)
    return hashlib.sha1(bytes(v >> 4 for v in img.tobytes())).hexdigest()


def drop_shared_logos(logos):
    digests = {slug: logo_digest(p) for slug, (p, _) in logos.items() if p}
    counts = Counter(digests.values())
    shared = {slug for slug, d in digests.items() if counts[d] >= SHARED_LOGO_LIMIT}
    for slug in shared:
        logos[slug] = (None, None)
    return shared, digests


# ============================================================ types

def description(c):
    d = f"{c.name}."
    if c.services:
        d += " Services: " + ", ".join(c.services[:10]) + "."
    if c.highlights:
        d += " " + "; ".join(c.highlights[:3]) + "."
    return d[:600]


def classify_retry(get_type, desc, attempts=6):
    import requests
    for i in range(attempts):
        try:
            ans = get_type(desc)
            if ans.get("choice"):
                return ans
        except requests.HTTPError as e:
            code = e.response.status_code if e.response is not None else 0
            if code in (400, 401, 402, 403):
                raise
        except (requests.RequestException, ValueError, KeyError):
            pass
        time.sleep(min(30, 2 ** i) + random.random())
    return None


def classify_all(companies, jobs):
    from sprite_parts import TYPES
    cache = read_json(TYPES_CACHE, {})
    todo = [c for c in companies if cache.get(c.slug, {}).get("desc") != description(c)]
    if todo:
        load_env()
        if not os.environ.get("OPENROUTER_API_KEY"):
            sys.exit(f"OPENROUTER_API_KEY is not set (environment or {ENV_FILE})")
        from business_pokemon_type import get_business_type
        lock = threading.Lock()
        t0 = time.time()
        with ThreadPoolExecutor(jobs) as ex:
            futs = {ex.submit(classify_retry, get_business_type, description(c)): c for c in todo}
            for i, f in enumerate(as_completed(futs), 1):
                c = futs[f]
                ans = f.result()
                if ans and ans["choice"].lower() in TYPES:
                    with lock:
                        cache[c.slug] = {"desc": description(c), "choice": ans["choice"].lower(),
                                         "probabilities": {k.lower(): v for k, v in ans.get("probabilities", {}).items()}}
                else:
                    log(f"  ! no type for {c.slug} (will retry next run)")
                if i % 100 == 0 or i == len(todo):
                    with lock:
                        write_json(TYPES_CACHE, cache)
                    log(f"  types {i}/{len(todo)} ({time.time() - t0:.0f}s)")
    return cache


# ============================================================ sprites

_lib = None


def _worker_init():
    setup_paths()


def render_one(name, ptype, probs, logo, colors):
    """(info, png bytes, anchor row) for one creature. Runs in a worker process."""
    global _lib
    import io
    import numpy as np
    import generate_sprite as gs
    if _lib is None:
        _lib = gs.load_library(LIBRARY)
    info, sprite = gs.generate(_lib, name, None, ptype, probs, logo=logo, colors=colors, write_files=False)
    if sprite.size != (SPRITE_SIZE, SPRITE_SIZE):
        raise ValueError(f"{name}: sprite is {sprite.size}, expected {SPRITE_SIZE}x{SPRITE_SIZE}")
    rows = np.nonzero(np.array(sprite)[:, :, 3].any(axis=1))[0]
    buf = io.BytesIO()
    sprite.save(buf, "PNG", optimize=True)
    keep = {"plan": info["plan"], "attributes": info["attributes"],
            "colors": {"main": info["colors"]["main"].lower(), "accent": info["colors"]["accent"].lower()}}
    return keep, buf.getvalue(), int(rows[-1]) if len(rows) else SPRITE_SIZE - 1


def render_all(jobs_in, jobs, rerender):
    """jobs_in: {slug: (name, ptype, probs, logo_path, colors, digest)} -> {slug: render entry}."""
    cache = {} if rerender else read_json(RENDER_CACHE, {})
    SPRITES.mkdir(parents=True, exist_ok=True)
    out, todo = {}, []
    for slug, (name, ptype, probs, logo, colors, digest) in jobs_in.items():
        key = hashlib.sha1(json.dumps([LIBRARY, name, ptype, digest, colors]).encode()).hexdigest()
        hit = cache.get(slug)
        if hit and hit.get("key") == key and (SPRITES / f"{slug}.png").exists():
            out[slug] = hit
        else:
            todo.append((slug, key, name, ptype, probs, logo, colors))
    if todo:
        t0 = time.time()
        with ProcessPoolExecutor(jobs, initializer=_worker_init) as ex:
            futs = {ex.submit(render_one, name, ptype, probs, str(logo) if logo else None, colors): (slug, key)
                    for slug, key, name, ptype, probs, logo, colors in todo}
            for i, f in enumerate(as_completed(futs), 1):
                slug, key = futs[f]
                info, png, anchor = f.result()
                (SPRITES / f"{slug}.png").write_bytes(png)
                out[slug] = cache[slug] = {"key": key, "anchorRow": anchor, **info}
                if i % 200 == 0 or i == len(todo):
                    log(f"  sprites {i}/{len(todo)} ({time.time() - t0:.0f}s)")
        write_json(RENDER_CACHE, {s: e for s, e in cache.items() if s in jobs_in})
    return out


# ============================================================ roster

def type_variant_colors(slug, ptype):
    """Colours for a creature with no logo and no colour note: its type colour, shifted per
    slug in hue / saturation / lightness so the many logo-less creatures of one type don't
    all share one palette."""
    from sprite_parts import TYPE_COLORS
    h = TYPE_COLORS[ptype].lstrip("#")
    r, g, b = (int(h[i:i + 2], 16) / 255 for i in (0, 2, 4))
    hue, light, sat = colorsys.rgb_to_hls(r, g, b)
    hue = (hue + (unit(slug, "hue") - 0.5) * 0.16) % 1.0
    sat = min(max(sat * (0.75 + unit(slug, "sat") * 0.6), 0.25), 0.95)
    light = min(max(light + (unit(slug, "light") - 0.5) * 0.24, 0.25), 0.72)
    rgb = colorsys.hls_to_rgb(hue, light, sat)
    return ["#" + "".join(f"{round(v * 255):02x}" for v in rgb)]


def rarity_for(slug, ptype):
    u = unit(slug, "rarity")
    rarity = next(r for r, cut in RARITY_CUTS if u < cut)
    if ptype == "dragon" and RARITY_RANK[rarity] < RARITY_RANK["rare"]:
        rarity = "rare"
    return rarity


def rates_for(slug, rarity):
    catch, flee = RATES[rarity]
    catch += (unit(slug, "catch") - 0.5) * 0.06
    flee += (unit(slug, "flee") - 0.5) * 0.04
    return round(min(max(catch, 0.0), 1.0), 3), round(min(max(flee, 0.0), 1.0), 3)


def habitat_for(slug, ptype):
    if ptype == "water":
        return "water"
    if ptype in BOTH_HABITAT_TYPES and unit(slug, "habitat") < BOTH_HABITAT_CHANCE:
        return "both"
    return "grass"


def second_type(ptype, probs):
    ranked = sorted(((p, t) for t, p in (probs or {}).items() if t != ptype), reverse=True)
    return ranked[0][1] if ranked and ranked[0][0] >= TYPE2_MIN else None


def assign_ids(current, previous, partial=False):
    """Species ids are stored in players' saves, so they never change once assigned.

    current: {slug: entry without id} for this run; previous: the species list of the existing
    roster.json ([] on a first run). Every slug keeps its old id; a slug that is gone stays
    in the list as its old entry with "retired": true (ids stay contiguous, id == index, and
    a retired slug that comes back gets its id back); new slugs get the next ids in slug order.
    With partial=True (a --limit run) slugs missing from this run keep their old entry as is.
    Returns (species, new id count, retired count)."""
    species = []
    for i, old in enumerate(sorted(previous, key=lambda s: s["id"])):
        if old["id"] != i:
            sys.exit(f"roster.json is corrupt: species {old.get('slug')} has id {old['id']} at index {i}")
        entry = current.get(old["slug"])
        if entry:
            species.append({"id": i, **entry, "retired": False})
        else:
            species.append({**old, "retired": bool(old.get("retired")) if partial else True})
    known = {s["slug"] for s in species}
    new = sorted(slug for slug in current if slug not in known)
    for slug in new:
        species.append({"id": len(species), **current[slug], "retired": False})
    return species, len(new), sum(1 for s in species if s["retired"])


def write_roster(species):
    lines = ",\n".join("  " + json.dumps(s, separators=(", ", ": ")) for s in species)
    text = f'{{"version": 1, "spriteSize": {SPRITE_SIZE}, "species": [\n{lines}\n]}}\n'
    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / "roster.json").write_text(text)


def contact_sheet(species, path, count=CONTACT_COUNT, cols=20, scale=2):
    from PIL import Image, ImageDraw, ImageFont
    pick = random.Random(7).sample(species, min(count, len(species)))
    cell_w, cell_h, label = SPRITE_SIZE * scale + 8, SPRITE_SIZE * scale + 30, 28
    rows = (len(pick) + cols - 1) // cols
    sheet = Image.new("RGB", (cols * cell_w + 8, rows * cell_h + 8), (58, 66, 62))
    d = ImageDraw.Draw(sheet)
    font = ImageFont.load_default(size=11)
    for i, s in enumerate(pick):
        x, y = 8 + (i % cols) * cell_w, 8 + (i // cols) * cell_h
        d.rectangle((x - 2, y - 2, x + SPRITE_SIZE * scale + 1, y + SPRITE_SIZE * scale + 1),
                    fill=(78, 90, 84), outline=(222, 200, 120) if s["logo"] else (78, 90, 84))
        spr = Image.open(OUT / s["sprite"]).convert("RGBA")
        spr = spr.resize((SPRITE_SIZE * scale, SPRITE_SIZE * scale), Image.NEAREST)
        sheet.paste(spr, (x, y), spr)
        name = s["name"] if len(s["name"]) <= 22 else s["name"][:21] + "…"
        d.text((x, y + SPRITE_SIZE * scale + 2), name, fill=(240, 240, 232), font=font)
        d.text((x, y + SPRITE_SIZE * scale + 2 + label // 2), f"{s['type']}{'/' + s['type2'] if s['type2'] else ''}"
               f" {s['rarity'][0].upper()}", fill=(180, 190, 180), font=font)
    sheet.save(path)


def main():
    ap = argparse.ArgumentParser(description="Build assets/creatures/ from the client CSV")
    ap.add_argument("--csv", required=True, help="ClickUp web-design export")
    ap.add_argument("--limit", type=int, help="Only this many companies (a fixed random sample, for testing)")
    ap.add_argument("--jobs", type=int, default=16, help="Parallel logo downloads and sprite renders")
    ap.add_argument("--type-jobs", type=int, default=8, help="Parallel classifier calls")
    ap.add_argument("--rerender", action="store_true", help="Render every sprite again")
    a = ap.parse_args()
    csv_path = Path(a.csv).resolve()
    setup_paths()
    import roster_csv

    log(f"reading {csv_path.name}")
    companies, stats = roster_csv.load_companies(csv_path)
    if a.limit:
        companies = sorted(random.Random(0).sample(companies, min(a.limit, len(companies))), key=lambda c: c.slug)
    log(f"  {stats['rows']} rows -> {stats['kept']} companies "
        f"({stats['not_company']} not companies, {stats['duplicates']} duplicates); using {len(companies)}")

    log("logos")
    logos = fetch_logos(companies, a.jobs)
    shared, digests = drop_shared_logos(logos)

    log("types")
    types = classify_all(companies, a.type_jobs)
    companies = [c for c in companies if c.slug in types]

    log("sprites")
    jobs_in = {}
    for c in companies:
        logo = logos[c.slug][0]
        colors = None if logo else (c.colors or type_variant_colors(c.slug, types[c.slug]["choice"]))
        jobs_in[c.slug] = (c.name, types[c.slug]["choice"], types[c.slug]["probabilities"], logo, colors,
                           digests.get(c.slug) if logo else None)
    renders = render_all(jobs_in, a.jobs, a.rerender)

    current = {}
    for c in companies:
        t = types[c.slug]
        r = renders[c.slug]
        rarity = rarity_for(c.slug, t["choice"])
        catch, flee = rates_for(c.slug, rarity)
        current[c.slug] = {
            "slug": c.slug, "name": c.name, "domain": c.domain,
            "type": t["choice"], "type2": second_type(t["choice"], t["probabilities"]),
            "rarity": rarity, "catchRate": catch, "fleeRate": flee, "habitat": habitat_for(c.slug, t["choice"]),
            "colors": r["colors"], "sprite": f"sprites/{c.slug}.png", "anchorRow": r["anchorRow"],
            "body": r["plan"], "attributes": r["attributes"], "logo": logos[c.slug][0] is not None,
        }
    species, new_ids, retired = assign_ids(current, read_json(OUT / "roster.json", {}).get("species", []),
                                           partial=bool(a.limit))
    write_roster(species)
    live = {Path(s["sprite"]).name for s in species}
    stale = [p for p in SPRITES.glob("*.png") if p.name not in live]
    for p in stale:
        p.unlink()
    active = [s for s in species if not s["retired"]]
    contact_sheet(active, OUT / "contact_sheet.png")

    methods = Counter("no-domain" if not (c.domain or c.staging) else "shared-template" if c.slug in shared
                      else (logos[c.slug][1] or "not-found") for c in companies)
    report = {
        "rows": stats["rows"], "notCompanies": stats["not_company"], "duplicates": stats["duplicates"],
        "companies": stats["kept"], "kept": len(active), "newIds": new_ids, "retired": retired,
        "withDomain": sum(1 for c in companies if c.domain),
        "withStagingSite": sum(1 for c in companies if c.staging),
        "logoSearched": sum(1 for c in companies if c.domain or c.staging),
        "logosFound": sum(1 for s in active if s["logo"]),
        "logosByMethod": dict(methods.most_common()),
        "logosRejectedShared": len(shared),
        "colorsFromLogoNote": sum(1 for c in companies if not logos[c.slug][0] and c.colors),
        "colorsFromType": sum(1 for c in companies if not logos[c.slug][0] and not c.colors),
        "types": dict(Counter(s["type"] for s in active).most_common()),
        "type2": dict(Counter(s["type2"] for s in active if s["type2"]).most_common()),
        "rarity": dict(Counter(s["rarity"] for s in active).most_common()),
        "habitat": dict(Counter(s["habitat"] for s in active).most_common()),
        "bodies": dict(Counter(s["body"] for s in active).most_common()),
    }
    write_json(REPORT, report, indent=2)
    log(f"{len(species)} species ({len(active)} active, {retired} retired, {new_ids} new ids) -> "
        f"{OUT / 'roster.json'} ({len(stale)} stale sprites removed)")
    log(json.dumps({k: report[k] for k in ("logosFound", "logosByMethod", "types")}))


if __name__ == "__main__":
    main()