from __future__ import annotations

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

from tools.console import PipelineProgressView, ProgressView
from tools.video_encode_output import (
    AudioEncodeOptions,
    VideoCodecOptions,
    VideoEncodeResult,
    add_audio_input_options,
    add_audio_options,
    encode_video_only_y4m,
    pixel_format_bit_depth,
    read_ffmpeg_progress,
    run_y4m_encoder,
)


def encode_video_with_timestamps(
    *,
    y4m_command: list[str],
    ffmpeg: Path,
    mkvmerge: Path,
    x264: Path | None,
    script: Path,
    frame_durations: list[float],
    audio_source: Path,
    output: Path,
    options: VideoCodecOptions,
    audio_options: AudioEncodeOptions,
    audio_start: float,
    audio_duration: float | None,
    audio_segments: list[tuple[float, float]] | None,
    audio_tempo: float,
    audio_sample_rate: int | None,
    source_has_audio: bool,
    allow_audio_copy: bool,
    pixel_format: str,
    colorspace: str | None,
    frame_count: int | None = None,
    progress_label: str | None = None,
) -> VideoEncodeResult:
    if not frame_durations:
        raise RuntimeError("Cannot encode VFR output with no frame durations")

    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 options.codec == "libx264" else "encoded.mkv"
        )
        if 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(
                y4m_command=y4m_command,
                x264=x264,
                script=script,
                timecodes=timecodes,
                output=intermediate,
                options=options,
                pixel_format=pixel_format,
                colorspace=colorspace,
                frame_count=len(frame_durations),
                progress_label=progress_label,
            )
        elif options.codec == "libx265":
            result = encode_video_only_y4m(
                y4m_command=y4m_command,
                ffmpeg=ffmpeg,
                script=script,
                audio_source=None,
                output=intermediate,
                options=options,
                audio_options=AudioEncodeOptions(mode="none"),
                pixel_format=pixel_format,
                colorspace=colorspace,
                frame_count=len(frame_durations),
                progress_label=progress_label,
            )
            renderer_returncode = result.renderer_returncode
            encoder_log = result.encoder_log
        else:
            raise RuntimeError(f"Unsupported timestamp-aware codec: {options.codec}")

        timed_video = temp_root / "timed.mkv"
        apply_video_timecodes(
            mkvmerge=mkvmerge,
            source=intermediate,
            timecodes=timecodes,
            output=timed_video,
        )
        ffmpeg_returncode = mux_timed_video_with_audio(
            ffmpeg=ffmpeg,
            timed_video=timed_video,
            audio_source=audio_source,
            output=output,
            audio_options=audio_options,
            audio_start=audio_start,
            audio_duration=audio_duration,
            audio_segments=audio_segments,
            audio_tempo=audio_tempo,
            audio_sample_rate=audio_sample_rate,
            source_has_audio=source_has_audio,
            allow_audio_copy=allow_audio_copy,
            frame_count=len(frame_durations),
            progress_label=progress_label,
        )

    return VideoEncodeResult(
        script=script,
        output=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(
    *,
    y4m_command: list[str],
    x264: Path,
    script: Path,
    timecodes: Path,
    output: Path,
    options: VideoCodecOptions,
    pixel_format: str,
    colorspace: str | None,
    frame_count: int | None,
    progress_label: str | None,
) -> tuple[int, str]:
    run = run_y4m_encoder(
        y4m_command=y4m_command,
        script=script,
        encoder_command=x264_command(
            x264=x264,
            timecodes=timecodes,
            output=output,
            options=options,
            pixel_format=pixel_format,
            colorspace=colorspace,
        ),
        frame_count=frame_count,
        progress_label=progress_label or script.name,
        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 {script} with exit {run.encoder_returncode}:\n"
            + details.strip()
        )
    if run.renderer_returncode != 0:
        raise RuntimeError(
            f"Avisynth runner failed for {script} with exit {run.renderer_returncode}:\n"
            + run.renderer_log
        )
    return run.renderer_returncode, run.encoder_log


def x264_command(
    *,
    x264: Path,
    timecodes: Path,
    output: Path,
    options: VideoCodecOptions,
    pixel_format: str,
    colorspace: str | None,
) -> list[str]:
    command = [
        str(x264),
        "--demuxer",
        "y4m",
        "--tcfile-in",
        str(timecodes),
        "--crf",
        str(options.crf),
        "--preset",
        options.preset,
        "--output-depth",
        str(pixel_format_bit_depth(pixel_format)),
        "--output-csp",
        _x264_output_csp(pixel_format),
        "--verbose",
        "--no-progress",
    ]
    if options.profile:
        command.extend(["--profile", options.profile])
    if options.level:
        command.extend(["--level", options.level])
    if options.threads is not None:
        command.extend(["--threads", str(options.threads)])
    if colorspace:
        command.extend(["--colormatrix", colorspace])
    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,
) -> None:
    completed = subprocess.run(
        [
            str(mkvmerge),
            "--timestamp-scale",
            "1000",
            "--no-date",
            "-o",
            str(output),
            "--timestamps",
            f"0:{timecodes}",
            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(
    *,
    ffmpeg: Path,
    timed_video: Path,
    audio_source: Path,
    output: Path,
    audio_options: AudioEncodeOptions,
    audio_start: float,
    audio_duration: float | None,
    audio_segments: list[tuple[float, float]] | None,
    audio_tempo: float,
    audio_sample_rate: int | None,
    source_has_audio: bool,
    allow_audio_copy: bool,
    frame_count: int | None,
    progress_label: str | None,
) -> int:
    command = [
        str(ffmpeg),
        "-y",
        "-hide_banner",
        "-loglevel",
        "error",
        "-nostats",
        "-progress",
        "pipe:2",
        "-i",
        str(timed_video),
    ]
    if audio_options.mode != "none" and source_has_audio:
        add_audio_input_options(
            command,
            audio_options=audio_options,
            audio_start=audio_start,
            audio_duration=audio_duration,
            allow_audio_copy=allow_audio_copy,
        )
        command.extend(["-i", str(audio_source)])

    command.extend(["-map", "0:v:0", "-c:v", "copy"])
    if output.suffix.lower() == ".mp4":
        command.extend(["-video_track_timescale", "90000"])
    add_audio_options(
        command,
        audio_options=audio_options,
        audio_start=audio_start,
        audio_duration=audio_duration,
        audio_segments=audio_segments,
        audio_tempo=audio_tempo,
        audio_sample_rate=audio_sample_rate,
        source_has_audio=source_has_audio,
        allow_audio_copy=allow_audio_copy,
    )
    command.append(str(output))

    process = subprocess.Popen(
        command,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
        text=True,
    )
    ffmpeg_stderr = read_ffmpeg_progress(
        process,
        total_frames=frame_count,
        label=f"Muxing {progress_label or output.name}",
    )
    if process.returncode != 0:
        raise RuntimeError(
            f"FFmpeg failed while muxing timestamped output {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)