import os
import argparse
from pathlib import Path

import imageio.v2 as imageio
import numpy as np
from PIL import Image, ImageDraw, ImageFont

from head_extension_app import (
    ct_window,
    fit_image,
    load_dicom_volume,
)


OUTPUT_DIR = Path("ppt_video")

FPS = 30
DURATION_SECONDS = 6
END_HOLD_SECONDS = 1


def get_font(size):
    for path in [r"C:\Windows\Fonts\arial.ttf", r"C:\Windows\Fonts\calibri.ttf"]:
        if os.path.exists(path):
            return ImageFont.truetype(path, size)
    return ImageFont.load_default()


def make_frame(slice_image, slice_index, slice_count, source_label, show_arrow=True):
    frame = Image.new("RGB", (1920, 1080), (0, 0, 0))
    draw = ImageDraw.Draw(frame)

    panel = fit_image(slice_image, 1160, 860)
    frame.paste(panel, (380, 145))

    title_font = get_font(48)
    small_font = get_font(28)
    draw.text((410, 70), f"{source_label} DICOM z-axis sequence", font=title_font, fill=(255, 255, 255))
    draw.text(
        (410, 1015),
        f"Slice {slice_index + 1} / {slice_count}",
        font=small_font,
        fill=(210, 210, 210),
    )

    arrow = (255, 210, 60)
    if show_arrow:
        x0, y0, x1, y1 = 1320, 275, 1450, 205
        draw.line((x0, y0, x1, y1), fill=arrow, width=8)
        draw.polygon([(x1, y1), (x1 - 34, y1 + 3), (x1 - 12, y1 + 29)], fill=arrow)

    bar_x, bar_y, bar_w, bar_h = 650, 1018, 870, 12
    draw.rounded_rectangle(
        (bar_x, bar_y, bar_x + bar_w, bar_y + bar_h),
        radius=6,
        fill=(70, 70, 70),
    )
    fill_w = int(bar_w * (slice_index + 1) / slice_count) if slice_count else 0
    draw.rounded_rectangle(
        (bar_x, bar_y, bar_x + fill_w, bar_y + bar_h),
        radius=6,
        fill=arrow,
    )

    return frame


def parse_args():
    parser = argparse.ArgumentParser(
        description="Generate a z-axis DICOM sequence MP4 video."
    )
    parser.add_argument(
        "--input",
        default="input_ct_2F",
        help="Input DICOM folder. Default: input_ct_2F",
    )
    parser.add_argument(
        "--output",
        default=str(OUTPUT_DIR / "head_extension_0_to_20deg.mp4"),
        help="Output MP4 file path.",
    )
    parser.add_argument(
        "--max-angle",
        type=float,
        default=20.0,
        help="Kept for compatibility; z-axis sequence videos do not use this value.",
    )
    parser.add_argument(
        "--duration",
        type=float,
        default=DURATION_SECONDS,
        help="Animation duration in seconds before final hold. Default: 6",
    )
    parser.add_argument(
        "--no-arrow",
        action="store_true",
        help="Hide the yellow direction arrow.",
    )
    parser.add_argument(
        "--source-label",
        default="DICOM",
        help="Label shown in the video title.",
    )
    return parser.parse_args()


def generate_video(input_dir, output_path, max_angle=20.0, duration_seconds=6.0, show_arrow=True, source_label="DICOM"):
    output_file = Path(output_path)
    output_file.parent.mkdir(parents=True, exist_ok=True)

    volume = load_dicom_volume(input_dir)
    slice_count = volume.shape[0]

    moving_frames = max(1, int(FPS * duration_seconds))
    hold_frames = FPS * END_HOLD_SECONDS

    with imageio.get_writer(
        output_file,
        format="FFMPEG",
        fps=FPS,
        codec="libx264",
        quality=8,
        macro_block_size=1,
    ) as writer:
        for index in range(moving_frames):
            t = index / (moving_frames - 1) if moving_frames > 1 else 0
            slice_index = min(slice_count - 1, int(round(t * (slice_count - 1))))
            slice_image = Image.fromarray(ct_window(volume[slice_index])).convert("RGB")
            writer.append_data(np.asarray(make_frame(slice_image, slice_index, slice_count, source_label, show_arrow)))

        for _ in range(hold_frames):
            slice_image = Image.fromarray(ct_window(volume[-1])).convert("RGB")
            writer.append_data(np.asarray(make_frame(slice_image, slice_count - 1, slice_count, source_label, show_arrow)))

    return output_file.resolve()


def main():
    args = parse_args()
    output_file = generate_video(
        args.input,
        args.output,
        args.max_angle,
        args.duration,
        not args.no_arrow,
        args.source_label,
    )
    print(output_file)


if __name__ == "__main__":
    main()