#!/usr/bin/env python3
"""Export EXE-proximity enemy sprite candidates for battle review scenes.

The mapping is deliberately non-promoting: it connects battle background
resource descriptors to nearby enemy/object sprite resource descriptors in the
EXE address space, so the browser can render an original sprite candidate
instead of a placeholder. It does not prove original enemy rows, formations,
stats, rewards, formulas, or battle-entry execution.
"""
from __future__ import annotations

import argparse
import html
import json
from pathlib import Path
from typing import Any

from decode_cns import decompress_cns, parse_image
from probe_exe_scene_tables import read_sections
from summarize_original_battle_resource_descriptors import build_summary as build_resource_descriptor_summary


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "out"
EXTRACT_FLD = ROOT / "extract_fld"
EXE = ROOT / "Hwanse2.exe"


def load_json(path: Path, fallback: Any) -> Any:
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except FileNotFoundError:
        return fallback


def ref_rows(resource_descriptors: dict, cls: str) -> list[dict]:
    return [
        row for row in resource_descriptors.get("resourceDescriptorRows") or []
        if row.get("class") == cls and isinstance(row.get("refVa"), int)
    ]


def cns_browser_path(name: str) -> str:
    src = EXTRACT_FLD / name
    if not src.exists():
        return ""
    return f"../extract_fld/{Path(name).stem}.cns"


def nearest_enemy(background_ref: dict, enemy_refs: list[dict]) -> dict | None:
    if not enemy_refs:
        return None
    ref_va = int(background_ref["refVa"])
    return min(
        enemy_refs,
        key=lambda row: (
            abs(int(row["refVa"]) - ref_va),
            str(row.get("name") or ""),
            int(row["refVa"]),
        ),
    )


def unpack_palette_indices(width: int, height: int, pixels: bytes, bpp: int) -> list[list[int]]:
    stride = ((width * bpp + 31) // 32) * 4
    rows: list[list[int]] = []
    for y in range(height):
        source_row = pixels[(height - 1 - y) * stride : (height - y) * stride]
        if bpp == 8:
            rows.append(list(source_row[:width]))
            continue
        values: list[int] = []
        for byte in source_row[: (width + 1) // 2]:
            values.append(byte >> 4)
            if len(values) < width:
                values.append(byte & 0x0F)
        rows.append(values)
    return rows


def row_candidates(height: int) -> list[int]:
    if height >= 390 and height % 4 == 0:
        return [4, 2, 5, 1]
    if height >= 300 and height % 3 == 0:
        return [3, 2, 4, 1]
    if height >= 160 and height % 2 == 0:
        return [2, 1, 4]
    return [1]


def target_frame_count(width: int, height: int, rows: int) -> int:
    if rows >= 4:
        return 16
    if rows == 3:
        return 12
    if rows == 2:
        return 8
    if height <= 64 and width >= 360:
        return 8
    return 6


def grid_score(index_rows: list[list[int]], width: int, height: int, columns: int, rows: int) -> tuple[float, dict] | None:
    if columns <= 0 or rows <= 0 or width % columns or height % rows:
        return None
    frame_width = width // columns
    frame_height = height // rows
    if frame_width < 24 or frame_height < 24:
        return None

    column_counts = [sum(1 for y in range(height) if index_rows[y][x] != 0) for x in range(width)]
    row_counts = [sum(1 for x in range(width) if index_rows[y][x] != 0) for y in range(height)]
    total_visible = sum(column_counts) or 1
    cut_penalty = 0
    cut_count = 0
    for column in range(1, columns):
        x = column * frame_width
        cut_penalty += sum(column_counts[max(0, x - 1) : min(width, x + 2)])
        cut_count += 1
    for row in range(1, rows):
        y = row * frame_height
        cut_penalty += sum(row_counts[max(0, y - 1) : min(height, y + 2)])
        cut_count += 1

    empty_cells = 0
    border_penalty = 0
    for row in range(rows):
        for column in range(columns):
            x0 = column * frame_width
            x1 = x0 + frame_width
            y0 = row * frame_height
            y1 = y0 + frame_height
            visible = 0
            for y in range(y0, y1):
                for x in range(x0, x1):
                    if index_rows[y][x] != 0:
                        visible += 1
            if visible == 0:
                empty_cells += 1
                continue
            for x in (x0, x1 - 1):
                border_penalty += sum(1 for y in range(y0, y1) if index_rows[y][x] != 0)
            for y in (y0, y1 - 1):
                border_penalty += sum(1 for x in range(x0, x1) if index_rows[y][x] != 0)

    cut_density = (cut_penalty / max(1, cut_count)) * 1000 / total_visible
    border_density = border_penalty * 100 / total_visible
    frame_count = columns * rows
    target_count = target_frame_count(width, height, rows)
    count_penalty = abs(frame_count - target_count) * 8

    aspect = frame_width / max(1, frame_height)
    aspect_penalty = 0.0
    if aspect < 0.55:
        aspect_penalty = (0.55 - aspect) * 120
    elif aspect > 2.2:
        aspect_penalty = (aspect - 2.2) * 80

    row_preference_penalty = 0
    preferred_rows = row_candidates(height)[0]
    if rows != preferred_rows:
        row_preference_penalty = 35

    score = cut_density + border_density * 0.18 + empty_cells * 90 + count_penalty + aspect_penalty + row_preference_penalty
    return score, {
        "columns": columns,
        "rows": rows,
        "sourceWidth": frame_width,
        "sourceHeight": frame_height,
        "frameCount": frame_count,
        "score": round(score, 3),
        "cutPenalty": int(cut_penalty),
        "emptyCells": empty_cells,
    }


def inferred_layout_candidates(name: str, payload: dict | None) -> list[dict]:
    width = int((payload or {}).get("width") or 0)
    height = int((payload or {}).get("height") or 0)
    if not width or not height:
        return []
    src = EXTRACT_FLD / name
    if not src.exists():
        return []
    try:
        decoded = decompress_cns(src.read_bytes())
        parsed_width, parsed_height, _palette, pixels, bpp = parse_image(decoded)
        if parsed_width != width or parsed_height != height:
            width, height = parsed_width, parsed_height
        index_rows = unpack_palette_indices(width, height, pixels, bpp)
    except Exception:
        return []

    candidates: list[tuple[float, dict]] = []
    row_options = row_candidates(height)
    column_options = [2, 3, 4, 5, 6, 7, 8, 9, 10, 12]
    if height >= 160:
        column_options = [4, 3, 5, 2, 6, 8, 10, 12]
    for rows in row_options:
        for columns in column_options:
            scored = grid_score(index_rows, width, height, columns, rows)
            if scored is not None:
                candidates.append(scored)
    candidates.sort(key=lambda item: (item[0], item[1]["columns"], item[1]["rows"]))
    return [candidate for _score, candidate in candidates[:8]]


def frame_hint(name: str, payload: dict | None) -> dict:
    width = int((payload or {}).get("width") or 0)
    height = int((payload or {}).get("height") or 0)
    if not width or not height:
        return {"sourceX": 0, "sourceY": 0, "sourceWidth": 0, "sourceHeight": 0}
    layout_candidates = inferred_layout_candidates(name, payload)
    source_width = min(width, 96 if width >= 96 else width)
    source_height = min(height, 96 if height >= 96 else height)
    if height <= 96 and width >= height:
        source_width = min(width, max(1, height))
        source_height = height
    return {
        "sourceX": 0,
        "sourceY": 0,
        "sourceWidth": source_width,
        "sourceHeight": source_height,
        "columns": max(1, width // max(1, source_width)),
        "rows": max(1, height // max(1, source_height)),
        "frameCount": max(1, (width // max(1, source_width)) * (height // max(1, source_height))),
        "inference": "legacy-dimension-fallback",
        "confidence": "low",
        "layoutCandidates": layout_candidates,
        "layoutCandidateSource": "palette0-grid-cut-score-review-only",
        "originalFrameRectsBound": False,
        "originalFrameOffsetsBound": False,
        "originalAnimationScriptBound": False,
    }


def all_enemy_sprite_asset_rows(
    payload_by_name: dict[str, dict],
    enemy_refs: list[dict],
    out_dir: Path,
) -> tuple[list[dict], dict[str, str]]:
    ref_by_name: dict[str, dict] = {}
    for row in sorted(enemy_refs, key=lambda item: (str(item.get("name") or ""), int(item.get("refVa") or 0))):
        name = str(row.get("name") or "")
        if name and name not in ref_by_name:
            ref_by_name[name] = row
    names = sorted(
        {
            path.name
            for pattern in ("boss*.cns", "z*.cns")
            for path in EXTRACT_FLD.glob(pattern)
        }
        | set(ref_by_name)
    )
    rows: list[dict] = []
    asset_paths: dict[str, str] = {}
    for name in names:
        asset_key = Path(name).stem
        asset_path = cns_browser_path(name)
        if asset_path:
            asset_paths[asset_key] = asset_path
        payload = payload_by_name.get(name)
        ref = ref_by_name.get(name) or {}
        rows.append({
            "enemyCns": name,
            "enemyAssetKey": asset_key,
            "enemyAssetPath": asset_path,
            "enemyRefVa": ref.get("refVa"),
            "enemyRefVaHex": ref.get("refVaHex") or (
                f"0x{int(ref['refVa']):08x}" if isinstance(ref.get("refVa"), int) else ""
            ),
            "payloadWidth": (payload or {}).get("width"),
            "payloadHeight": (payload or {}).get("height"),
            "frameHint": frame_hint(name, payload),
            "selectionKind": "extracted-enemy-object-sprite",
            "originalEnemyRowBound": False,
            "originalStatsOrRewardsBound": False,
        })
    return rows, asset_paths


def build_summary(
    resource_descriptors: dict,
    battle_event_candidates: dict,
    cns_payloads: list[dict],
    out_dir: Path,
) -> dict:
    payload_by_name = {row.get("name"): row for row in cns_payloads if row.get("name")}
    backgrounds = ref_rows(resource_descriptors, "battle-background-tilemap")
    enemy_refs = ref_rows(resource_descriptors, "enemy-object-sprite-image")
    sprite_asset_rows, asset_paths = all_enemy_sprite_asset_rows(payload_by_name, enemy_refs, out_dir)
    background_by_name: dict[str, list[dict]] = {}
    for row in backgrounds:
        background_by_name.setdefault(str(row.get("name") or ""), []).append(row)

    candidate_rows = []
    candidate_asset_paths: dict[str, str] = {}
    for candidate in battle_event_candidates.get("candidates") or []:
        background_key = str(candidate.get("battleBackground") or "")
        background_name = f"{background_key}.cns" if background_key and not background_key.endswith(".cns") else background_key
        matching_backgrounds = background_by_name.get(background_name) or []
        best_pair = None
        for background_ref in matching_backgrounds:
            enemy = nearest_enemy(background_ref, enemy_refs)
            if not enemy:
                continue
            distance = abs(int(enemy["refVa"]) - int(background_ref["refVa"]))
            pair = (distance, background_ref, enemy)
            if best_pair is None or pair[0] < best_pair[0]:
                best_pair = pair
        if not best_pair:
            continue
        distance, background_ref, enemy = best_pair
        enemy_name = str(enemy.get("name") or "")
        asset_key = Path(enemy_name).stem
        asset_path = cns_browser_path(enemy_name)
        if asset_path:
            asset_paths[asset_key] = asset_path
            candidate_asset_paths[asset_key] = asset_path
        payload = payload_by_name.get(enemy_name)
        hint = frame_hint(enemy_name, payload)
        candidate_rows.append({
            "candidateId": candidate.get("id") or "",
            "blockId": candidate.get("blockId") or "",
            "battleBackground": background_key,
            "battleBackgroundRefVa": background_ref.get("refVa"),
            "battleBackgroundRefVaHex": background_ref.get("refVaHex") or f"0x{int(background_ref['refVa']):08x}",
            "enemyCns": enemy_name,
            "enemyAssetKey": asset_key,
            "enemyAssetPath": asset_path,
            "enemyRefVa": enemy.get("refVa"),
            "enemyRefVaHex": enemy.get("refVaHex") or f"0x{int(enemy['refVa']):08x}",
            "refDistance": distance,
            "payloadWidth": (payload or {}).get("width"),
            "payloadHeight": (payload or {}).get("height"),
            "frameHint": hint,
            "selectionKind": "nearest-exe-resource-descriptor",
            "originalEnemyRowBound": False,
            "originalStatsOrRewardsBound": False,
        })

    by_candidate = {row["candidateId"]: row for row in candidate_rows if row.get("candidateId")}
    by_background = {}
    for row in sorted(candidate_rows, key=lambda item: (item["refDistance"], item["candidateId"])):
        by_background.setdefault(row["battleBackground"], row)
    by_enemy_asset_key = {
        row["enemyAssetKey"]: row
        for row in sprite_asset_rows
        if row.get("enemyAssetKey")
    }
    candidate_asset_keys = {row["enemyAssetKey"] for row in candidate_rows if row.get("enemyAssetKey")}

    return {
        "scope": "Battle enemy sprite candidates chosen by EXE resource descriptor proximity.",
        "source": [
            "in-memory EXE resource descriptor proximity",
            "out/battle_event_candidates.json",
            "out/cns_payloads.json",
            "extract_fld/*.cns",
        ],
        "promotionStatus": "sprite-proximity-candidate-not-enemy-row",
        "candidateCount": len(candidate_rows),
        "uniqueEnemySpriteCount": len({row["enemyCns"] for row in candidate_rows}),
        "candidateAssetCount": len(candidate_asset_paths),
        "extractedEnemySpriteAssetCount": len(sprite_asset_rows),
        "assetCount": len(asset_paths),
        "assets": asset_paths,
        "spriteAssets": sprite_asset_rows,
        "candidates": candidate_rows,
        "byCandidateId": by_candidate,
        "byBattleBackground": by_background,
        "byEnemyAssetKey": by_enemy_asset_key,
        "checks": {
            "battleEnemySpriteCandidatesMapped": len(candidate_rows) > 0,
            "candidateEnemyAssetsHaveCnsSource": len(candidate_asset_paths) == len(candidate_asset_keys),
            "enemyAssetsHaveCnsSource": len(asset_paths) == len({row["enemyAssetKey"] for row in sprite_asset_rows if row["enemyAssetPath"]}),
            "allExtractedEnemySpriteAssetsRuntimeSelectable": len(sprite_asset_rows) == len(asset_paths) and len(sprite_asset_rows) >= len(candidate_asset_paths),
            "originalEnemyRowsBound": False,
            "originalStatsOrRewardsBound": False,
            "originalBattleEntryProven": False,
        },
        "conclusion": (
            f"{len(candidate_rows)} battle review candidates now have an EXE-nearest enemy/object sprite visual candidate. "
            f"The runtime asset table also exposes {len(sprite_asset_rows)} extracted boss*/z* sprite assets for prototype battle selection. "
            "This improves the browser battle review scene with original sprite assets, but it is proximity evidence only: "
            "enemy rows, stats, formations, rewards, formulas, and battle-entry execution remain unproven."
        ),
    }


def html_page(summary: dict) -> str:
    check_rows = "".join(
        f"<tr><td>{html.escape(key)}</td><td>{html.escape(str(value))}</td></tr>"
        for key, value in summary["checks"].items()
    )
    candidate_rows = "".join(
        "<tr>"
        f"<td>{html.escape(row['candidateId'])}</td>"
        f"<td><code>{html.escape(row['battleBackground'])}</code></td>"
        f"<td><code>{html.escape(row['enemyCns'])}</code></td>"
        f"<td><code>{html.escape(row['battleBackgroundRefVaHex'])}</code></td>"
        f"<td><code>{html.escape(row['enemyRefVaHex'])}</code></td>"
        f"<td>{row['refDistance']}</td>"
        f"<td><code>{html.escape(row['enemyAssetKey'])}</code></td>"
        "</tr>"
        for row in summary["candidates"]
    )
    sprite_rows = "".join(
        "<tr>"
        f"<td><code>{html.escape(row['enemyCns'])}</code></td>"
        f"<td>{html.escape(str(row.get('payloadWidth') or '-'))}x{html.escape(str(row.get('payloadHeight') or '-'))}</td>"
        f"<td><code>{html.escape(row.get('enemyRefVaHex') or '-')}</code></td>"
        f"<td><code>{html.escape(row['enemyAssetKey'])}</code></td>"
        "</tr>"
        for row in summary["spriteAssets"]
    )
    return "\n".join([
        "<!doctype html>",
        '<html lang="en">',
        "<head>",
        '  <meta charset="utf-8">',
        "  <title>Battle Enemy Sprite Candidates</title>",
        "  <style>body{font-family:system-ui,sans-serif;margin:24px;line-height:1.45;max-width:1200px}table{border-collapse:collapse;width:100%;margin:16px 0}td,th{border:1px solid #ddd;padding:6px 8px;text-align:left;vertical-align:top}th{background:#f5f5f5}code{white-space:nowrap}</style>",
        "</head>",
        "<body>",
        "  <h1>Battle Enemy Sprite Candidates</h1>",
        f"  <p>{html.escape(summary['conclusion'])}</p>",
        f"  <p>Promotion status: <code>{html.escape(summary['promotionStatus'])}</code></p>",
        "  <h2>Checks</h2>",
        f"  <table><thead><tr><th>check</th><th>value</th></tr></thead><tbody>{check_rows}</tbody></table>",
        "  <h2>Candidates</h2>",
        f"  <table><thead><tr><th>battle candidate</th><th>background</th><th>enemy sprite</th><th>bg ref</th><th>enemy ref</th><th>distance</th><th>asset</th></tr></thead><tbody>{candidate_rows}</tbody></table>",
        "  <h2>Extracted Sprite Assets</h2>",
        f"  <table><thead><tr><th>sprite</th><th>size</th><th>ref</th><th>runtime asset</th></tr></thead><tbody>{sprite_rows}</tbody></table>",
        "</body>",
        "</html>",
        "",
    ])


def js_json(value: Any) -> str:
    return json.dumps(value, ensure_ascii=False, separators=(",", ":"))


def write_outputs(summary: dict, out_dir: Path) -> None:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "battle_enemy_candidates.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    (out_dir / "battle_enemy_candidates.js").write_text(
        "window.HWANSE_BATTLE_ENEMY_CANDIDATES = "
        + js_json({
            "promotionStatus": summary["promotionStatus"],
            "candidateCount": summary["candidateCount"],
            "uniqueEnemySpriteCount": summary["uniqueEnemySpriteCount"],
            "candidateAssetCount": summary["candidateAssetCount"],
            "extractedEnemySpriteAssetCount": summary["extractedEnemySpriteAssetCount"],
            "candidates": summary["candidates"],
            "spriteAssets": summary["spriteAssets"],
            "byCandidateId": summary["byCandidateId"],
            "byBattleBackground": summary["byBattleBackground"],
            "byEnemyAssetKey": summary["byEnemyAssetKey"],
        })
        + ";\nwindow.HWANSE_BATTLE_ENEMY_ASSETS = "
        + js_json(summary["assets"])
        + ";\n",
        encoding="utf-8",
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out-dir", type=Path, default=OUT)
    args = parser.parse_args()
    cns_payloads = load_json(args.out_dir / "cns_payloads.json", [])
    resource_descriptors = load_json(args.out_dir / "original_battle_resource_descriptors.json", {})
    if not resource_descriptors:
        exe_data = EXE.read_bytes()
        sections = read_sections(exe_data)
        resource_descriptors = build_resource_descriptor_summary(
            exe_data,
            sections,
            cns_payloads,
        )
    summary = build_summary(
        resource_descriptors,
        load_json(args.out_dir / "battle_event_candidates.json", {}),
        cns_payloads,
        args.out_dir,
    )
    write_outputs(summary, args.out_dir)
    print(f"wrote battle enemy sprite candidates -> {args.out_dir / 'battle_enemy_candidates.js'}")


if __name__ == "__main__":
    main()