import gradio as gr
import os
import subprocess
import uuid
import zipfile
from datetime import datetime
import pandas as pd

# Ensure output directory exists
OUTPUT_DIR = "outputs"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ---------------------------
# Helpers & Configuration
# ---------------------------

def map_level_to_crf(level: str, hevc: bool) -> int:
    """
    Map user-friendly compression levels to CRF values.
    HEVC often achieves similar quality at slightly higher CRF than H.264,
    but we keep scales aligned for simplicity.
    """
    level = (level or "Medium").strip().lower()
    if level == "low":        # Best quality, larger files
        return 22 if not hevc else 24
    elif level == "high":     # Smallest files, lower quality
        return 32 if not hevc else 34
    else:                     # Medium
        return 28 if not hevc else 30

def ffmpeg_available() -> bool:
    try:
        subprocess.run(["ffmpeg", "-version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True)
        return True
    except Exception:
        return False

def resolution_to_height(res_choice: str) -> int | None:
    """
    Map UI selection to target height for downscaling.
    Using '-vf scale=-2:HEIGHT' ensures even width while preserving aspect.
    """
    if not res_choice or "No downscaling" in res_choice:
        return None
    table = {
        "720p (1280x720)": 720,
        "1080p (1920x1080)": 1080,
        "1440p (2560x1440)": 1440,
        "4K (3840x2160)": 2160,
    }
    return table.get(res_choice)

# ---------------------------
# Core compression
# ---------------------------

def compress_single_video(
    infile_path: str,
    use_hevc: bool,
    crf: int,
    preset: str = "fast",
    audio_bitrate_kbps: int = 128,
    target_height: int | None = None,
    progress: gr.Progress | None = None,
    fallback_to_h264_on_error: bool = True
):
    """
    Compress a single MP4 using FFmpeg.
    - Video codec: H.264 (libx264) or H.265 (libx265)
    - Optional downscale using '-vf scale=-2:HEIGHT'
    Returns: (output_path, original_size_bytes, compressed_size_bytes, codec_used)
    """
    base = os.path.basename(infile_path)
    out_name = f"{os.path.splitext(base)[0]}_compressed_{uuid.uuid4().hex[:8]}.mp4"
    out_path = os.path.join(OUTPUT_DIR, out_name)

    original_size = os.path.getsize(infile_path)

    # Build common parts
    vf_args = []
    if target_height:
        # Limit maximal height; width auto-calculated as even number to maintain aspect ratio
        vf_args = ["-vf", f"scale=-2:{target_height}"]

    vcodec = "libx265" if use_hevc else "libx264"
    codec_used = "H.265 (libx265)" if use_hevc else "H.264 (libx264)"

    # Base ffmpeg command
    cmd = [
        "ffmpeg", "-y",
        "-i", infile_path,
        *(vf_args),
        "-vcodec", vcodec,
        "-crf", str(crf),
        "-preset", preset,
        # Improve playback compatibility
        "-pix_fmt", "yuv420p",
        # AAC audio with selected bitrate
        "-acodec", "aac", "-b:a", f"{audio_bitrate_kbps}k",
        # Better streaming behavior
        "-movflags", "+faststart",
    ]

    # Tag HEVC properly for Apple players
    if use_hevc:
        cmd += ["-tag:v", "hvc1"]

    cmd += [out_path]

    if progress:
        progress(0.2, desc=f"Compressing {base} with {codec_used} (CRF {crf}, preset {preset})…")

    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if proc.returncode != 0 and use_hevc and fallback_to_h264_on_error:
        # Fallback to H.264 if HEVC encoder isn't available or failed
        vcodec = "libx264"
        codec_used = "H.264 (libx264)"
        fallback_cmd = [
            "ffmpeg", "-y",
            "-i", infile_path,
            *(vf_args),
            "-vcodec", vcodec,
            "-crf", str(crf),
            "-preset", preset,
            "-pix_fmt", "yuv420p",
            "-acodec", "aac", "-b:a", f"{audio_bitrate_kbps}k",
            "-movflags", "+faststart",
            out_path
        ]
        proc2 = subprocess.run(fallback_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if proc2.returncode != 0:
            raise RuntimeError(
                f"FFmpeg failed on {base}.\n\n"
                f"HEVC STDERR:\n{proc.stderr.decode(errors='ignore')}\n\n"
                f"H.264 STDERR:\n{proc2.stderr.decode(errors='ignore')}"
            )

    elif proc.returncode != 0:
        raise RuntimeError(
            f"FFmpeg failed on {base}.\n\nSTDERR:\n{proc.stderr.decode(errors='ignore')}\n\n"
            f"STDOUT:\n{proc.stdout.decode(errors='ignore')}"
        )

    compressed_size = os.path.getsize(out_path)
    if progress:
        progress(0.95, desc=f"Finalizing {base}")

    return out_path, original_size, compressed_size, codec_used

def validate_inputs(files):
    if not files or len(files) == 0:
        raise gr.Error("Please upload at least one MP4 file.")
    paths = []
    for f in files:
        path = getattr(f, "name", None) or (isinstance(f, dict) and f.get("name")) or None
        if not path:
            raise gr.Error("Could not read uploaded file path.")
        if not path.lower().endswith(".mp4"):
            raise gr.Error(f"Unsupported file type: {os.path.basename(path)}. Please upload .mp4 only.")
        paths.append(path)
    return paths

def compress_videos(
    files,
    level: str,
    use_advanced: bool,
    crf_value: int,
    preset: str,
    audio_bitrate_kbps: int,
    use_hevc: bool,
    max_resolution_choice: str,
    progress: gr.Progress = gr.Progress()
):
    """
    Multi-file compression with:
    - H.264 vs H.265 toggle
    - Optional downscaling
    - Progress & summary report
    """
    if not ffmpeg_available():
        raise gr.Error("FFmpeg is not installed. Add it in packages.txt.")

    input_paths = validate_inputs(files)
    total_files = len(input_paths)
    progress(0, desc=f"Starting compression for {total_files} file(s)…")

    target_height = resolution_to_height(max_resolution_choice)
    crf = crf_value if use_advanced else map_level_to_crf(level, hevc=use_hevc)
    preset = preset or "fast"
    audio_bitrate_kbps = max(64, int(audio_bitrate_kbps or 128))

    records = []
    compressed_paths = []

    for idx, in_path in enumerate(input_paths, start=1):
        base = os.path.basename(in_path)
        try:
            progress((idx - 1) / total_files, desc=f"[{idx}/{total_files}] Processing {base}")
            out_path, orig_size, comp_size, codec_used = compress_single_video(
                infile_path=in_path,
                use_hevc=use_hevc,
                crf=crf,
                preset=preset,
                audio_bitrate_kbps=audio_bitrate_kbps,
                target_height=target_height,
                progress=progress
            )
            reduction_bytes = orig_size - comp_size
            reduction_pct = (reduction_bytes / orig_size * 100.0) if orig_size > 0 else 0.0
            records.append({
                "File": base,
                "Codec": codec_used,
                "Downscale": max_resolution_choice if target_height else "None",
                "Original Size (MB)": round(orig_size / (1024 * 1024), 3),
                "Compressed Size (MB)": round(comp_size / (1024 * 1024), 3),
                "Size Reduced (MB)": round(reduction_bytes / (1024 * 1024), 3),
                "Reduction (%)": round(reduction_pct, 2),
                "Output": os.path.basename(out_path)
            })
            compressed_paths.append(out_path)
            progress(idx / total_files, desc=f"[{idx}/{total_files}] Done {base} (↓ {round(reduction_pct,2)}%)")
        except Exception as e:
            records.append({
                "File": base,
                "Codec": "—",
                "Downscale": max_resolution_choice if target_height else "None",
                "Original Size (MB)": None,
                "Compressed Size (MB)": None,
                "Size Reduced (MB)": None,
                "Reduction (%)": None,
                "Output": f"ERROR: {str(e).splitlines()[0][:180]}"
            })

    df = pd.DataFrame.from_records(records)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    report_name = f"compression_report_{timestamp}.csv"
    report_path = os.path.join(OUTPUT_DIR, report_name)
    df.to_csv(report_path, index=False)

    zip_name = f"compressed_bundle_{timestamp}.zip"
    zip_path = os.path.join(OUTPUT_DIR, zip_name)
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        zf.write(report_path, arcname=report_name)
        for p in compressed_paths:
            zf.write(p, arcname=os.path.basename(p))

    progress(1.0, desc="All done! Download your results below.")
    return zip_path, df, compressed_paths

# ---------------------------
# Gradio UI
# ---------------------------

with gr.Blocks(title = "BG CoRe.Vi MP4 (Multi-file)", css=""".summary-table {max-height: 280px; overflow-y: auto;}""") as demo:
    gr.Markdown("# BG CoRe.Vi MP4 (H.264 / H.265, Downscaling, Multi-file)")
    gr.Markdown(
        "Upload one or more MP4s, choose compression level or advanced settings, optionally downscale, and download the compressed files.\n"
        "Includes a progress bar and a CSV summary (sizes & reduction)."
    )

    with gr.Row():
        with gr.Column():
            files_in = gr.Files(
                label="Upload MP4 Videos",
                file_count="multiple",
                file_types=[".mp4"]
            )

            level = gr.Radio(
                label="Compression Level",
                choices=["Low", "Medium", "High"],
                value="Medium",
                info="Low = best quality, High = smallest files"
            )

            # Codec & downscaling
            use_hevc = gr.Checkbox(label="Use HEVC (H.265)", value=False,
                                   info="Smaller files, slower to encode")
            max_res = gr.Dropdown(
                label="Optional Downscaling",
                choices=[
                    "No downscaling",
                    "720p (1280x720)",
                    "1080p (1920x1080)",
                    "1440p (2560x1440)",
                    "4K (3840x2160)"
                ],
                value="No downscaling"
            )

            with gr.Accordion("Advanced Options (optional)", open=False):
                use_adv = gr.Checkbox(label="Use advanced settings", value=False)
                crf_val = gr.Slider(minimum=18, maximum=36, step=1, value=28, label="CRF (Quality/Compression)")
                preset = gr.Dropdown(
                    label="FFmpeg Preset",
                    choices=["ultrafast", "superfast", "veryfast", "faster", "fast", "medium", "slow", "slower", "veryslow"],
                    value="fast",
                    allow_custom_value=False
                )
                audio_bitrate = gr.Slider(
                    minimum=64, maximum=320, step=16, value=128, label="Audio Bitrate (kbps)"
                )

            run_btn = gr.Button("🔄 Compress", variant="primary")

        with gr.Column():
            zip_out = gr.File(label="Download ZIP (compressed videos + CSV report)")
            summary = gr.Dataframe(label="Compression Summary", interactive=False, wrap=True, elem_classes="summary-table")
            files_out = gr.Files(label="Download Individual Compressed Files")

    run_btn.click(
        fn=compress_videos,
        inputs=[files_in, level, use_adv, crf_val, preset, audio_bitrate, use_hevc, max_res],
        outputs=[zip_out, summary, files_out]
    )

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