"""Extract every recoverable scene on the Glaze drive to one master JSON."""
import os, re, json, hashlib, traceback
from collections import defaultdict
import divformats as df

SDA4 = r"c:\VWE\TeslaRel410\sda4"
MIN_COVERAGE = 0.6

# ---------------------------------------------------------------- file index
# sda4 dev drive + the release repo's CONTENT (same-era Division files)
ROOTS = [SDA4, r"c:\VWE\TeslaRel410\CONTENT"]
index = defaultdict(list)
for r in ROOTS:
    for root, dirs, files in os.walk(r):
        dirs[:] = [d for d in dirs if d.upper() not in ("EDITBACK", "RECYCLED")]
        for f in files:
            base, ext = os.path.splitext(f.lower())
            index[(base, ext)].append(os.path.join(root, f))

def nearest(cands, refdir):
    """Pick candidate path sharing the longest prefix with refdir."""
    def score(p):
        a, b = p.lower(), refdir.lower()
        n = 0
        while n < min(len(a), len(b)) and a[n] == b[n]: n += 1
        return n
    return max(cands, key=score)

def find(base, exts, refdir):
    for ext in exts:
        c = index.get((base.lower(), ext))
        if c: return nearest(c, refdir)
    return None

# ---------------------------------------------------------------- caches
model_cache = {}    # abspath -> model id
models = {}         # id -> groups
image_cache = {}    # abspath -> image id
images = {}         # id -> datauri
mat_file_cache = {} # abspath -> (textures, materials, ramps)

def load_model(path):
    if path in model_cache: return model_cache[path]
    mid = "m%d" % len(models)
    if path.lower().endswith(".vgf"):
        groups = df.parse_vgf(path)
        extra = ({}, {}, {})
    else:
        groups, m, t, r = df.parse_bgf(path)
        extra = (m, t, r)
    models[mid] = groups
    model_cache[path] = (mid, extra)
    return model_cache[path]

def load_image(path):
    if path in image_cache: return image_cache[path]
    try:
        uri = df.tga_datauri(path) if path.lower().endswith(".tga") else df.svt_datauri(path)
    except Exception:
        return None
    iid = "i%d" % len(images)
    images[iid] = uri
    image_cache[path] = iid
    return iid

def load_mat_file(path):
    if path in mat_file_cache: return mat_file_cache[path]
    if path.lower().endswith(".vmf"):
        t, m, r = df.parse_vmf(path, "x")
        m = {k.split(":", 1)[1]: v for k, v in m.items()}
    else:
        g, m, t, r = df.parse_bgf(path)
    mat_file_cache[path] = (t, m, r)
    return mat_file_cache[path]

def resolve_texture_image(tex, refdir):
    """tex: {map, slice} -> image id or None"""
    if not tex or not tex.get("map"): return None
    m, sl = tex["map"], tex.get("slice")
    tries = []
    if sl is not None: tries.append(("%s%d" % (m, sl + 1), (".tga",)))
    dm = re.match(r"(.+?)(\d)$", m)
    tries.append((m, (".tga", ".svt")))
    if dm: tries.append((dm.group(1) + dm.group(2), (".tga", ".svt")))
    for base, exts in tries:
        p = find(base, exts, refdir)
        if p: return load_image(p)
    return None

def resolve_material(ref, refdir, embedded):
    """ref like 'file:name' or 'name'. Returns resolved material dict or None."""
    if not ref: return None
    if ":" in ref: fbase, name = ref.split(":", 1)
    else: fbase, name = None, ref
    # 1) material file <fbase>.vmf/.bmf
    if fbase:
        p = find(fbase, (".vmf", ".bmf"), refdir)
        if p:
            t, m, r = load_mat_file(p)
            if name in m:
                return finish_material(m[name], t, r, refdir)
    # 2) embedded in the BGF itself
    em, et, er = embedded
    if name in em:
        return finish_material(em[name], et, er, refdir)
    if fbase:  # sometimes embedded keyed with prefix
        if ref in em: return finish_material(em[ref], et, er, refdir)
    return None

def finish_material(m, texdict, rampdict, refdir):
    out = {"ambient": m.get("ambient"), "diffuse": m.get("diffuse"),
           "emissive": m.get("emissive"), "immune": m.get("immune", False)}
    tex = m.get("texture")
    if tex:
        tinfo = texdict.get(tex)
        if tinfo is None and re.match(r".*_tex$", tex):
            tinfo = {"map": tex[:-4], "slice": None}
        # heuristic: trailing digit in texture name selects TGA slice
        if tinfo and tinfo.get("slice") is None:
            dm = re.match(r"(.+?)(\d)_?tex$", tex)
            if dm and tinfo.get("map") and not re.search(r"\d$", tinfo["map"]):
                tinfo = {"map": tinfo["map"], "slice": int(dm.group(2)) - 1}
        out["img"] = resolve_texture_image(tinfo, refdir)
    rname = m.get("ramp")
    if rname and rname in rampdict:
        out["ramp"] = rampdict[rname]
    return out

# ---------------------------------------------------------------- scenes
def dos_to_real(dospath, scndir):
    """Map a scene's GEOMETRY/MATERIAL dir (DOS path) onto the dump."""
    if not dospath: return scndir
    p = dospath.replace("\\", "/").strip()
    p = re.sub(r"^[A-Za-z]:", "", p).lstrip("/")
    cand = os.path.join(SDA4, *p.split("/"))
    if os.path.isdir(cand): return cand
    return scndir

seen_hashes = set()
scenes_out = {}
skipped = []

PRIORITY = ["STDAVE", "HPDAVE", "CANYON", "BTDAVE", "BTRAVINE", "RPDAVE", "ERIC",
            "FX", "RD", "CONVERT", "DPL3RLS", "DPL3", "PROBLEMS", "SPHERES"]
def prio(rel):
    top = rel.split(os.sep)[0].upper()
    return PRIORITY.index(top) if top in PRIORITY else 99

all_scn = []
for root, dirs, files in os.walk(SDA4):
    dirs[:] = [d for d in dirs if d.upper() not in ("EDITBACK", "RECYCLED")]
    for f in files:
        if f.lower().endswith(".scn"):
            all_scn.append(os.path.join(root, f))
all_scn.sort(key=lambda p: (prio(os.path.relpath(p, SDA4)), p))

for scn_path in all_scn:
    rel = os.path.relpath(scn_path, SDA4)
    try:
        content = open(scn_path, "rb").read()
        h = hashlib.md5(content).hexdigest()
        if h in seen_hashes:
            skipped.append((rel, "duplicate")); continue
        scene = df.parse_scn(scn_path)
        pls = scene["placements"]
        if not pls:
            skipped.append((rel, "no placements")); continue
        scndir = os.path.dirname(scn_path)
        geodir = dos_to_real(scene.get("geometry"), scndir)

        needed = []
        for p in pls:
            needed.append(p["geo"].split(":")[0])
            for c in p["children"]: needed.append(c["geo"].split(":")[0])
        uniq = sorted(set(needed))
        found_models, missing = {}, []
        for g in uniq:
            fp = find(g, (".vgf",), geodir) or find(g, (".bgf",), geodir)
            if fp: found_models[g] = fp
            else: missing.append(g)
        if len(found_models) == 0 or len(found_models) / len(uniq) < MIN_COVERAGE:
            skipped.append((rel, "coverage %d/%d" % (len(found_models), len(uniq))))
            continue

        # load models + materials
        used_models, used_mats = {}, {}
        for g, fp in found_models.items():
            mid, embedded = load_model(fp)
            used_models[g] = mid
            for grp in models[mid]:
                mref = grp.get("material")
                if mref and mref not in used_mats:
                    r = resolve_material(mref, os.path.dirname(fp), embedded)
                    if r is None:
                        r = resolve_material(mref, dos_to_real(scene.get("material"), scndir), embedded)
                    used_mats[mref] = r or {}
        # splines
        spls = {}
        for p in pls:
            if p["kind"] == "dynamic":
                sp = find(os.path.splitext(p["spl"])[0], (".spl",), scndir)
                if sp: spls[p["spl"]] = df.parse_spl(sp)
        placements = [p for p in pls if p["geo"].split(":")[0] in used_models
                      and (p["kind"] != "dynamic" or p["spl"] in spls)]
        for p in placements:
            p["children"] = [c for c in p["children"] if c["geo"].split(":")[0] in used_models]

        scenes_out[rel.replace(os.sep, "/")] = {
            "fov": scene.get("fov", 60), "clip": scene.get("clip", [1, 5000]),
            "fog": scene.get("fog"), "bg": scene.get("bg", [0, 0, 0]),
            "ambient": scene.get("ambient", [0.2, 0.2, 0.2]),
            "lights": scene.get("lights", []), "start": scene.get("start"),
            "placements": placements, "modelmap": used_models,
            "materials": used_mats, "splines": spls, "missing": missing,
        }
        seen_hashes.add(h)
        nv = sum(len(g.get("v", g.get("points", []))) for m in used_models.values() for g in models[m])
        print("OK  %-46s models=%d/%d verts=%d" % (rel, len(found_models), len(uniq), nv))
    except Exception as e:
        skipped.append((rel, "ERROR " + str(e)))
        print("ERR %-46s %s" % (rel, e))
        traceback.print_exc()

out = {"models": models, "images": images, "scenes": scenes_out}
dst = os.path.join(os.path.dirname(os.path.abspath(__file__)), "allscenes.json")
with open(dst, "w") as fh:
    json.dump(out, fh, separators=(",", ":"))
print("\n%d scenes extracted, %d skipped" % (len(scenes_out), len(skipped)))
print("models=%d images=%d  size=%.1f MB" % (len(models), len(images), os.path.getsize(dst) / 1e6))
with open("extract_report.txt", "w") as fh:
    for rel, why in skipped: fh.write("%s\t%s\n" % (rel, why))