#!/usr/bin/env python3
"""
Test encode + decode + trace verification for enc_av1.

Encodes 8 test sequences (including crop variants), decodes with MSAC trace,
and verifies that disabled tools do NOT appear in the bitstream.
"""

import json
import os
import shutil
import subprocess
import sys
import tempfile
from collections import Counter
from pathlib import Path

# ============================================================================
# Paths
# ============================================================================

SCRIPT_DIR = Path(__file__).parent
PROJECT_DIR = SCRIPT_DIR.parent
ENCODE_PY = SCRIPT_DIR / "encode.py"
DECODER = PROJECT_DIR / "_ref_av1_aom" / "build_msac" / "aomdec"
SEQ_DIR = PROJECT_DIR / "test_av1" / "sequences"

# ============================================================================
# Enum mappings (from AV1 spec / libaom enums.h)
# ============================================================================

PARTITION_ENUM = {
    0: "PARTITION_NONE",
    1: "PARTITION_HORZ",
    2: "PARTITION_VERT",
    3: "PARTITION_SPLIT",
    4: "PARTITION_HORZ_A",
    5: "PARTITION_HORZ_B",
    6: "PARTITION_VERT_A",
    7: "PARTITION_VERT_B",
    8: "PARTITION_HORZ_4",
    9: "PARTITION_VERT_4",
}

INTRA_MODE_ENUM = {
    0: "DC_PRED",
    1: "V_PRED",
    2: "H_PRED",
    3: "D45_PRED",
    4: "D135_PRED",
    5: "D113_PRED",
    6: "D157_PRED",
    7: "D203_PRED",
    8: "D67_PRED",
    9: "SMOOTH_PRED",
    10: "SMOOTH_V_PRED",
    11: "SMOOTH_H_PRED",
    12: "PAETH_PRED",
}

# Inter modes start at value 13 in the PREDICTION_MODE enum
INTER_MODE_ENUM = {
    13: "NEARESTMV",
    14: "NEARMV",
    15: "GLOBALMV",
    16: "NEWMV",
    17: "NEAREST_NEARESTMV",
    18: "NEAR_NEARMV",
    19: "NEAREST_NEWMV",
    20: "NEW_NEARESTMV",
    21: "NEAR_NEWMV",
    22: "NEW_NEARMV",
    23: "GLOBAL_GLOBALMV",
    24: "NEW_NEWMV",
}

TX_TYPE_ENUM = {
    0: "DCT_DCT",
    1: "ADST_DCT",
    2: "DCT_ADST",
    3: "ADST_ADST",
    4: "FLIPADST_DCT",
    5: "DCT_FLIPADST",
    6: "FLIPADST_FLIPADST",
    7: "ADST_FLIPADST",
    8: "FLIPADST_ADST",
    9: "IDTX",
    10: "V_DCT",
    11: "H_DCT",
    12: "V_ADST",
    13: "H_ADST",
    14: "V_FLIPADST",
    15: "H_FLIPADST",
}

COMPOUND_TYPE_ENUM = {
    0: "COMPOUND_AVERAGE",
    1: "COMPOUND_DISTWTD",
    2: "COMPOUND_WEDGE",
    3: "COMPOUND_DIFFWTD",
}

# ============================================================================
# Test cases
# ============================================================================

TESTS = [
    {
        "name": "01_B0_black",
        "ini": "B0_infrastructure.ini",
        "input": "test_256x256_black_1fr.yuv",
        "args": ["-n", "1"],
        "desc": "B0: minimal — DC_PRED + DCT_DCT only",
    },
    {
        "name": "02_B1_intra",
        "ini": "B1_intra.ini",
        "input": "test_256x256.yuv",
        "args": ["-n", "1"],
        "desc": "B1: all intra modes + 7 TX types",
    },
    {
        "name": "03_B2_features",
        "ini": "B2_intra_features.ini",
        "input": "test_256x256_palette.yuv",
        "args": ["-n", "1"],
        "desc": "B2: intra features (angle delta, CfL, palette, filter intra)",
    },
    {
        "name": "04_B3_inter",
        "ini": "B3_inter.ini",
        "input": "test_256x256_motion.yuv",
        "args": ["-n", "5"],
        "desc": "B3: inter prediction, all refs, compound modes",
    },
    {
        "name": "05_B4_deblock",
        "ini": "B4_deblocking.ini",
        "input": "test_256x256_motion.yuv",
        "args": ["-n", "5"],
        "desc": "B4: deblocking ON, CDEF OFF, LR OFF",
    },
    {
        "name": "06_B5_full",
        "ini": "B5_full.ini",
        "input": "test_256x256_motion.yuv",
        "args": ["-n", "5"],
        "desc": "B5: full codec — all tools ON",
    },
    {
        "name": "07_B1_crop",
        "ini": "B1_intra.ini",
        "input": "test_256x256.yuv",
        "args": ["-n", "1"],
        "desc": "B1 + crop: 256x256 -> 128x128",
        "crop": {"enabled": True, "width": 128, "height": 128, "offset_x": 64, "offset_y": 64},
    },
    {
        "name": "08_B3_crop",
        "ini": "B3_inter.ini",
        "input": "test_256x256_motion.yuv",
        "args": ["-n", "3"],
        "desc": "B3 + crop: 256x256 -> 128x64",
        "crop": {"enabled": True, "width": 128, "height": 64, "offset_x": 0, "offset_y": 96},
    },
]


# ============================================================================
# INI reading (lightweight, matches encode.py logic)
# ============================================================================

import configparser

def read_ini(path):
    """Read INI config and return tool settings."""
    cfg = configparser.ConfigParser()
    cfg.read(path)

    result = {
        "filters": {},
        "features": {},
        "tools": {},
    }

    # Filters
    if cfg.has_section("filters"):
        result["filters"] = {
            "deblocking": cfg.getboolean("filters", "deblocking", fallback=True),
            "cdef": cfg.getboolean("filters", "cdef", fallback=True),
            "loop_restoration": cfg.getboolean("filters", "loop_restoration", fallback=True),
        }

    # Features
    if cfg.has_section("features"):
        result["features"] = {
            "filter_intra": cfg.getboolean("features", "filter_intra", fallback=True),
            "palette": cfg.getboolean("features", "palette", fallback=True),
            "cfl_intra": cfg.getboolean("features", "cfl_intra", fallback=True),
            "angle_delta": cfg.getboolean("features", "angle_delta", fallback=True),
        }

    # Tool sections
    tool_sections = {
        "partitions": list(PARTITION_ENUM.values()),
        "intra_modes": list(INTRA_MODE_ENUM.values()),
        "inter_modes": list(INTER_MODE_ENUM.values()),
        "compound_types": list(COMPOUND_TYPE_ENUM.values()),
        "tx_types": list(TX_TYPE_ENUM.values()),
    }
    for section, tools in tool_sections.items():
        result["tools"][section] = {}
        if cfg.has_section(section):
            for tool in tools:
                result["tools"][section][tool] = cfg.getboolean(
                    section, tool, fallback=True)
        else:
            for tool in tools:
                result["tools"][section][tool] = True

    return result


# ============================================================================
# Encode
# ============================================================================

def encode(test, out_dir):
    """Run encode.py and return path to .ivf file."""
    ivf_path = out_dir / f"{test['name']}.ivf"
    ini_path = SCRIPT_DIR / test["ini"]

    # If test has crop, create a modified INI
    if "crop" in test:
        tmp_ini = out_dir / f"{test['name']}.ini"
        # Read original INI, modify crop section
        cfg = configparser.ConfigParser()
        cfg.read(ini_path)
        crop = test["crop"]
        cfg.set("crop", "enabled", "true")
        cfg.set("crop", "width", str(crop["width"]))
        cfg.set("crop", "height", str(crop["height"]))
        cfg.set("crop", "offset_x", str(crop["offset_x"]))
        cfg.set("crop", "offset_y", str(crop["offset_y"]))
        with open(tmp_ini, "w") as f:
            cfg.write(f)
        ini_path = tmp_ini

    input_path = SEQ_DIR / test["input"]
    cmd = [
        sys.executable, str(ENCODE_PY),
        "-i", str(input_path),
        "-o", str(ivf_path),
        "-cfg", str(ini_path),
    ] + test.get("args", [])

    result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
    if result.returncode != 0:
        print(f"  ENCODE FAILED!")
        print(result.stderr[:500])
        return None

    size = ivf_path.stat().st_size
    print(f"  Encoded: {ivf_path.name} ({size} bytes)")
    return ivf_path


# ============================================================================
# Decode with trace
# ============================================================================

def decode_with_trace(ivf_path, trace_dir):
    """Decode .ivf with MSAC trace and CDQ output."""
    cdq_dir = trace_dir / "cdq"
    msac_dir = trace_dir / "msac"
    cdq_dir.mkdir(parents=True, exist_ok=True)
    msac_dir.mkdir(parents=True, exist_ok=True)

    cmd = [
        str(DECODER),
        "--noblit",
        f"--out-cdq={cdq_dir}",
        f"--out-msac={msac_dir}",
        str(ivf_path),
    ]

    result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    if result.returncode != 0:
        print(f"  DECODE FAILED!")
        print(result.stderr[:500])
        return None

    print(f"  Decoded with trace OK")
    return {"cdq": cdq_dir, "msac": msac_dir}


# ============================================================================
# Parse traces
# ============================================================================

def parse_msac_traces(msac_dir):
    """Parse all MSAC JSONL trace files and collect symbol statistics."""
    stats = {
        "partition_values": Counter(),
        "intra_mode_values": Counter(),
        "tx_type_values": Counter(),
        "inter_mode_values": Counter(),
        "compound_type_values": Counter(),
        "symbols": Counter(),
        "angle_delta_nonzero": 0,
    }

    for f in sorted(msac_dir.glob("*_msac.jsonl")):
        with open(f) as fh:
            for line in fh:
                try:
                    d = json.loads(line)
                except json.JSONDecodeError:
                    continue
                name = d.get("name", "")
                value = d.get("value", 0)
                stats["symbols"][name] += 1

                if name == "angle_delta" and value != 3:
                    # MAX_ANGLE_DELTA=3 means delta=0; anything else is nonzero
                    stats["angle_delta_nonzero"] += 1

                if name == "partition":
                    stats["partition_values"][value] += 1
                elif name == "intra_mode":
                    stats["intra_mode_values"][value] += 1
                elif name == "tx_type_resolved":
                    # Use resolved TX_TYPE enum (not raw CDF index from "tx_type")
                    stats["tx_type_values"][value] += 1
                elif name == "inter_mode":
                    stats["inter_mode_values"][value] += 1
                elif name == "compound_type":
                    stats["compound_type_values"][value] += 1

    return stats


def parse_cdq_frame_headers(cdq_dir):
    """Parse CDQ frame header JSONL files for filter/feature info."""
    headers = {}
    for f in sorted(cdq_dir.glob("obu_*_frame*.jsonl")):
        fields = {}
        with open(f) as fh:
            for line in fh:
                try:
                    d = json.loads(line)
                except json.JSONDecodeError:
                    continue
                name = d.get("name", "")
                value = d.get("value", 0)
                fields[name] = value
        headers[f.name] = fields
    return headers


# ============================================================================
# Verification
# ============================================================================

def verify(test_name, ini_cfg, msac_stats, cdq_headers, desc):
    """Verify trace against INI config. Returns (pass_count, fail_count, messages)."""
    passes = 0
    fails = 0
    msgs = []

    def check(condition, msg):
        nonlocal passes, fails
        if condition:
            passes += 1
        else:
            fails += 1
            msgs.append(f"  FAIL: {msg}")

    # --- Filter checks from frame headers ---
    for fname, fields in cdq_headers.items():
        flt = ini_cfg["filters"]

        # Deblocking: loop_filter_level[0] should be 0 when OFF
        if not flt.get("deblocking", True):
            lf0 = fields.get("loop_filter_level[0]", -1)
            lf1 = fields.get("loop_filter_level[1]", -1)
            check(lf0 == 0 and lf1 == 0,
                  f"deblocking OFF but loop_filter_level = [{lf0},{lf1}] in {fname}")
        else:
            # When deblocking ON, we just note it's configured
            passes += 1

        # CDEF
        if not flt.get("cdef", True):
            cdef_bits = fields.get("cdef_bits", -1)
            cdef_y = fields.get("cdef_y_strength", -1)
            cdef_uv = fields.get("cdef_uv_strength", -1)
            check(cdef_y == 0 and cdef_uv == 0,
                  f"CDEF OFF but cdef_y_strength={cdef_y} cdef_uv_strength={cdef_uv} in {fname}")
        else:
            passes += 1

        # Loop restoration
        if not flt.get("loop_restoration", True):
            # All 3 lr_type fields should be 0
            lr_count = sum(1 for k, v in fields.items() if k == "lr_type" and v != 0)
            # Actually lr_type appears multiple times, check from raw
            lr_values = [v for k, v in fields.items() if k == "lr_type"]
            # With our simple parser, we only get the last value for duplicate keys
            # Just check the stored value
            lr = fields.get("lr_type", -1)
            check(lr == 0, f"LR OFF but lr_type={lr} in {fname}")
        else:
            passes += 1

        # Only check first frame header
        break

    # --- Feature checks from MSAC symbols ---
    feat = ini_cfg["features"]

    if not feat.get("filter_intra", True):
        fi_count = msac_stats["symbols"].get("filter_intra_enable", 0)
        fi_mode = msac_stats["symbols"].get("filter_intra_mode", 0)
        # filter_intra_enable may still appear (to signal it's not used), but mode should be 0
        check(fi_mode == 0,
              f"filter_intra OFF but filter_intra_mode appears {fi_mode}x in trace")
    else:
        passes += 1

    if not feat.get("palette", True):
        pal_count = msac_stats["symbols"].get("palette_mode", 0)
        check(pal_count == 0,
              f"palette OFF but palette_mode appears {pal_count}x in trace")
    else:
        passes += 1

    if not feat.get("cfl_intra", True):
        cfl_count = msac_stats["symbols"].get("cfl:signs", 0)
        check(cfl_count == 0,
              f"CfL OFF but cfl:signs appears {cfl_count}x in trace")
    else:
        passes += 1

    if not feat.get("angle_delta", True):
        # angle_delta is a mandatory syntax element for directional modes — always
        # present in bitstream. When disabled, encoder should always code delta=0
        # (value=MAX_ANGLE_DELTA=3, center of range [-3,+3]).
        ad_nonzero = msac_stats.get("angle_delta_nonzero", 0)
        check(ad_nonzero == 0,
              f"angle_delta OFF but {ad_nonzero} symbols have delta != 0")
    else:
        passes += 1

    # --- Partition type checks ---
    tools_part = ini_cfg["tools"].get("partitions", {})
    for val, name in PARTITION_ENUM.items():
        if not tools_part.get(name, True):
            count = msac_stats["partition_values"].get(val, 0)
            check(count == 0,
                  f"{name} disabled but appears {count}x in partition trace (value={val})")
        elif msac_stats["partition_values"].get(val, 0) > 0:
            passes += 1

    # --- Intra mode checks ---
    tools_intra = ini_cfg["tools"].get("intra_modes", {})
    for val, name in INTRA_MODE_ENUM.items():
        if not tools_intra.get(name, True):
            count = msac_stats["intra_mode_values"].get(val, 0)
            check(count == 0,
                  f"{name} disabled but appears {count}x in intra_mode trace (value={val})")
        elif msac_stats["intra_mode_values"].get(val, 0) > 0:
            passes += 1

    # --- TX type checks ---
    tools_tx = ini_cfg["tools"].get("tx_types", {})
    for val, name in TX_TYPE_ENUM.items():
        if not tools_tx.get(name, True):
            count = msac_stats["tx_type_values"].get(val, 0)
            check(count == 0,
                  f"{name} disabled but appears {count}x in tx_type trace (value={val})")
        elif msac_stats["tx_type_values"].get(val, 0) > 0:
            passes += 1

    # --- Inter mode checks (only for inter presets) ---
    tools_inter = ini_cfg["tools"].get("inter_modes", {})
    for val, name in INTER_MODE_ENUM.items():
        if not tools_inter.get(name, True):
            count = msac_stats["inter_mode_values"].get(val, 0)
            check(count == 0,
                  f"{name} disabled but appears {count}x in inter_mode trace (value={val})")
        elif msac_stats["inter_mode_values"].get(val, 0) > 0:
            passes += 1

    # --- Compound type checks ---
    tools_comp = ini_cfg["tools"].get("compound_types", {})
    for val, name in COMPOUND_TYPE_ENUM.items():
        if not tools_comp.get(name, True):
            count = msac_stats["compound_type_values"].get(val, 0)
            check(count == 0,
                  f"{name} disabled but appears {count}x in compound_type trace (value={val})")
        elif msac_stats["compound_type_values"].get(val, 0) > 0:
            passes += 1

    return passes, fails, msgs


# ============================================================================
# Summary: print trace statistics
# ============================================================================

def print_trace_summary(msac_stats):
    """Print a concise summary of what tools were found in the trace."""
    parts_used = [PARTITION_ENUM[v] for v in sorted(msac_stats["partition_values"]) if v in PARTITION_ENUM]
    intra_used = [INTRA_MODE_ENUM[v] for v in sorted(msac_stats["intra_mode_values"]) if v in INTRA_MODE_ENUM]
    tx_used = [TX_TYPE_ENUM[v] for v in sorted(msac_stats["tx_type_values"]) if v in TX_TYPE_ENUM]
    inter_used = [INTER_MODE_ENUM[v] for v in sorted(msac_stats["inter_mode_values"]) if v in INTER_MODE_ENUM]
    comp_used = [COMPOUND_TYPE_ENUM[v] for v in sorted(msac_stats["compound_type_values"]) if v in COMPOUND_TYPE_ENUM]

    if parts_used:
        print(f"    Partitions: {', '.join(parts_used)}")
    if intra_used:
        print(f"    Intra:      {', '.join(intra_used)}")
    if tx_used:
        print(f"    TX types:   {', '.join(tx_used)}")
    if inter_used:
        print(f"    Inter:      {', '.join(inter_used)}")
    if comp_used:
        print(f"    Compound:   {', '.join(comp_used)}")

    # Notable symbols
    for sym in ["angle_delta", "cfl:signs", "filter_intra_mode", "filter_intra_enable"]:
        c = msac_stats["symbols"].get(sym, 0)
        if c > 0:
            print(f"    {sym}: {c}x")


# ============================================================================
# Main
# ============================================================================

def main():
    if not DECODER.exists():
        print(f"ERROR: Decoder not found: {DECODER}")
        sys.exit(1)
    if not ENCODE_PY.exists():
        print(f"ERROR: encode.py not found: {ENCODE_PY}")
        sys.exit(1)

    work_dir = Path(tempfile.mkdtemp(prefix="enc_av1_test_"))
    print(f"Work dir: {work_dir}")
    print(f"{'='*72}")

    total_pass = 0
    total_fail = 0
    all_results = []

    for i, test in enumerate(TESTS):
        print(f"\n[{i+1}/{len(TESTS)}] {test['name']}: {test['desc']}")
        print(f"{'-'*72}")

        test_dir = work_dir / test["name"]
        test_dir.mkdir()

        # 1. Encode
        ivf = encode(test, test_dir)
        if ivf is None:
            total_fail += 1
            all_results.append((test["name"], 0, 1, ["  FAIL: encode failed"]))
            continue

        # 2. Decode with trace
        trace_dir = test_dir / "trace"
        trace = decode_with_trace(ivf, trace_dir)
        if trace is None:
            total_fail += 1
            all_results.append((test["name"], 0, 1, ["  FAIL: decode failed"]))
            continue

        # 3. Parse traces
        msac_stats = parse_msac_traces(trace["msac"])
        cdq_headers = parse_cdq_frame_headers(trace["cdq"])

        print_trace_summary(msac_stats)

        # 4. Read INI config
        ini_path = SCRIPT_DIR / test["ini"]
        ini_cfg = read_ini(ini_path)

        # 5. Verify
        p, f, msgs = verify(test["name"], ini_cfg, msac_stats, cdq_headers, test["desc"])
        total_pass += p
        total_fail += f
        all_results.append((test["name"], p, f, msgs))

        if f == 0:
            print(f"  RESULT: PASS ({p} checks)")
        else:
            print(f"  RESULT: {f} FAILS / {p} passes")
            for m in msgs:
                print(m)

    # Summary
    print(f"\n{'='*72}")
    print(f"SUMMARY: {total_pass} passed, {total_fail} failed")
    print(f"{'='*72}")
    for name, p, f, msgs in all_results:
        status = "PASS" if f == 0 else f"FAIL({f})"
        print(f"  {name}: {status} ({p} checks)")
        for m in msgs:
            print(m)

    if total_fail > 0:
        print(f"\nWork dir preserved: {work_dir}")
    else:
        print(f"\nAll tests passed! Cleaning up...")
        shutil.rmtree(work_dir, ignore_errors=True)

    sys.exit(1 if total_fail > 0 else 0)


if __name__ == "__main__":
    main()
