#!/usr/bin/env python3
"""
AV1 Encoder Wrapper — encode YUV using an INI config file.

Usage:
  # Generate default config:
  ./encode.py --init my_config.ini

  # Encode:
  ./encode.py -i input.yuv -o output.ivf -cfg my_config.ini

  # Source dimensions are in the INI (or override with -W/-H):
  ./encode.py -i big.yuv -o out.ivf -cfg my_config.ini -W 1920 -H 1080
"""

import argparse
import configparser
import subprocess
import sys
import tempfile
from pathlib import Path

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

SCRIPT_DIR = Path(__file__).parent
PROJECT_DIR = SCRIPT_DIR.parent
REF_AOM = PROJECT_DIR / "_ref_av1_aom"
ENCODER = REF_AOM / "build_enc" / "aomenc"

# ============================================================================
# Tool definitions — order matches AV1 spec
# ============================================================================

TOOLS = {
    "partitions": [
        "PARTITION_NONE", "PARTITION_HORZ", "PARTITION_VERT", "PARTITION_SPLIT",
        "PARTITION_HORZ_A", "PARTITION_HORZ_B", "PARTITION_VERT_A", "PARTITION_VERT_B",
        "PARTITION_HORZ_4", "PARTITION_VERT_4",
    ],
    "intra_modes": [
        "DC_PRED", "V_PRED", "H_PRED",
        "D45_PRED", "D135_PRED", "D113_PRED", "D157_PRED", "D203_PRED", "D67_PRED",
        "SMOOTH_PRED", "SMOOTH_V_PRED", "SMOOTH_H_PRED", "PAETH_PRED",
    ],
    "inter_refs": [
        "LAST_FRAME", "LAST2_FRAME", "LAST3_FRAME", "GOLDEN_FRAME",
        "BWDREF_FRAME", "ALTREF2_FRAME", "ALTREF_FRAME", "COMPOUND",
    ],
    "inter_modes": [
        "NEARESTMV", "NEARMV", "GLOBALMV", "NEWMV",
        "NEAREST_NEARESTMV", "NEAR_NEARMV", "NEW_NEWMV",
        "NEAREST_NEWMV", "NEW_NEARESTMV", "NEAR_NEWMV", "NEW_NEARMV",
        "GLOBAL_GLOBALMV",
    ],
    "compound_types": [
        "COMPOUND_AVERAGE", "COMPOUND_DISTWTD", "COMPOUND_WEDGE", "COMPOUND_DIFFWTD",
    ],
    "tx_sizes": [
        "TX_4X4", "TX_8X8", "TX_16X16", "TX_32X32", "TX_64X64",
    ],
    "tx_types": [
        "DCT_DCT", "ADST_DCT", "DCT_ADST", "ADST_ADST",
        "FLIPADST_DCT", "DCT_FLIPADST", "FLIPADST_FLIPADST",
        "ADST_FLIPADST", "FLIPADST_ADST", "IDTX",
        "V_DCT", "H_DCT", "V_ADST", "H_ADST", "V_FLIPADST", "H_FLIPADST",
    ],
}

# CSV category names (match what tools_config.c expects)
CSV_CATEGORIES = {
    "partitions": "partition",
    "intra_modes": "intra",
    "inter_refs": "inter",
    "inter_modes": "inter_mode",
    "compound_types": "compound_type",
    "tx_sizes": "tx_size",
    "tx_types": "transform",
}

# ============================================================================
# Generate default INI
# ============================================================================

def generate_default_ini(path):
    """Write a default .ini config with all tools listed."""
    lines = []
    lines.append("; AV1 Encoder Configuration")
    lines.append("; Edit true/false to enable/disable individual tools")
    lines.append(";")
    lines.append("; Usage:  ./encode.py -i input.yuv -o output.ivf -cfg this_file.ini")
    lines.append("")

    lines.append("[source]")
    lines.append("; Source YUV dimensions (override with -W/-H on command line)")
    lines.append("width = 256")
    lines.append("height = 256")
    lines.append("; frames = 0 means encode all frames in file")
    lines.append("frames = 0")
    lines.append("")

    lines.append("[crop]")
    lines.append("; Set enabled = true to auto-crop (ffmpeg) before encoding")
    lines.append("enabled = false")
    lines.append("width = 128")
    lines.append("height = 128")
    lines.append("offset_x = 0")
    lines.append("offset_y = 0")
    lines.append("")

    lines.append("[encode]")
    lines.append("cq_level = 30")
    lines.append("bit_depth = 8")
    lines.append("sb_size = 64")
    lines.append("cpu_used = 3")
    lines.append("kf_max_dist = 10")
    lines.append("; Extra aomenc flags (appended at the end)")
    lines.append("extra_flags =")
    lines.append("")

    lines.append("[filters]")
    lines.append("deblocking = true")
    lines.append("cdef = true")
    lines.append("loop_restoration = true")
    lines.append("")

    lines.append("[features]")
    lines.append("angle_delta = true")
    lines.append("cfl_intra = true")
    lines.append("palette = true")
    lines.append("filter_intra = true")
    lines.append("rect_tx = true")
    lines.append("global_motion = true")
    lines.append("obmc = true")
    lines.append("warped_motion = true")
    lines.append("")

    for section, tool_list in TOOLS.items():
        lines.append(f"[{section}]")
        for tool in tool_list:
            lines.append(f"{tool} = true")
        lines.append("")

    Path(path).write_text('\n'.join(lines))
    print(f"Created: {path}")
    print(f"Encode:  ./encode.py -i input.yuv -o output.ivf -cfg {path}")


# ============================================================================
# Read INI
# ============================================================================

def read_config(path):
    """Read .ini config, return parsed dict."""
    cfg = configparser.ConfigParser()
    cfg.read(path)

    result = {}

    # Source dimensions
    result["source"] = {
        "width": cfg.getint("source", "width", fallback=256),
        "height": cfg.getint("source", "height", fallback=256),
        "frames": cfg.getint("source", "frames", fallback=0),
    }

    # Crop
    result["crop"] = {
        "enabled": cfg.getboolean("crop", "enabled", fallback=False),
        "width": cfg.getint("crop", "width", fallback=0),
        "height": cfg.getint("crop", "height", fallback=0),
        "offset_x": cfg.getint("crop", "offset_x", fallback=0),
        "offset_y": cfg.getint("crop", "offset_y", fallback=0),
    }

    # Encode params
    result["encode"] = {
        "cq_level": cfg.getint("encode", "cq_level", fallback=30),
        "bit_depth": cfg.getint("encode", "bit_depth", fallback=8),
        "sb_size": cfg.getint("encode", "sb_size", fallback=64),
        "cpu_used": cfg.getint("encode", "cpu_used", fallback=3),
        "kf_max_dist": cfg.getint("encode", "kf_max_dist", fallback=10),
        "extra_flags": cfg.get("encode", "extra_flags", fallback="").strip(),
    }

    # Filters
    result["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
    result["features"] = {}
    if cfg.has_section("features"):
        result["features"] = {
            "angle_delta": cfg.getboolean("features", "angle_delta", fallback=True),
            "cfl_intra": cfg.getboolean("features", "cfl_intra", fallback=True),
            "palette": cfg.getboolean("features", "palette", fallback=True),
            "filter_intra": cfg.getboolean("features", "filter_intra", fallback=True),
            "rect_tx": cfg.getboolean("features", "rect_tx", fallback=True),
            "global_motion": cfg.getboolean("features", "global_motion", fallback=True),
            "obmc": cfg.getboolean("features", "obmc", fallback=True),
            "warped_motion": cfg.getboolean("features", "warped_motion", fallback=True),
        }

    # Tools
    result["tools"] = {}
    for section, tool_list in TOOLS.items():
        result["tools"][section] = {}
        if cfg.has_section(section):
            for tool in tool_list:
                result["tools"][section][tool] = cfg.getboolean(
                    section, tool, fallback=True)
        else:
            for tool in tool_list:
                result["tools"][section][tool] = True

    return result


# ============================================================================
# Crop YUV
# ============================================================================

def crop_yuv(src_file, src_w, src_h, crop_w, crop_h, off_x, off_y,
             bit_depth, num_frames):
    """Crop raw YUV using ffmpeg, return path to temp file."""
    pix_fmt = "yuv420p" if bit_depth == 8 else "yuv420p10le"
    out_path = tempfile.mktemp(suffix=".yuv", prefix="crop_")

    cmd = [
        "ffmpeg", "-y",
        "-f", "rawvideo", "-pix_fmt", pix_fmt,
        "-s", f"{src_w}x{src_h}",
        "-i", src_file,
        "-vf", f"crop={crop_w}:{crop_h}:{off_x}:{off_y}",
        "-f", "rawvideo", "-pix_fmt", pix_fmt,
    ]
    if num_frames > 0:
        cmd.extend(["-frames:v", str(num_frames)])
    cmd.append(out_path)

    print(f"Crop: {src_w}x{src_h} -> {crop_w}x{crop_h} @ ({off_x},{off_y})")
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    if result.returncode != 0:
        print(f"CROP FAILED: {result.stderr[:500]}")
        sys.exit(1)
    return out_path


# ============================================================================
# Write tools CSV
# ============================================================================

def write_tools_csv(path, tools_dict):
    """Write encoder tools CSV from tools dict."""
    lines = ["# Auto-generated by encode.py"]
    for section, tool_list in TOOLS.items():
        csv_cat = CSV_CATEGORIES[section]
        for tool in tool_list:
            enabled = tools_dict.get(section, {}).get(tool, True)
            lines.append(f"{csv_cat},{tool},{1 if enabled else 0}")
    Path(path).write_text('\n'.join(lines) + '\n')


# ============================================================================
# Build flags from [filters] and [features]
# ============================================================================

def build_feature_flags(cfg):
    """Build aomenc CLI flags from [filters] and [features]."""
    flags = []

    flt = cfg.get("filters", {})
    if flt:
        if not flt.get("deblocking", True):
            flags.append("--loopfilter-control=0")
        flags.append(f"--enable-cdef={1 if flt.get('cdef', True) else 0}")
        flags.append(f"--enable-restoration={1 if flt.get('loop_restoration', True) else 0}")

    feat = cfg.get("features", {})
    if feat:
        for key, flag in [
            ("angle_delta", "--enable-angle-delta"),
            ("cfl_intra", "--enable-cfl-intra"),
            ("palette", "--enable-palette"),
            ("filter_intra", "--enable-filter-intra"),
            ("rect_tx", "--enable-rect-tx"),
            ("global_motion", "--enable-global-motion"),
            ("obmc", "--enable-obmc"),
            ("warped_motion", "--enable-warped-motion"),
        ]:
            flags.append(f"{flag}={1 if feat.get(key, True) else 0}")

    return ' '.join(flags)


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

def run_encode(input_file, output_file, cfg, width_override=None, height_override=None):
    """Run full pipeline: crop (optional) -> CSV -> aomenc."""

    src_w = width_override or cfg["source"]["width"]
    src_h = height_override or cfg["source"]["height"]
    crop = cfg["crop"]
    enc = cfg["encode"]
    num_frames = cfg["source"]["frames"]
    bit_depth = enc["bit_depth"]

    enc_input = input_file
    enc_w, enc_h = src_w, src_h
    cropped_tmp = None

    # Crop
    if crop["enabled"]:
        cropped_tmp = crop_yuv(
            input_file, src_w, src_h,
            crop["width"], crop["height"],
            crop["offset_x"], crop["offset_y"],
            bit_depth, num_frames)
        enc_input = cropped_tmp
        enc_w, enc_h = crop["width"], crop["height"]

    # Tools CSV
    csv_tmp = tempfile.mktemp(suffix=".csv", prefix="tools_")
    write_tools_csv(csv_tmp, cfg["tools"])

    # Feature flags
    feature_flags = build_feature_flags(cfg)
    extra = enc["extra_flags"]
    if feature_flags:
        extra = f"{feature_flags} {extra}".strip()

    # aomenc command
    cmd = [
        str(ENCODER),
        f"--tools-cfg={csv_tmp}",
        "-w", str(enc_w), "-h", str(enc_h),
        "--fps=30/1",
        f"--cpu-used={enc['cpu_used']}",
        "--end-usage=q", f"--cq-level={enc['cq_level']}",
        f"--kf-max-dist={enc['kf_max_dist']}",
        "--passes=1",
        f"--sb-size={enc['sb_size']}",
        "--enable-flip-idtx=1",
        "--reduced-tx-type-set=0",
        f"--bit-depth={bit_depth}",
        f"--input-bit-depth={bit_depth}",
    ]
    if num_frames > 0:
        cmd.append(f"--limit={num_frames}")
    if extra:
        cmd.extend(extra.split())
    cmd.extend(["-o", str(output_file), str(enc_input)])

    Path(output_file).parent.mkdir(parents=True, exist_ok=True)

    # Summary
    print(f"Input:    {input_file} ({src_w}x{src_h})")
    if crop["enabled"]:
        print(f"Crop:     {crop['width']}x{crop['height']} "
              f"@ ({crop['offset_x']},{crop['offset_y']})")
    print(f"Encode:   {enc_w}x{enc_h}, cq={enc['cq_level']}, "
          f"bd={bit_depth}, kf={enc['kf_max_dist']}")
    print(f"Output:   {output_file}")

    disabled = [t for s in cfg["tools"].values() for t, on in s.items() if not on]
    if disabled:
        print(f"Disabled: {len(disabled)} tools — "
              + ", ".join(disabled[:8])
              + ("..." if len(disabled) > 8 else ""))
    print()

    # Run
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
    if result.stderr:
        print(result.stderr.strip())

    # Cleanup
    Path(csv_tmp).unlink(missing_ok=True)
    if cropped_tmp:
        Path(cropped_tmp).unlink(missing_ok=True)

    if result.returncode != 0:
        print(f"\nENCODE FAILED (exit {result.returncode})")
        return False

    size = Path(output_file).stat().st_size
    print(f"\nOK: {output_file} ({size} bytes)")
    return True


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

def main():
    parser = argparse.ArgumentParser(
        description="AV1 Encoder — INI config based",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  %(prog)s --init config.ini
  %(prog)s -i input.yuv -o output.ivf -cfg config.ini
  %(prog)s -i input.yuv -o output.ivf -cfg config.ini -W 1920 -H 1080
""")
    parser.add_argument("-i", "--input", help="Input YUV file")
    parser.add_argument("-o", "--output", help="Output IVF file")
    parser.add_argument("-cfg", "--config", help="INI config file")
    parser.add_argument("-W", "--width", type=int, default=None,
                        help="Override source width (from INI)")
    parser.add_argument("-H", "--height", type=int, default=None,
                        help="Override source height (from INI)")
    parser.add_argument("-n", "--frames", type=int, default=None,
                        help="Override frame count (from INI)")
    parser.add_argument("--init", metavar="FILE",
                        help="Generate default INI config and exit")

    args = parser.parse_args()

    if args.init:
        generate_default_ini(args.init)
        return

    if not args.input or not args.output or not args.config:
        parser.error("Required: -i input.yuv -o output.ivf -cfg config.ini")

    if not Path(args.input).exists():
        print(f"ERROR: Input not found: {args.input}")
        sys.exit(1)
    if not Path(args.config).exists():
        print(f"ERROR: Config not found: {args.config}")
        sys.exit(1)
    if not ENCODER.exists():
        print(f"ERROR: Encoder not found: {ENCODER}")
        print(f"Build: cd {REF_AOM}/build_enc && cmake --build . -j8")
        sys.exit(1)

    cfg = read_config(args.config)

    # CLI overrides
    if args.frames is not None:
        cfg["source"]["frames"] = args.frames

    ok = run_encode(args.input, args.output, cfg,
                    width_override=args.width, height_override=args.height)
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()
