#!/usr/bin/env python3

# Create one MP4 video for each deployment from timestamped still images.
#
# Usage:
#   python 02_create_videos.py
#
# You will be prompted to enter path to the campaign folder.
#
# Expected structure:
#
#   campaign/
#   ├── deployment/
#   │   ├── deployment.csv
#   │   └── any_folder_name/
#   │       ├── image1.jpg
#   │       ├── image2.jpg
#   │       └── ...
#
# or:
#
#   campaign/
#   ├── platform/
#   │   └── deployment/
#   │       ├── deployment.csv
#   │       └── any_folder_name/
#   │           ├── image1.jpg
#   │           ├── image2.jpg
#   │           └── ...
#
# The parent of the folder containing images is treated as the deployment folder.
#
# Each deployment folder must contain exactly one CSV with:
#   key             = image filename without extension
#   timestamp_start = date/time used to order the images
#
# Final videos are written directly to the selected campaign folder as:
#   <campaign>_<deployment>.mp4
#
# Install:
#   pip install numpy pandas opencv-python
#
# ffmpeg must also be installed and available on PATH.
# NVIDIA GPU encoding requires an ffmpeg build with h264_nvenc support.

from pathlib import Path
import shutil
import subprocess

import cv2
import numpy as np
import pandas as pd


# =========================
# CONFIG
# =========================

seconds_per_image = 1.0
fade_seconds = 0.15
output_fps = 20
output_height = 1080

extensions = {".jpg", ".jpeg", ".png"}

ffmpeg_exe = "ffmpeg"  # or set as r"path\to\ffmpeg.exe"

# NVIDIA GPU encode settings
nvenc_preset = "p5"   # p1 fastest, p7 slowest/best
nvenc_cq = 23         # higher = smaller/lower quality

delete_temp_mp4 = True


# =========================
# GET CAMPAIGN FOLDER
# =========================

campaign_input = input("Enter path to campaign folder: ").strip().strip('"')
campaign_folder = Path(campaign_input).expanduser().resolve()

if not campaign_folder.is_dir():
    raise SystemExit(f"Campaign folder not found: {campaign_folder}")

print(f"\nCampaign folder: {campaign_folder}")


# =========================
# DERIVED SETTINGS
# =========================

hold_frames = int(round(seconds_per_image * output_fps))
fade_frames = int(round(fade_seconds * output_fps))

if hold_frames < 1:
    raise SystemExit("seconds_per_image is too small.")

if fade_frames < 1:
    raise SystemExit("fade_seconds is too small for the chosen output_fps.")

if fade_frames >= hold_frames:
    raise SystemExit("fade_seconds must be shorter than seconds_per_image.")

fourcc = cv2.VideoWriter_fourcc(*"mp4v")


# =========================
# CHECK FFMPEG
# =========================

if Path(ffmpeg_exe).is_file():
    ffmpeg_path = str(Path(ffmpeg_exe))
else:
    ffmpeg_path = shutil.which(ffmpeg_exe)

if ffmpeg_path is None:
    raise SystemExit(
        "ffmpeg not found. Add ffmpeg to PATH or set ffmpeg_exe "
        "to the full path to ffmpeg.exe."
    )


# =========================
# FIND DEPLOYMENT IMAGE FOLDERS
# =========================

deployment_folders = []

# Find any folder below the campaign that directly contains image files.
# Its parent is treated as the deployment folder.
for folder in sorted(p for p in campaign_folder.rglob("*") if p.is_dir()):

    has_images = any(
        p.is_file() and p.suffix.lower() in extensions
        for p in folder.iterdir()
    )

    if not has_images:
        continue

    deployment_folder = folder.parent
    deployment_name = deployment_folder.name

    deployment_folders.append(
        (deployment_name, deployment_folder, folder)
    )

if not deployment_folders:
    raise SystemExit(
        "No image folders found below the selected campaign folder."
    )

print(f"Found {len(deployment_folders)} deployment(s).")


# =========================
# HELPER
# =========================

def make_canvas(img_path, output_width, output_height):
    img = cv2.imread(str(img_path))

    if img is None:
        return None

    h, w = img.shape[:2]

    scale = output_height / h
    resized_width = int(round(w * scale))

    img_resized = cv2.resize(
        img,
        (resized_width, output_height),
        interpolation=cv2.INTER_AREA
    )

    canvas = np.zeros(
        (output_height, output_width, 3),
        dtype=np.uint8
    )

    x_offset = (output_width - resized_width) // 2

    canvas[:, x_offset:x_offset + resized_width] = img_resized

    return canvas


# =========================
# PROCESS EACH DEPLOYMENT
# =========================

processed_count = 0
skipped_count = 0
failed_count = 0

for deployment_name, deployment_folder, images_folder in deployment_folders:

    print(f"\n=== Processing: {deployment_folder} ===")
    print(f"Image folder: {images_folder}")

    csv_files = [
        p
        for p in deployment_folder.iterdir()
        if p.is_file() and p.suffix.lower() == ".csv"
    ]

    if len(csv_files) == 0:
        print(f"Skipped (no CSV found): {deployment_folder}")
        skipped_count += 1
        continue

    if len(csv_files) > 1:
        print(f"Skipped (more than one CSV found): {deployment_folder}")
        skipped_count += 1
        continue

    csv_path = csv_files[0]

    temp_mp4 = (
        campaign_folder
        / f"{campaign_folder.name}_{deployment_name}_temp.mp4"
    )

    final_mp4 = (
        campaign_folder
        / f"{campaign_folder.name}_{deployment_name}.mp4"
    )

    if final_mp4.exists():
        print(f"Skipped (final MP4 already exists): {final_mp4}")
        skipped_count += 1
        continue

    try:

        # =========================
        # READ CSV
        # =========================

        try:
            df = pd.read_csv(csv_path)

        except Exception as e:
            print(f"Skipped (could not read CSV): {csv_path} -> {e}")
            skipped_count += 1
            continue

        required_cols = {"key", "timestamp_start"}

        missing = required_cols - set(df.columns)

        if missing:
            print(
                f"Skipped (missing required columns {missing}): "
                f"{csv_path}"
            )
            skipped_count += 1
            continue

        df["key"] = df["key"].astype(str).str.strip()

        df["timestamp_start"] = pd.to_datetime(
            df["timestamp_start"],
            errors="coerce",
            utc=True
        )

        df = df.dropna(
            subset=["key", "timestamp_start"]
        ).copy()

        if df.empty:
            print(
                f"Skipped (no valid key/timestamp_start rows): "
                f"{csv_path}"
            )
            skipped_count += 1
            continue

        df = df.sort_values(
            "timestamp_start",
            ascending=True
        )


        # =========================
        # BUILD IMAGE LOOKUP
        # =========================

        image_lookup = {}

        for img_path in images_folder.iterdir():

            if not img_path.is_file():
                continue

            if img_path.suffix.lower() not in extensions:
                continue

            stem = img_path.stem

            if stem not in image_lookup:
                image_lookup[stem] = img_path

            else:
                print(
                    "Warning: duplicate image stem found, "
                    f"keeping first: {stem}"
                )


        # =========================
        # ORDER IMAGES FROM CSV
        # =========================

        ordered_image_files = []
        missing_images = 0

        for _, row in df.iterrows():

            key = row["key"]

            img_path = image_lookup.get(key)

            if img_path is not None:
                ordered_image_files.append(img_path)

            else:
                missing_images += 1

        if not ordered_image_files:
            print(
                f"Skipped (no images matched CSV keys): "
                f"{deployment_folder}"
            )
            skipped_count += 1
            continue

        if missing_images:
            print(
                f"Warning: {missing_images} CSV row(s) had no "
                f"matching image in {images_folder}"
            )


        # =========================
        # WORK OUT OUTPUT WIDTH
        # =========================

        resized_widths = []
        readable_images = []

        for img_path in ordered_image_files:

            img = cv2.imread(str(img_path))

            if img is None:
                print(
                    f"Skipped unreadable image during sizing: "
                    f"{img_path}"
                )
                continue

            h, w = img.shape[:2]

            scale = output_height / h

            resized_width = int(round(w * scale))

            resized_widths.append(resized_width)
            readable_images.append(img_path)

        if not readable_images:
            print(
                f"No readable matched images in: {images_folder}"
            )
            skipped_count += 1
            continue

        output_width = max(resized_widths)

        if output_width % 2 != 0:
            output_width += 1

        if output_height % 2 != 0:
            raise SystemExit(
                "output_height must be even for MP4 safety."
            )


        # =========================
        # OPEN WRITER
        # =========================

        writer = cv2.VideoWriter(
            str(temp_mp4),
            fourcc,
            output_fps,
            (output_width, output_height)
        )

        if not writer.isOpened():
            print(
                f"Could not open video writer for: {temp_mp4}"
            )
            failed_count += 1
            continue


        # =========================
        # PREBUILD CANVASES
        # =========================

        canvases = []

        for img_path in readable_images:

            canvas = make_canvas(
                img_path,
                output_width,
                output_height
            )

            if canvas is None:
                print(
                    f"Skipped unreadable image: {img_path}"
                )
                continue

            canvases.append(
                (img_path.name, canvas)
            )

        if not canvases:

            writer.release()

            if temp_mp4.exists():
                temp_mp4.unlink()

            print(
                "No usable images after canvas generation."
            )

            skipped_count += 1
            continue


        # =========================
        # WRITE TEMP VIDEO
        # =========================

        if len(canvases) == 1:

            _, canvas = canvases[0]

            for _ in range(hold_frames):
                writer.write(canvas)

        else:

            static_frames = hold_frames - fade_frames

            for i in range(len(canvases) - 1):

                current_name, current_canvas = canvases[i]
                next_name, next_canvas = canvases[i + 1]

                # Clean hold
                for _ in range(static_frames):
                    writer.write(current_canvas)

                # Crossfade
                for j in range(1, fade_frames + 1):

                    alpha = j / (fade_frames + 1)

                    blended = cv2.addWeighted(
                        current_canvas,
                        1.0 - alpha,
                        next_canvas,
                        alpha,
                        0.0
                    )

                    writer.write(blended)

                print(
                    f"Added with fade: "
                    f"{current_name} -> {next_name}"
                )

            # Hold final image
            _, final_canvas = canvases[-1]

            for _ in range(hold_frames):
                writer.write(final_canvas)

        writer.release()

        print(f"Saved temp MP4: {temp_mp4}")


        # =========================
        # GPU RE-ENCODE WITH FFMPEG
        # =========================

        ffmpeg_cmd = [
            ffmpeg_path,
            "-y",
            "-i", str(temp_mp4),
            "-c:v", "h264_nvenc",
            "-preset", nvenc_preset,
            "-cq", str(nvenc_cq),
            "-pix_fmt", "yuv420p",
            "-movflags", "+faststart",
            str(final_mp4)
        ]

        print(
            "Re-encoding with ffmpeg "
            "(NVIDIA GPU)..."
        )

        result = subprocess.run(
            ffmpeg_cmd,
            capture_output=True,
            text=True
        )

        if result.returncode != 0:

            print(result.stderr)

            print(
                f"Failed ffmpeg re-encode: "
                f"{deployment_folder}"
            )

            failed_count += 1
            continue

        print(f"Saved final MP4: {final_mp4}")

        if delete_temp_mp4 and temp_mp4.exists():

            temp_mp4.unlink()

            print(
                f"Deleted temp MP4: {temp_mp4}"
            )

        processed_count += 1

    except Exception as e:

        print(
            f"Failed: {deployment_folder} -> {e}"
        )

        failed_count += 1
        continue


# =========================
# SUMMARY
# =========================

print("\nDone.")
print(f"Processed successfully: {processed_count}")
print(f"Skipped: {skipped_count}")
print(f"Failed: {failed_count}")