#!/usr/bin/env python3

import argparse
import json
import os
import shutil
import subprocess
import sys
import tempfile
import time

import cv2
import numpy as np
import openvino as ov


SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
FFMPEG_BIN = os.path.join(SCRIPT_DIR, "bin", "ffmpeg")
FFPROBE_BIN = os.path.join(SCRIPT_DIR, "bin", "ffprobe")
MODEL_XML = os.path.join(SCRIPT_DIR, "model", "realesr-general-x4v3_480x270.xml")
MODEL_WIDTH = 480
MODEL_HEIGHT = 270
SCALE = 4
OUTPUT_WIDTH = MODEL_WIDTH * SCALE
OUTPUT_HEIGHT = MODEL_HEIGHT * SCALE
ENCODE_PRESET = "medium"
ENCODE_CRF = "15"


def get_video_info(video_path):
    cmd = [
        FFPROBE_BIN,
        "-v",
        "quiet",
        "-print_format",
        "json",
        "-show_streams",
        "-show_format",
        video_path,
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"ffprobe failed: {result.stderr}")

    info = json.loads(result.stdout)
    video_stream = next((stream for stream in info["streams"] if stream["codec_type"] == "video"), None)
    if video_stream is None:
        raise RuntimeError("No video stream found")

    width = int(video_stream["width"])
    height = int(video_stream["height"])
    r_frame_rate = video_stream.get("r_frame_rate", "30/1")
    num, den = map(int, r_frame_rate.split("/"))
    fps = num / den
    has_audio = any(stream["codec_type"] == "audio" for stream in info["streams"])
    return width, height, fps, has_audio


def main():
    parser = argparse.ArgumentParser(description="Experiment-faithful Real-ESRGAN general-x4v3 1080p video pipeline")
    parser.add_argument("input", help="Input video path")
    parser.add_argument("-o", "--output", help="Output video path")
    args = parser.parse_args()

    input_video = args.input
    if not os.path.isfile(input_video):
        print(f"Error: input video not found: {input_video}")
        sys.exit(1)

    output_video = args.output or f"{os.path.splitext(input_video)[0]}_smart_v2_{OUTPUT_WIDTH}x{OUTPUT_HEIGHT}.mp4"

    if not os.path.isfile(FFMPEG_BIN) or not os.path.isfile(FFPROBE_BIN):
        print("Error: ffmpeg/ffprobe not found in repo bin/. Run install.sh first.")
        sys.exit(1)
    if not os.path.isfile(MODEL_XML):
        print(f"Error: model not found: {MODEL_XML}")
        sys.exit(1)

    print(f"Analyzing input: {input_video}")
    width, height, fps, has_audio = get_video_info(input_video)
    print(f"Source: {width}x{height}, {fps:.2f} fps, audio: {'yes' if has_audio else 'no'}")
    if width != MODEL_WIDTH or height != MODEL_HEIGHT:
        print(f"Pre-scale: in-memory lanczos -> {MODEL_WIDTH}x{MODEL_HEIGHT}")
    else:
        print(f"Pre-scale: skipped, model input stays at {MODEL_WIDTH}x{MODEL_HEIGHT}")
    print(f"Input:  {MODEL_WIDTH}x{MODEL_HEIGHT}")
    print(f"Output: {OUTPUT_WIDTH}x{OUTPUT_HEIGHT}")
    print(f"Loading OpenVINO model: {MODEL_XML}")

    core = ov.Core()
    compile_config = {
        "PERFORMANCE_HINT": "LATENCY",
        "INFERENCE_PRECISION_HINT": "bf16",
    }

    model = core.read_model(MODEL_XML)
    model.reshape({0: [1, 3, MODEL_HEIGHT, MODEL_WIDTH]})
    compiled = core.compile_model(model, "CPU", compile_config)
    infer_request = compiled.create_infer_request()
    input_port = compiled.input(0)
    print("Model loaded on CPU")
    print(
        "OpenVINO config: "
        f"precision={compile_config['INFERENCE_PRECISION_HINT']} "
        f"performance={compile_config['PERFORMANCE_HINT']}"
    )

    dummy = np.zeros((1, 3, MODEL_HEIGHT, MODEL_WIDTH), dtype=np.float32)
    infer_request.infer({input_port: dummy})
    print("Warmup done")

    tmp_dir = tempfile.mkdtemp()
    tmp_video = os.path.join(tmp_dir, "sr_video_noaudio.mp4")

    try:
        cap = cv2.VideoCapture(input_video)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        print(f"Total frames: {total_frames}")
        print(f"Encoder: libx264 preset={ENCODE_PRESET} crf={ENCODE_CRF}")

        ffmpeg_cmd = [
            FFMPEG_BIN,
            "-y",
            "-f",
            "rawvideo",
            "-vcodec",
            "rawvideo",
            "-s",
            f"{OUTPUT_WIDTH}x{OUTPUT_HEIGHT}",
            "-pix_fmt",
            "bgr24",
            "-r",
            str(fps),
            "-i",
            "-",
            "-c:v",
            "libx264",
            "-preset",
            ENCODE_PRESET,
            "-crf",
            ENCODE_CRF,
            "-pix_fmt",
            "yuv420p",
            tmp_video,
        ]
        pipe = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE, stderr=subprocess.DEVNULL)

        frame_idx = 0
        total_infer_time = 0.0
        start_time = time.time()

        while True:
            ret, frame = cap.read()
            if not ret:
                break

            if frame.shape[1] != MODEL_WIDTH or frame.shape[0] != MODEL_HEIGHT:
                frame = cv2.resize(frame, (MODEL_WIDTH, MODEL_HEIGHT), interpolation=cv2.INTER_LANCZOS4)

            img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
            img_tensor = np.transpose(img, (2, 0, 1))[np.newaxis, ...]

            t0 = time.time()
            infer_request.infer({input_port: img_tensor})
            result = infer_request.get_output_tensor(0).data.copy()
            t1 = time.time()
            total_infer_time += t1 - t0

            output = result.squeeze(0).transpose(1, 2, 0)
            output = np.clip(output, 0, 1) * 255.0
            output = output.astype(np.uint8)
            output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)

            pipe.stdin.write(output.tobytes())
            frame_idx += 1

            if frame_idx % 10 == 0 or frame_idx == total_frames:
                elapsed = time.time() - start_time
                fps_proc = frame_idx / elapsed if elapsed > 0 else 0.0
                infer_fps = frame_idx / total_infer_time if total_infer_time > 0 else 0.0
                print(
                    f"\rProgress: {frame_idx}/{total_frames} "
                    f"({frame_idx * 100 / max(total_frames, 1):.1f}%) "
                    f"total: {fps_proc:.2f} fps, infer: {infer_fps:.2f} fps",
                    end="",
                    flush=True,
                )

        if pipe.stdin is not None:
            pipe.stdin.close()
        pipe.wait()
        cap.release()
        print()

        total_time = time.time() - start_time
        print(f"Done! {frame_idx} frames in {total_time:.2f}s")
        print(f"  Total pipeline: {frame_idx / total_time:.2f} fps")
        print(f"  Pure inference: {frame_idx / total_infer_time:.2f} fps ({total_infer_time:.2f}s)")
        print(f"  Avg infer/frame: {total_infer_time / frame_idx * 1000:.1f} ms")

        if has_audio:
            print("Merging audio...")
            merge_cmd = [
                FFMPEG_BIN,
                "-y",
                "-i",
                tmp_video,
                "-i",
                input_video,
                "-c:v",
                "copy",
                "-c:a",
                "copy",
                "-map",
                "0:v:0",
                "-map",
                "1:a:0?",
                "-shortest",
                output_video,
            ]
            result = subprocess.run(merge_cmd, capture_output=True, text=True)
            if result.returncode != 0:
                print(f"Audio merge warning: {result.stderr}")
                shutil.copy2(tmp_video, output_video)
        else:
            shutil.copy2(tmp_video, output_video)

        print(f"Saved: {output_video}")

    finally:
        shutil.rmtree(tmp_dir, ignore_errors=True)


if __name__ == "__main__":
    main()
