import argparse, os, csv, json, struct, collections
from PIL import Image

Image.MAX_IMAGE_PIXELS = None


def tga_header(b):
    return {"w": struct.unpack_from("<H", b, 12)[0],
            "h": struct.unpack_from("<H", b, 14)[0],
            "bpp": b[16], "type": b[2]}


def dds_header(b):
    h = struct.unpack_from("<I", b, 12)[0]
    w = struct.unpack_from("<I", b, 16)[0]
    mip = struct.unpack_from("<I", b, 28)[0]
    four = b[84:88].decode("latin1").strip()
    return {"w": w, "h": h, "fourcc": four, "mip": mip}


def classify(name):
    low = name.lower().replace("\\", "/")
    flags = []
    if any(k in low for k in ("shrift", "debugfont")) or low.endswith("stats/digits.tga"):
        flags.append("font_atlas")
    if "overhead" in low or "mainmenu" in low:
        flags.append("ui_atlas")
    if any(k in low for k in ("menu", "credits", "tutorial", "game-name", "button-")):
        flags.append("text_baked")
    if low.startswith(("landscape/", "sky/", "waterfall")) or "cloud" in low:
        flags.append("tiling_likely")
    if "spells" in low or "/fx/" in low or low.startswith("fx/"):
        flags.append("fx")
    group = low.split("/")[0] if "/" in low else "(root)"
    return group, ";".join(flags)


def main():
    ap = argparse.ArgumentParser(description="Convert Dragon textures to PNG and build an upscale-planning report.")
    ap.add_argument("--csv", required=True, help="index csv produced by res_unpack.py")
    ap.add_argument("--indir", required=True, help="extracted original directory")
    ap.add_argument("--outdir", required=True, help="PNG output directory")
    ap.add_argument("--report", required=True, help="output report csv")
    args = ap.parse_args()

    rows = list(csv.DictReader(open(args.csv, newline="", encoding="utf-8")))
    out_rows = []
    groups = collections.Counter()
    flags_count = collections.Counter()
    fmt_count = collections.Counter()
    total_px = 0
    failed = []

    for i, r in enumerate(rows):
        name = r["name"]
        src = os.path.join(args.indir, r["file"])
        ext = os.path.splitext(name)[1].lower()
        rec = {"name": name, "ext": ext, "src": r["file"], "png": "", "orig_w": "", "orig_h": "",
               "mode": "", "has_alpha": "", "dds_fourcc": "", "dds_mips": "",
               "group": "", "flags": "", "orig_bytes": r["size"]}
        group, flags = classify(name)
        rec["group"] = group
        rec["flags"] = flags
        for f in flags.split(";"):
            if f:
                flags_count[f] += 1
        groups[group] += 1

        try:
            if ext in (".tga", ".dds"):
                raw = open(src, "rb").read()
                if ext == ".tga":
                    hi = tga_header(raw)
                    fmt_count["TGA%s/%dbpp" % (hi["type"], hi["bpp"])] += 1
                    rec["orig_w"], rec["orig_h"] = hi["w"], hi["h"]
                else:
                    hi = dds_header(raw)
                    fmt_count["DDS:" + hi["fourcc"]] += 1
                    rec["orig_w"], rec["orig_h"] = hi["w"], hi["h"]
                    rec["dds_fourcc"], rec["dds_mips"] = hi["fourcc"], hi["mip"]
                im = Image.open(src)
                im.load()
                if im.mode == "P":
                    im = im.convert("RGBA")
                rec["mode"] = im.mode
                alpha = im.mode in ("RGBA", "LA")
                if alpha:
                    a = im.getchannel("A") if im.mode == "RGBA" else im.getchannel("A")
                    alpha = a.getextrema()[0] < 255
                rec["has_alpha"] = int(bool(alpha))
                total_px += im.size[0] * im.size[1]
                png_path = os.path.join(args.outdir, os.path.splitext(r["file"])[0] + ".png")
                os.makedirs(os.path.dirname(png_path), exist_ok=True)
                im.save(png_path)
                rec["png"] = os.path.relpath(png_path, args.outdir)
            else:
                fmt_count["other"] += 1
                rec["flags"] = (flags + ";non_texture").strip(";")
        except Exception as e:
            failed.append((name, str(e)))
            rec["mode"] = "ERROR:" + str(e)[:60]
        out_rows.append(rec)
        if (i + 1) % 100 == 0:
            print("processed %d/%d" % (i + 1, len(rows)))

    with open(args.report, "w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(out_rows[0].keys()))
        w.writeheader()
        w.writerows(out_rows)

    summary = {
        "total": len(out_rows),
        "formats": dict(fmt_count),
        "groups": dict(groups.most_common()),
        "flags": dict(flags_count),
        "total_megapixels": round(total_px / 1e6, 2),
        "failed": failed,
    }
    print(json.dumps(summary, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()