import gradio as gr
import subprocess
import tempfile
import shutil
import os
import struct

# ─────────────────────────────────────────────
# MP4 ATOM PATCHING (Python port of C# code)
# ─────────────────────────────────────────────

def index_of_sequence(buffer: bytes, pattern: bytes, start: int) -> int:
    if not buffer or not pattern or len(pattern) > len(buffer) - start:
        return -1
    for i in range(start, len(buffer) - len(pattern) + 1):
        if buffer[i:i+len(pattern)] == pattern:
            return i
    return -1


def patch_atom(atom_name: str, data: bytearray, scale_factor: float | None, log: list) -> int:
    patch_count = 0
    start_index = 0
    atom_bytes = atom_name.encode("utf-8")

    while True:
        found_index = index_of_sequence(bytes(data), atom_bytes, start_index)
        if found_index == -1:
            break

        header_offset = found_index - 4
        if header_offset < 0:
            start_index = found_index + 4
            continue

        if header_offset + 4 > len(data):
            start_index = found_index + 4
            continue

        box_size = struct.unpack_from(">I", data, header_offset)[0]
        if box_size < 8:
            start_index = found_index + 4
            continue

        if header_offset + 8 >= len(data):
            start_index = found_index + 4
            continue

        version = data[header_offset + 8]

        if version == 0:
            ts_off, dur_off = header_offset + 20, header_offset + 24
            if dur_off + 4 > len(data):
                start_index = found_index + 4
                continue
            old_ts  = struct.unpack_from(">I", data, ts_off)[0]
            old_dur = struct.unpack_from(">I", data, dur_off)[0]
            chosen  = scale_factor if scale_factor is not None else (30000.0 / old_ts)
            new_ts  = int(old_ts  * chosen)
            new_dur = int(old_dur * chosen)
            struct.pack_into(">I", data, ts_off,  new_ts)
            struct.pack_into(">I", data, dur_off, new_dur)
            log.append(f"  ✔ Patched '{atom_name}' (v0): timescale {old_ts}→{new_ts}, duration {old_dur}→{new_dur}")
            patch_count += 1

        elif version == 1:
            ts_off, dur_off = header_offset + 28, header_offset + 32
            if dur_off + 8 > len(data):
                start_index = found_index + 4
                continue
            old_ts  = struct.unpack_from(">I", data, ts_off)[0]
            old_dur = struct.unpack_from(">Q", data, dur_off)[0]
            chosen  = scale_factor if scale_factor is not None else (30000.0 / old_ts)
            new_ts  = int(old_ts  * chosen)
            new_dur = int(old_dur * chosen)
            struct.pack_into(">I", data, ts_off,  new_ts)
            struct.pack_into(">Q", data, dur_off, new_dur)
            log.append(f"  ✔ Patched '{atom_name}' (v1): timescale {old_ts}→{new_ts}, duration {old_dur}→{new_dur}")
            patch_count += 1
        else:
            log.append(f"  ⚠ '{atom_name}' (v{version}) skipped — unknown version")

        start_index = found_index + 4

    return patch_count


def patch_mp4(input_path: str, output_path: str, scale_factor: float | None) -> list[str]:
    log = []
    with open(input_path, "rb") as f:
        data = bytearray(f.read())
    log.append(f"📂 Read {len(data):,} bytes")

    n_mvhd = patch_atom("mvhd", data, scale_factor, log)
    n_mdhd = patch_atom("mdhd", data, scale_factor, log)
    total  = n_mvhd + n_mdhd

    if total == 0:
        log.append("⚠ No atoms patched — file may be unchanged")
    else:
        log.append(f"✅ Total atoms patched: {total}")

    with open(output_path, "wb") as f:
        f.write(data)
    log.append(f"💾 Saved to {output_path}")
    return log


# ─────────────────────────────────────────────
# FFMPEG HELPERS
# ─────────────────────────────────────────────

def ffmpeg_available() -> bool:
    return shutil.which("ffmpeg") is not None


def detect_fps(input_path: str) -> float:
    try:
        result = subprocess.run(
            ["ffprobe", "-v", "error", "-select_streams", "v:0",
             "-show_entries", "stream=r_frame_rate",
             "-of", "default=noprint_wrappers=1:nokey=1", input_path],
            capture_output=True, text=True, timeout=30
        )
        num, den = result.stdout.strip().split("/")
        return round(float(num) / float(den), 2)
    except Exception:
        return 0.0


def get_video_info(input_path: str) -> dict:
    try:
        import json
        result = subprocess.run(
            ["ffprobe", "-v", "error", "-print_format", "json",
             "-show_streams", "-show_format", input_path],
            capture_output=True, text=True, timeout=30
        )
        data   = json.loads(result.stdout)
        fmt    = data.get("format", {})
        vstream = next((s for s in data.get("streams", []) if s.get("codec_type") == "video"), {})
        return {
            "duration": float(fmt.get("duration", 0)),
            "size_mb":  round(os.path.getsize(input_path) / 1e6, 2),
            "width":    vstream.get("width", "?"),
            "height":   vstream.get("height", "?"),
            "codec":    vstream.get("codec_name", "?"),
            "fps":      vstream.get("r_frame_rate", "?"),
        }
    except Exception:
        return {}


def reencode_video(input_path: str, output_path: str,
                   crf: int, fps_out: float,
                   use_nvenc: bool, log: list) -> bool:
    codec  = "h264_nvenc" if use_nvenc else "libx264"
    preset = "p4"         if use_nvenc else "slow"

    cmd = [
        "ffmpeg", "-y", "-i", input_path,
        "-c:v", codec,
        "-preset", preset,
        "-r", str(fps_out),
        "-pix_fmt", "yuv420p",
        "-movflags", "+faststart",
    ]
    if use_nvenc:
        cmd += ["-cq", str(crf)]
    else:
        cmd += ["-crf", str(crf)]

    cmd += ["-c:a", "aac", "-b:a", "192k", output_path]
    log.append(f"▶ FFmpeg: {' '.join(cmd)}")

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=300)
        if result.returncode != 0:
            log.append(f"❌ FFmpeg error:\n{result.stderr[-800:]}")
            return False
        log.append("✅ FFmpeg encoding complete")
        return True
    except subprocess.TimeoutExpired:
        log.append("❌ FFmpeg timed out")
        return False


# ─────────────────────────────────────────────
# MAIN PIPELINE
# ─────────────────────────────────────────────

FPS_SCALE = {
    "120fps → 60fps": (120, 0.5),
    "120fps → 30fps": (120, 0.25),
    "90fps → 60fps":  (90,  0.5),
    "90fps → 30fps":  (90,  0.25),
    "60fps → 60fps":  (60,  1.0),
    "60fps → 30fps":  (60,  0.5),
    "30fps → 30fps":  (30,  1.0),
    "Auto (target 30k timescale)": (None, None),
}


def process_video(
    video_file,
    fps_mode: str,
    crf: int,
    use_nvenc: bool,
    enable_reencode: bool,
    enable_patch: bool,
):
    if video_file is None:
        return None, "⚠ Please upload a video first."

    input_path = video_file
    logs: list[str] = []
    logs.append("═" * 48)
    logs.append("🎬 TikTok HD Video Compressor — Starting...")
    logs.append("═" * 48)

    if not ffmpeg_available():
        logs.append("⚠ ffmpeg / ffprobe not found in PATH — skipping encode step")
        enable_reencode = False

    info = get_video_info(input_path)
    if info:
        logs.append(f"📊 Input  : {info.get('width')}×{info.get('height')}  "
                    f"{info.get('fps')} fps  {info.get('codec')}  "
                    f"{info.get('size_mb')} MB  {info.get('duration'):.1f}s")

    tmp_dir = tempfile.mkdtemp()
    try:
        step1 = os.path.join(tmp_dir, "step1_encoded.mp4")
        step2 = os.path.join(tmp_dir, "step2_patched.mp4")

        # Determine output FPS from fps_mode
        _, scale_factor = FPS_SCALE.get(fps_mode, (None, None))
        fps_map = {
            "120fps → 60fps": 60,
            "120fps → 30fps": 30,
            "90fps → 60fps":  60,
            "90fps → 30fps":  30,
            "60fps → 60fps":  60,
            "60fps → 30fps":  30,
            "30fps → 30fps":  30,
            "Auto (target 30k timescale)": 30,
        }
        fps_out = fps_map.get(fps_mode, 30)

        # ── Step 1: Re-encode with FFmpeg ──────────────────
        if enable_reencode:
            logs.append("\n── Step 1: FFmpeg Re-encode ──────────────────")
            ok = reencode_video(input_path, step1, crf, fps_out, use_nvenc, logs)
            if not ok:
                return None, "\n".join(logs)
            encoded_path = step1
        else:
            logs.append("\n── Step 1: Re-encode skipped ─────────────────")
            encoded_path = input_path

        # ── Step 2: Atom Patching ──────────────────────────
        if enable_patch:
            logs.append("\n── Step 2: MP4 Atom Patching ─────────────────")
            patch_logs = patch_mp4(encoded_path, step2, scale_factor)
            logs.extend(patch_logs)
            final_path = step2
        else:
            logs.append("\n── Step 2: Atom patch skipped ────────────────")
            final_path = encoded_path

        # ── Final output ───────────────────────────────────
        out_dir  = tempfile.mkdtemp()
        out_path = os.path.join(out_dir, "tiktok_hd_output.mp4")
        shutil.copy2(final_path, out_path)

        out_info = get_video_info(out_path)
        if out_info:
            logs.append(f"\n📊 Output : {out_info.get('width')}×{out_info.get('height')}  "
                        f"{out_info.get('fps')} fps  {out_info.get('codec')}  "
                        f"{out_info.get('size_mb')} MB")

        size_in  = info.get("size_mb", 0)
        size_out = out_info.get("size_mb", 0) if out_info else 0
        if size_in and size_out:
            pct = round((1 - size_out / size_in) * 100, 1)
            logs.append(f"📉 Size reduction: {size_in} MB → {size_out} MB  ({pct}% smaller)")

        logs.append("\n✅ Done! Download your TikTok-ready video below.")
        return out_path, "\n".join(logs)

    except Exception as e:
        logs.append(f"\n❌ Unexpected error: {e}")
        return None, "\n".join(logs)
    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)


# ─────────────────────────────────────────────
# GRADIO UI
# ─────────────────────────────────────────────

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&family=DM+Sans:wght@300;400;600&display=swap');

:root {
    --bg: #0a0a0f;
    --surface: #12121a;
    --surface2: #1c1c28;
    --accent: #ff2d55;
    --accent2: #0af;
    --text: #f0f0f8;
    --muted: #888;
    --border: #2a2a3a;
    --radius: 12px;
    --mono: 'Space Mono', monospace;
    --sans: 'DM Sans', sans-serif;
}

body, .gradio-container {
    background: var(--bg) !important;
    font-family: var(--sans) !important;
    color: var(--text) !important;
}

.gradio-container { max-width: 860px !important; margin: 0 auto !important; }

/* Header */
.app-header {
    text-align: center;
    padding: 2.5rem 1rem 1.5rem;
    border-bottom: 1px solid var(--border);
    margin-bottom: 2rem;
}
.app-header h1 {
    font-family: var(--mono) !important;
    font-size: 2rem;
    letter-spacing: -0.02em;
    color: var(--text);
    margin: 0 0 .4rem;
}
.app-header h1 span { color: var(--accent); }
.app-header p { color: var(--muted); font-size: .9rem; margin: 0; }

/* Panels */
.panel {
    background: var(--surface) !important;
    border: 1px solid var(--border) !important;
    border-radius: var(--radius) !important;
    padding: 1.25rem !important;
    margin-bottom: 1rem !important;
}
.panel-title {
    font-family: var(--mono);
    font-size: .7rem;
    letter-spacing: .12em;
    text-transform: uppercase;
    color: var(--accent2);
    margin-bottom: .75rem;
}

/* Inputs */
label span, .gr-form label span {
    font-family: var(--sans) !important;
    font-size: .8rem !important;
    color: var(--muted) !important;
    font-weight: 600 !important;
    letter-spacing: .04em !important;
    text-transform: uppercase !important;
}
input, select, textarea, .gr-input, .gr-dropdown {
    background: var(--surface2) !important;
    border: 1px solid var(--border) !important;
    color: var(--text) !important;
    border-radius: 8px !important;
    font-family: var(--sans) !important;
}

/* Buttons */
button.primary, .gr-button-primary {
    background: var(--accent) !important;
    border: none !important;
    color: #fff !important;
    font-family: var(--mono) !important;
    font-size: .85rem !important;
    letter-spacing: .05em !important;
    border-radius: 8px !important;
    padding: .7rem 1.5rem !important;
    transition: opacity .2s !important;
}
button.primary:hover { opacity: .85 !important; }

/* Log box */
.log-box textarea {
    font-family: var(--mono) !important;
    font-size: .75rem !important;
    background: #08080e !important;
    color: #7dff8a !important;
    border: 1px solid var(--border) !important;
    border-radius: 8px !important;
    line-height: 1.7 !important;
}

/* Info badge */
.info-badge {
    display: inline-flex;
    gap: .4rem;
    align-items: center;
    background: var(--surface2);
    border: 1px solid var(--border);
    border-radius: 6px;
    padding: .3rem .7rem;
    font-family: var(--mono);
    font-size: .7rem;
    color: var(--muted);
    margin-right: .5rem;
    margin-bottom: .5rem;
}
"""

HEADER_HTML = """
<div class="app-header">
  <h1>🎬 TikTok<span>HD</span> Compressor</h1>
  <p>MP4 Atom Patching · FFmpeg H.264 · Smooth FPS for TikTok upload</p>
</div>
"""

INFO_HTML = """
<div style="display:flex;flex-wrap:wrap;gap:.5rem;margin-bottom:1rem">
  <span class="info-badge">⚡ MP4 Atom Patch</span>
  <span class="info-badge">🎞 FFmpeg H.264</span>
  <span class="info-badge">🎯 TikTok-optimised</span>
  <span class="info-badge">🖥 CPU + NVIDIA GPU</span>
</div>
"""

FPS_CHOICES = list(FPS_SCALE.keys())

with gr.Blocks(css=CSS, title="TikTok HD Compressor") as demo:

    gr.HTML(HEADER_HTML)
    gr.HTML(INFO_HTML)

    with gr.Row():
        with gr.Column(scale=1):

            gr.HTML('<div class="panel-title">📁 Input Video</div>')
            video_input = gr.Video(label="Upload MP4 / MOV / AVI", sources=["upload"])

            gr.HTML('<div class="panel-title" style="margin-top:1.25rem">⚙ Encoding Settings</div>')

            fps_mode = gr.Dropdown(
                label="FPS Mode (Atom Patch Scale)",
                choices=FPS_CHOICES,
                value="60fps → 30fps",
            )

            crf_slider = gr.Slider(
                label="CRF Quality (lower = better, larger file)",
                minimum=15, maximum=35, value=23, step=1,
            )

            with gr.Row():
                enable_reencode = gr.Checkbox(label="Re-encode with FFmpeg", value=True)
                enable_patch    = gr.Checkbox(label="MP4 Atom Patch",        value=True)
                use_nvenc       = gr.Checkbox(label="NVIDIA GPU (h264_nvenc)", value=False)

            run_btn = gr.Button("▶ Compress for TikTok", variant="primary", size="lg")

        with gr.Column(scale=1):
            gr.HTML('<div class="panel-title">📤 Output</div>')
            video_output = gr.Video(label="Processed Video")

            gr.HTML('<div class="panel-title" style="margin-top:1.25rem">📋 Log</div>')
            log_output = gr.Textbox(
                label="",
                lines=18,
                interactive=False,
                elem_classes=["log-box"],
                placeholder="Logs will appear here after processing...",
            )

    gr.HTML("""
    <div style="margin-top:1.5rem;border-top:1px solid #2a2a3a;padding-top:1rem;
                text-align:center;font-size:.72rem;color:#555;font-family:'Space Mono',monospace;">
      MP4 Atom Patching based on Editing News / Evilfy (GPLv3) ·
      FFmpeg licensed under LGPL/GPL ·
      Built for TikTok HD upload optimisation
    </div>
    """)

    run_btn.click(
        fn=process_video,
        inputs=[video_input, fps_mode, crf_slider, use_nvenc, enable_reencode, enable_patch],
        outputs=[video_output, log_output],
    )

if __name__ == "__main__":
    demo.launch()