#!/usr/bin/env python3
"""
Web UI server for enc_av1 — configure tools, run encoding, view results.
Usage: python3 server.py [--port 8080]
"""

import configparser
import json
import os
import subprocess
import sys
import tempfile
import threading
import time
from http.server import HTTPServer, SimpleHTTPRequestHandler
from pathlib import Path
from urllib.parse import parse_qs, urlparse

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

# ============================================================================
# Encoding job state
# ============================================================================

job_lock = threading.Lock()
current_job = {
    "running": False,
    "log": [],
    "result": None,
    "ivf_path": None,
    "trace": None,
}


def list_sequences():
    """List available YUV sequences with metadata."""
    seqs = []
    if SEQ_DIR.exists():
        for f in sorted(SEQ_DIR.glob("*.yuv")):
            name = f.stem
            # Parse dimensions from filename: test_WxH_desc.yuv
            parts = name.split("_")
            w, h = 256, 256
            for p in parts:
                if "x" in p:
                    try:
                        wh = p.split("x")
                        w, h = int(wh[0]), int(wh[1])
                        break
                    except (ValueError, IndexError):
                        pass
            frames = max(1, f.stat().st_size // (w * h * 3 // 2))
            bit_depth = 10 if "10b" in name else 8
            if bit_depth == 10:
                frames = max(1, f.stat().st_size // (w * h * 3))
            seqs.append({
                "path": str(f),
                "name": name,
                "width": w,
                "height": h,
                "frames": frames,
                "bit_depth": bit_depth,
                "size": f.stat().st_size,
            })
    return seqs


def list_presets():
    """List available INI presets."""
    presets = []
    for f in sorted(SCRIPT_DIR.glob("B*.ini")):
        presets.append({"name": f.stem, "path": str(f)})
    return presets


def load_preset(path):
    """Load INI preset into dict."""
    cfg = configparser.ConfigParser()
    cfg.read(path)
    result = {}
    for section in cfg.sections():
        result[section] = {}
        for key, val in cfg.items(section):
            result[section][key] = val
    return result


def run_encode_job(config_dict, input_path, job_name):
    """Run encoding in background thread."""
    global current_job

    def log(msg):
        with job_lock:
            current_job["log"].append(msg)

    try:
        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
        out_ivf = OUTPUT_DIR / f"{job_name}.ivf"

        # Write temp INI
        ini_tmp = OUTPUT_DIR / f"{job_name}.ini"
        cfg = configparser.ConfigParser()
        for section, items in config_dict.items():
            cfg[section] = items
        with open(ini_tmp, "w") as f:
            cfg.write(f)

        log(f"Config: {ini_tmp}")
        log(f"Input:  {input_path}")
        log(f"Output: {out_ivf}")

        # Count disabled tools
        disabled = []
        for section in ["partitions", "intra_modes", "inter_modes", "inter_refs",
                        "compound_types", "tx_sizes", "tx_types"]:
            if section in config_dict:
                for k, v in config_dict[section].items():
                    if v.lower() == "false":
                        disabled.append(k)
        if disabled:
            log(f"Disabled: {len(disabled)} tools")

        # Get frames from config
        frames = config_dict.get("source", {}).get("frames", "0")

        # Build encode command
        cmd = [
            sys.executable, str(ENCODE_PY),
            "-i", str(input_path),
            "-o", str(out_ivf),
            "-cfg", str(ini_tmp),
        ]
        if int(frames) > 0:
            cmd.extend(["-n", str(frames)])

        log(f"\n--- ENCODING ---")
        t0 = time.time()
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
        elapsed = time.time() - t0

        if result.stdout:
            for line in result.stdout.strip().split("\n"):
                log(line)
        if result.stderr:
            for line in result.stderr.strip().split("\n"):
                log(line)

        if result.returncode != 0:
            log(f"\nENCODE FAILED (exit {result.returncode})")
            with job_lock:
                current_job["result"] = "error"
                current_job["running"] = False
            return

        size = out_ivf.stat().st_size
        log(f"\nEncoded OK: {size} bytes in {elapsed:.1f}s")

        # Decode with trace
        if DECODER.exists():
            trace_dir = OUTPUT_DIR / f"{job_name}_trace"
            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)

            log(f"\n--- DECODING WITH TRACE ---")
            dec_cmd = [
                str(DECODER), "--noblit",
                f"--out-cdq={cdq_dir}",
                f"--out-msac={msac_dir}",
                str(out_ivf),
            ]
            dec_result = subprocess.run(dec_cmd, capture_output=True, text=True, timeout=120)
            if dec_result.returncode == 0:
                log("Decode OK")
                # Parse trace summary
                trace_summary = parse_trace_summary(msac_dir)
                with job_lock:
                    current_job["trace"] = trace_summary
                log(f"\n--- TRACE SUMMARY ---")
                for k, v in trace_summary.items():
                    if v:
                        log(f"  {k}: {v}")
            else:
                log(f"Decode failed: {dec_result.stderr[:200]}")
        else:
            log("(decoder not found, skipping trace)")

        with job_lock:
            current_job["result"] = "ok"
            current_job["ivf_path"] = str(out_ivf)

    except Exception as e:
        log(f"ERROR: {e}")
        with job_lock:
            current_job["result"] = "error"
    finally:
        with job_lock:
            current_job["running"] = False
        log("--- DONE ---")


def parse_trace_summary(msac_dir):
    """Quick parse of MSAC trace for summary."""
    from collections import Counter
    partition_vals = Counter()
    intra_vals = Counter()
    tx_type_vals = Counter()
    inter_mode_vals = Counter()
    symbols = Counter()

    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_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",
    }
    TX_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",
    }

    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)
                symbols[name] += 1
                if name == "partition":
                    partition_vals[value] += 1
                elif name == "intra_mode":
                    intra_vals[value] += 1
                elif name == "tx_type_resolved":
                    tx_type_vals[value] += 1
                elif name == "inter_mode":
                    inter_mode_vals[value] += 1

    return {
        "partitions": ", ".join(f"{PARTITION_ENUM.get(v, f'?{v}')}({c})"
                                for v, c in sorted(partition_vals.items())),
        "intra_modes": ", ".join(f"{INTRA_ENUM.get(v, f'?{v}')}({c})"
                                 for v, c in sorted(intra_vals.items())),
        "tx_types": ", ".join(f"{TX_ENUM.get(v, f'?{v}')}({c})"
                              for v, c in sorted(tx_type_vals.items())),
        "angle_delta": str(symbols.get("angle_delta", 0)),
        "cfl_signs": str(symbols.get("cfl:signs", 0)),
        "filter_intra": str(symbols.get("filter_intra_mode", 0)),
        "total_symbols": str(sum(symbols.values())),
    }


# ============================================================================
# HTTP Handler
# ============================================================================

class Handler(SimpleHTTPRequestHandler):
    def do_GET(self):
        parsed = urlparse(self.path)
        path = parsed.path

        if path == "/" or path == "/index.html":
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.end_headers()
            html_path = SCRIPT_DIR / "web_ui.html"
            self.wfile.write(html_path.read_bytes())

        elif path == "/api/sequences":
            self.json_response(list_sequences())

        elif path == "/api/presets":
            self.json_response(list_presets())

        elif path == "/api/preset":
            qs = parse_qs(parsed.query)
            preset_path = qs.get("path", [None])[0]
            if preset_path and Path(preset_path).exists():
                self.json_response(load_preset(preset_path))
            else:
                self.json_response({"error": "not found"}, 404)

        elif path == "/api/status":
            with job_lock:
                self.json_response({
                    "running": current_job["running"],
                    "log": current_job["log"],
                    "result": current_job["result"],
                    "trace": current_job["trace"],
                    "ivf_path": current_job["ivf_path"],
                })

        else:
            self.send_error(404)

    def do_POST(self):
        parsed = urlparse(self.path)

        if parsed.path == "/api/encode":
            content_len = int(self.headers.get("Content-Length", 0))
            body = json.loads(self.rfile.read(content_len))

            with job_lock:
                if current_job["running"]:
                    self.json_response({"error": "encoding in progress"}, 409)
                    return
                current_job["running"] = True
                current_job["log"] = []
                current_job["result"] = None
                current_job["trace"] = None
                current_job["ivf_path"] = None

            config = body.get("config", {})
            input_path = body.get("input_path", "")
            job_name = body.get("name", f"encode_{int(time.time())}")

            t = threading.Thread(target=run_encode_job,
                                 args=(config, input_path, job_name),
                                 daemon=True)
            t.start()
            self.json_response({"status": "started", "name": job_name})

        else:
            self.send_error(404)

    def json_response(self, data, code=200):
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(json.dumps(data).encode())

    def log_message(self, format, *args):
        pass  # suppress request logs


def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--port", type=int, default=8080)
    args = parser.parse_args()

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    server = HTTPServer(("0.0.0.0", args.port), Handler)
    print(f"enc_av1 Web UI: http://localhost:{args.port}")
    print(f"Output dir:     {OUTPUT_DIR}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down.")


if __name__ == "__main__":
    main()
