from __future__ import annotations

import re
import subprocess
import sys
import tempfile
from dataclasses import replace
from pathlib import Path

from tools.console import PipelineProgressView, ProgressView
from tools.video_color import ColorMetadata
from tools.video_encode_output import (
    AudioEncodeOptions,
    VideoEncodeRequest,
    VideoEncodeResult,
    add_audio_input_options,
    add_audio_options,
    encode_video_only_y4m,
    pixel_format_bit_depth,
    read_ffmpeg_progress,
    run_y4m_encoder,
    x264_static_hdr_options,
)


def encode_video_with_timestamps(
    *,
    request: VideoEncodeRequest,
    mkvmerge: Path,
    x264: Path | None,
    frame_durations: list[float],
) -> VideoEncodeResult:
    if not frame_durations:
        raise RuntimeError("Cannot encode VFR output with no frame durations")

    request.output.parent.mkdir(parents=True, exist_ok=True)
    # Network shares may allow the final output but deny creating temporary folders.
    with tempfile.TemporaryDirectory(prefix="mbt-vfr-") as temp_dir:
        temp_root = Path(temp_dir)
        timecodes = temp_root / "frames.timecodes.txt"
        write_timecode_v2(timecodes, frame_durations)
        intermediate = temp_root / (
            "encoded.h264" if request.options.codec == "libx264" else "encoded.mkv"
        )
        if request.options.codec == "libx264":
            if x264 is None:
                raise RuntimeError("x264 CLI is required for timestamp-aware x264 encoding")
            renderer_returncode, encoder_log = encode_x264_y4m(
                request=request,
                x264=x264,
                timecodes=timecodes,
                output=intermediate,
            )
        elif request.options.codec == "libx265":
            result = encode_video_only_y4m(
                replace(request, output=intermediate, audio_options=AudioEncodeOptions(mode="none"))
            )
            renderer_returncode = result.renderer_returncode
            encoder_log = result.encoder_log
        else:
            raise RuntimeError(f"Unsupported timestamp-aware codec: {request.options.codec}")

        timed_video = temp_root / "timed.mkv"
        apply_video_timecodes(
            mkvmerge=mkvmerge,
            source=intermediate,
            timecodes=timecodes,
            output=timed_video,
            color_metadata=request.color_metadata,
        )
        ffmpeg_returncode = mux_timed_video_with_audio(
            request=request,
            timed_video=timed_video,
        )

    return VideoEncodeResult(
        script=request.script,
        output=request.output,
        renderer_returncode=renderer_returncode,
        ffmpeg_returncode=ffmpeg_returncode,
        encoder_log=encoder_log,
    )


def write_timecode_v2(path: Path, frame_durations: list[float]) -> None:
    timestamps = [0.0]
    for duration in frame_durations:
        if duration <= 0:
            raise RuntimeError(f"Invalid video frame duration: {duration}")
        timestamps.append(timestamps[-1] + duration)
    lines = ["# timecode format v2"]
    lines.extend(f"{timestamp * 1000:.9f}" for timestamp in timestamps)
    path.write_text("\n".join(lines) + "\n", encoding="utf-8")


def encode_x264_y4m(
    *,
    request: VideoEncodeRequest,
    x264: Path,
    timecodes: Path,
    output: Path,
) -> tuple[int, str]:
    run = run_y4m_encoder(
        y4m_command=request.y4m_command,
        script=request.script,
        encoder_command=x264_command(
            request=request,
            x264=x264,
            timecodes=timecodes,
            output=output,
        ),
        frame_count=request.frame_count,
        progress_label=request.progress_label,
        read_progress=read_x264_progress,
    )
    if run.encoder_returncode != 0:
        details = run.encoder_log
        if run.renderer_returncode != 0:
            details += "\nAvisynth runner:\n" + run.renderer_log
        raise RuntimeError(
            f"x264 failed for {request.script} with exit {run.encoder_returncode}:\n"
            + details.strip()
        )
    if run.renderer_returncode != 0:
        raise RuntimeError(
            f"Avisynth runner failed for {request.script} with exit {run.renderer_returncode}:\n"
            + run.renderer_log
        )
    return run.renderer_returncode, run.encoder_log


def x264_command(
    *,
    request: VideoEncodeRequest,
    x264: Path,
    timecodes: Path,
    output: Path,
) -> list[str]:
    command = [
        str(x264),
        "--demuxer",
        "y4m",
        "--tcfile-in",
        str(timecodes),
        "--crf",
        str(request.options.crf),
        "--preset",
        request.options.preset,
        "--output-depth",
        str(pixel_format_bit_depth(request.pixel_format)),
        "--output-csp",
        _x264_output_csp(request.pixel_format),
        "--verbose",
        "--no-progress",
    ]
    if request.options.profile:
        command.extend(["--profile", request.options.profile])
    if request.options.level:
        command.extend(["--level", request.options.level])
    if request.options.threads is not None:
        command.extend(["--threads", str(request.options.threads)])
    if request.color_metadata is not None:
        command.extend(request.color_metadata.x264_args())
    elif request.colorspace:
        command.extend(["--colormatrix", request.colorspace])
    command.extend(x264_static_hdr_options(request.static_hdr))
    command.extend(["--output", str(output), "-"])
    return command


def _x264_output_csp(pixel_format: str) -> str:
    if "444" in pixel_format:
        return "i444"
    if "422" in pixel_format:
        return "i422"
    return "i420"


def apply_video_timecodes(
    *,
    mkvmerge: Path,
    source: Path,
    timecodes: Path,
    output: Path,
    color_metadata: ColorMetadata | None = None,
) -> None:
    completed = subprocess.run(
        [
            str(mkvmerge),
            "--timestamp-scale",
            "1000",
            "--no-date",
            "-o",
            str(output),
            "--timestamps",
            f"0:{timecodes}",
            *(color_metadata.mkvmerge_args() if color_metadata else []),
            str(source),
        ],
        capture_output=True,
        text=True,
    )
    if completed.returncode != 0:
        details = (completed.stderr or completed.stdout).strip()
        raise RuntimeError(
            f"mkvmerge failed to apply video timestamps with exit "
            f"{completed.returncode}:\n{details}"
        )


def mux_timed_video_with_audio(
    *,
    request: VideoEncodeRequest,
    timed_video: Path,
) -> int:
    command = [
        str(request.ffmpeg),
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-nostats",
        "-progress",
        "pipe:2",
        "-i",
        str(timed_video),
    ]
    if request.audio_options.mode != "none" and request.source_has_audio:
        add_audio_input_options(
            command,
            audio_options=request.audio_options,
            audio_start=request.audio_start,
            audio_duration=request.audio_duration,
            allow_audio_copy=request.allow_audio_copy,
        )
        command.extend(["-i", str(request.audio_source)])

    command.extend(["-map", "0:v:0", "-c:v", "copy"])
    if request.color_metadata:
        command.extend(request.color_metadata.ffmpeg_args())
    if request.options.codec == "libx265":
        bitstream_filter = request.color_metadata.hevc_bitstream_filter() if request.color_metadata else None
        if bitstream_filter:
            command.extend(["-bsf:v", bitstream_filter])
    if request.output.suffix.lower() == ".mp4":
        command.extend(["-video_track_timescale", "90000", "-movflags", "+write_colr"])
    add_audio_options(
        command,
        audio_options=request.audio_options,
        audio_start=request.audio_start,
        audio_duration=request.audio_duration,
        audio_segments=request.audio_segments,
        audio_tempo=request.audio_tempo,
        audio_sample_rate=request.audio_sample_rate,
        source_has_audio=request.source_has_audio,
        allow_audio_copy=request.allow_audio_copy,
    )
    command.append(str(request.output))

    process = subprocess.Popen(
        command,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
        text=True,
    )
    ffmpeg_stderr = read_ffmpeg_progress(
        process,
        total_frames=request.frame_count,
        label=f"Muxing {request.progress_label}",
    )
    if process.returncode != 0:
        raise RuntimeError(
            f"FFmpeg failed while muxing timestamped output {request.output} with exit "
            f"{process.returncode}:\n{ffmpeg_stderr}"
        )
    return process.returncode


def read_x264_progress(
    process: subprocess.Popen,
    *,
    total_frames: int | None,
    label: str,
    pipeline_progress: PipelineProgressView | None = None,
) -> str:
    assert process.stderr is not None
    messages: list[str] = []
    progress = (
        ProgressView(total_frames or 0, label, embedded_percent=True, show_rate=True)
        if total_frames and sys.stdout.isatty() and pipeline_progress is None
        else None
    )
    if progress:
        progress.update(0)
    for line in process.stderr:
        match = re.search(r"x264 \[debug\]: frame=\s*(\d+)", line)
        if match:
            if progress:
                progress.update(int(match.group(1)) + 1)
            elif pipeline_progress is not None:
                pipeline_progress.update_encoding(int(match.group(1)) + 1)
        else:
            messages.append(line.rstrip())
    process.wait()
    if progress:
        progress.finish(keep=True)
    return "\n".join(messages)