from __future__ import annotations

import os
import shutil
import sys
import time
from datetime import timedelta, timezone
from pathlib import Path
from statistics import median

from tools.video_formatting import (
    format_frame_count,
    format_framerate,
    format_seconds,
    join_parts,
    video_parts,
)
from tools.video_inputs import VideoInput
from tools.video_probe import VideoProbe


WORKSPACE_DIR = Path("video_workspace")
SOURCE_DIR = ".source"
HELPER_SCRIPT_NAME = "avisynth_helpers.avs"


def workspace_root(script_path: Path) -> Path:
    return script_path.resolve().parent / WORKSPACE_DIR


def visible_script_path(root: Path, item: VideoInput) -> Path:
    return root / item.collapsed_relative.with_name(
        f"{item.collapsed_relative.name}.avs"
    )


def hidden_source_script_path(root: Path, item: VideoInput) -> Path:
    return _source_sidecar_path(root, item, ".source.avs")


def frame_timestamp_values_path(root: Path, item: VideoInput) -> Path:
    return _source_sidecar_path(root, item, ".timestamps.txt")


def frame_time_values_path(root: Path, item: VideoInput) -> Path:
    return _source_sidecar_path(root, item, ".frame-times.txt")


def frame_time_label_values_path(root: Path, item: VideoInput) -> Path:
    return _source_sidecar_path(root, item, ".frame-time-labels.txt")


def frame_duration_values_path(root: Path, item: VideoInput) -> Path:
    return _source_sidecar_path(root, item, ".durations.txt")


def frame_duration_label_values_path(root: Path, item: VideoInput) -> Path:
    return _source_sidecar_path(root, item, ".duration-labels.txt")


def validation_script_path(root: Path, item: VideoInput) -> Path:
    return _source_sidecar_path(root, item, ".validate.avs")


def validation_csv_path(root: Path, item: VideoInput) -> Path:
    return _source_sidecar_path(root, item, ".frames.csv")


def editable_audio_script_path(root: Path, item: VideoInput) -> Path:
    visible = visible_script_path(root, item)
    return visible.with_name(f"{visible.stem}_audio.avs")


def rendered_audio_wav_path(root: Path, item: VideoInput) -> Path:
    return _source_sidecar_path(root, item, ".audio.wav")


def _source_sidecar_path(root: Path, item: VideoInput, suffix: str) -> Path:
    return root / SOURCE_DIR / item.collapsed_relative.with_name(
        f"{item.collapsed_relative.name}{suffix}"
    )


def existing_visible_scripts(root: Path, items: list[VideoInput]) -> list[Path]:
    return [path for path in (visible_script_path(root, item) for item in items) if path.exists()]


def write_workspace(
    *,
    root: Path,
    items: list[VideoInput],
    probes_by_path: dict[Path, VideoProbe],
    overwrite_visible: bool | set[Path],
    ffms2_plugin_path: Path | None,
) -> tuple[list[Path], list[Path]]:
    written: list[Path] = []
    skipped_visible: list[Path] = []

    source_root = root / SOURCE_DIR
    _clear_source_root(source_root)

    helper_script = helper_script_path()

    for source_id, item in enumerate(items):
        probe = probes_by_path[item.path.resolve()]
        source_script = hidden_source_script_path(root, item)
        timestamp_values = frame_timestamp_values_path(root, item)
        frame_time_values = frame_time_values_path(root, item)
        frame_time_label_values = frame_time_label_values_path(root, item)
        duration_values = frame_duration_values_path(root, item)
        duration_label_values = frame_duration_label_values_path(root, item)
        validation_script = validation_script_path(root, item)
        visible_script = visible_script_path(root, item)

        source_script.parent.mkdir(parents=True, exist_ok=True)
        visible_script.parent.mkdir(parents=True, exist_ok=True)

        timestamp_values.write_text(_timestamp_values_file(probe), encoding="utf-8")
        frame_time_values.write_text(_frame_time_values_file(probe), encoding="utf-8")
        frame_time_label_values.write_text(
            _frame_time_label_values_file(probe), encoding="utf-8"
        )
        duration_values.write_text(_duration_values_file(probe), encoding="utf-8")
        duration_label_values.write_text(_duration_label_values_file(probe), encoding="utf-8")
        source_script.write_text(
            _hidden_source_script(
                item,
                probe,
                source_id,
                source_script,
                helper_script,
                ffms2_plugin_path,
                timestamp_values,
                frame_time_values,
                frame_time_label_values,
                duration_values,
                duration_label_values,
            ),
            encoding="utf-8",
        )
        validation_script.write_text(
            _validation_script(visible_script, validation_script),
            encoding="utf-8",
        )
        written.extend([
            timestamp_values,
            frame_time_values,
            frame_time_label_values,
            duration_values,
            duration_label_values,
            source_script,
            validation_script,
        ])

        if visible_script.exists() and not _should_overwrite_visible(visible_script, overwrite_visible):
            skipped_visible.append(visible_script)
        else:
            visible_script.write_text(
                _visible_script(
                    visible_script,
                    source_script,
                    item,
                    probe,
                ),
                encoding="utf-8",
            )
            written.append(visible_script)

    return written, skipped_visible


def _clear_source_root(source_root: Path) -> None:
    if not source_root.exists():
        return
    for attempt in range(5):
        try:
            shutil.rmtree(source_root)
            return
        except OSError as exc:
            if attempt == 4:
                raise RuntimeError(
                    f"Could not reset {source_root}. Close any program using files in "
                    "video_workspace/.source and try again."
                ) from exc
            time.sleep(0.25)


def _should_overwrite_visible(path: Path, overwrite_visible: bool | set[Path]) -> bool:
    if isinstance(overwrite_visible, bool):
        return overwrite_visible
    return path.resolve() in overwrite_visible


def _visible_script(
    visible_script: Path,
    source_script: Path,
    item: VideoInput,
    probe: VideoProbe,
) -> str:
    import_path = _avs_path(os.path.relpath(source_script, visible_script.parent))
    audio_script_name = f"{visible_script.stem}_audio.avs"
    timing_kind = probe.timing.timing_kind.upper()
    final_clip = _default_final_clip(probe)
    windows_32bit_note = (
        "# End this file with a commented #32bit marker to use 32-bit AviSynth.\n"
        if sys.platform == "win32"
        else ""
    )
    return (
        "# Don't touch.\n"
        f'Import("{import_path}")\n'
        "\n"
        f"# Source: {item.path}\n"
        f"# Video: {_format_video_summary(probe)}\n"
        f"# Duration: {format_seconds(probe.duration_seconds)}, {_format_frame_count(probe)}\n"
        f"# Framerate: {format_framerate(timing_kind, probe.avg_frame_rate)}\n"
        "\n"
        "# Trimmed frames also trim audio according to the timestamps of the trimmed frames.\n"
        "# MBT_Drop marks frames for dropping without changing other frame timings.\n"
        "# Frame properties must be preserved.\n"
        "\n"
        "# Bundled helper functions:\n"
        "# MBT_Drop(last, int start, int \"end\"), MBT_Undrop(...)\n"
        "# MBT_DropEvery(last, int cycle, int offset0, ...), MBT_UndropEvery(...)\n"
        "#   Similar syntax as Trim/SelectEvery. end defaults to -1, meaning one frame.\n"
        "#   Examples: MBT_Drop(last, 100), MBT_Drop(last, 100, -20),\n"
        "#             MBT_DropEvery(last, 4, 1, 3)\n"
        "# MBT_Info(last, int \"size\", int \"align\")\n"
        "#   Overlay source frame, video time, absolute time, duration, and drop marker.\n"
        "# Resize(clip v, int \"w\", int \"h\", float \"dar\", float \"xcrop\", float \"ycrop\", int \"pos\", bool \"linear\", string \"input_matrix\", string \"output_matrix\")\n"
        "#   Resize/crop. linear auto-enables for large downscales unless set explicitly.\n"
        "# MBT_CorrectMatrix(last)\n"
        "#   Convert to the resolution-expected matrix using _Matrix as input metadata.\n"
        "# DeShake(clip c, int \"w\", int \"h\", float \"dar\", float \"crop\", float \"xcrop\", float \"ycrop\", int \"pos\", bool \"zoom\", bool \"rot\", float \"freq\")\n"
        "#   Stabilize through MVTools + DePan; those plugins must be loaded separately.\n"
        "# Cropf(...)\n"
        "#   Crop using fractional dimensions.\n"
        "# RotateCrop(clip c, float angle, float \"dar\")\n"
        "#   Rotate with manyPlus Turn, then crop the largest background-free rectangle.\n"
        "# FixContrast(...)\n"
        "#   Remap luma levels to limited-range video luma.\n"
        "\n"
        "# Optional audio edits:\n"
        f"#   Copy this script to {audio_script_name} next to this file.\n"
        "#   Edit that copy for audio fades/channel changes/etc. Its video output is ignored.\n"
        "#   Its audio output will be trimmed according to this script's frame edits.\n"
        f"{windows_32bit_note}"
        "\n"
        "# User edits below this line. The imported source clip will already be in \"last\".\n"
        "\n"
        f"{final_clip}\n"
    )


def _hidden_source_script(
    item: VideoInput,
    probe: VideoProbe,
    source_id: int,
    source_script: Path,
    helper_script: Path,
    ffms2_plugin_path: Path | None,
    timestamp_values: Path,
    frame_time_values: Path,
    frame_time_label_values: Path,
    duration_values: Path,
    duration_label_values: Path,
) -> str:
    source_path = _avs_path(str(item.path.resolve()))
    import_path = _avs_path(os.path.relpath(helper_script, source_script.parent))
    cache_path = _avs_path(
        str(source_script.with_name(f"{item.collapsed_relative.name}.ffindex").resolve())
    )
    audio_cache_path = _avs_path(
        str(source_script.with_name(f"{item.collapsed_relative.name}.audio.ffindex").resolve())
    )
    timestamp_values_path = _avs_path(str(timestamp_values.resolve()))
    frame_time_values_path = _avs_path(str(frame_time_values.resolve()))
    frame_time_label_values_path = _avs_path(str(frame_time_label_values.resolve()))
    duration_values_path = _avs_path(str(duration_values.resolve()))
    duration_label_values_path = _avs_path(str(duration_label_values.resolve()))
    if ffms2_plugin_path is None:
        plugin_setup = "# FFMS2 is expected from Avisynth autoload or a previously loaded plugin.\n"
    else:
        plugin_setup = f'LoadPlugin("{_avs_path(str(ffms2_plugin_path.resolve()))}")\n'
    audio_setup = (
        "try {\n"
        "    mbt_video_only = mbt_video_only\n"
        "} catch (err_msg) {\n"
        "    mbt_video_only = false\n"
        "}\n"
        f'source = mbt_video_only ? video : MBT_AudioSource(video, "{source_path}", "{audio_cache_path}")\n'
        if probe.audio_stream_count > 0
        else "source = video\n"
    )
    validation_blank_setup = (
        "try {\n"
        "    mbt_validation_blank = mbt_validation_blank\n"
        "} catch (err_msg) {\n"
        "    mbt_validation_blank = false\n"
        "}\n"
        "video = mbt_validation_blank ? BlankClip(video, color=$000000) : video\n"
    )
    return (
        "# Generated by media-batch-tools. This file is regenerated by video_encode.py.\n"
        "# Do not put manual edits here; edit the matching script in video_workspace/.\n"
        f"# Source: {source_path}\n"
        f"# Source id: {source_id}\n"
        f"# Detected timing: {probe.timing.timing_kind.upper()}\n"
        "# Requires FFMS2 to provide FFVideoSource.\n"
        "\n"
        f"{plugin_setup}"
        f'Import("{import_path}")\n'
        "\n"
        f'video = FFVideoSource("{source_path}", cachefile="{cache_path}")\n'
        f"{validation_blank_setup}"
        f"{audio_setup}"
        f"last = MBT_MarkSource(source, {source_id}, {_matrix_code_for_source(probe)})\n"
        f'last = ConditionalReader(last, "{timestamp_values_path}", "mbt_absolute_timestamp", false)\n'
        f'last = ConditionalReader(last, "{frame_time_values_path}", "mbt_frame_time_seconds", false)\n'
        f'last = ConditionalReader(last, "{frame_time_label_values_path}", "mbt_relative_timestamp", false)\n'
        f'last = ConditionalReader(last, "{duration_values_path}", "mbt_frame_duration_seconds", false)\n'
        f'last = ConditionalReader(last, "{duration_label_values_path}", "mbt_frame_duration", false)\n'
        "last\n"
    )


def _validation_script(
    visible_script: Path,
    validation_script: Path,
) -> str:
    import_path = _avs_path(os.path.relpath(visible_script, validation_script.parent))
    return (
        "# Generated by media-batch-tools. This file is regenerated by video_encode.py.\n"
        "# The media-batch-tools runner reads frame identity properties from this clip.\n"
        f'Import("{import_path}")\n'
        "last\n"
    )


def _timestamp_values_file(probe: VideoProbe) -> str:
    lines = [
        "# Generated by media-batch-tools for AviSynth ConditionalReader.",
        "Type string",
        "Default",
    ]
    if probe.metadata_end is None or probe.duration_seconds is None:
        return "\n".join(lines) + "\n"

    source_end = probe.metadata_end
    if source_end.tzinfo is None:
        source_end = source_end.replace(tzinfo=timezone.utc)
    source_begin = source_end - timedelta(seconds=probe.duration_seconds)
    for frame, timestamp in enumerate(probe.frame_timestamps):
        value = (source_begin + timedelta(seconds=timestamp)).astimezone(timezone.utc)
        lines.append(f"{frame} {value.isoformat().replace('+00:00', 'Z')}")
    return "\n".join(lines) + "\n"


def _frame_time_values_file(probe: VideoProbe) -> str:
    lines = [
        "# Generated by media-batch-tools for AviSynth ConditionalReader.",
        "Type float",
        "Default 0.0",
    ]
    for frame, timestamp in enumerate(probe.frame_timestamps):
        lines.append(f"{frame} {timestamp:.9f}")
    return "\n".join(lines) + "\n"


def _frame_time_label_values_file(probe: VideoProbe) -> str:
    lines = [
        "# Generated by media-batch-tools for AviSynth ConditionalReader.",
        "Type string",
        "Default 0.000000",
    ]
    for frame, timestamp in enumerate(probe.frame_timestamps):
        lines.append(f"{frame} {_format_relative_timestamp(timestamp)}")
    return "\n".join(lines) + "\n"


def _duration_values_file(probe: VideoProbe) -> str:
    lines = [
        "# Generated by media-batch-tools for AviSynth ConditionalReader.",
        "Type float",
        "Default 0.0",
    ]
    for frame, duration in enumerate(_source_frame_durations(probe)):
        lines.append(f"{frame} {duration:.9f}")
    return "\n".join(lines) + "\n"


def _duration_label_values_file(probe: VideoProbe) -> str:
    lines = [
        "# Generated by media-batch-tools for AviSynth ConditionalReader.",
        "Type string",
        "Default",
    ]
    for frame, duration in enumerate(_source_frame_durations(probe)):
        lines.append(f"{frame} {_format_duration_label(duration)}")
    return "\n".join(lines) + "\n"


def _format_duration_label(duration: float) -> str:
    if duration <= 0:
        return f"{duration:.6f} (1/inf)"
    reciprocal = 1.0 / duration
    return f"{duration:.6f} (1/{reciprocal:.2g})"


def _format_relative_timestamp(seconds: float) -> str:
    if seconds < 60.0:
        return f"{seconds:.6f}"

    total = int(seconds)
    fraction = seconds - total
    fractional = f"{fraction:.6f}"[1:]
    if total < 3600:
        minutes = total // 60
        remaining_seconds = total % 60
        return f"{minutes}:{remaining_seconds:02d}{fractional} ({seconds:.6f})"

    hours = total // 3600
    minutes = (total % 3600) // 60
    remaining_seconds = total % 60
    return f"{hours}:{minutes:02d}:{remaining_seconds:02d}{fractional} ({seconds:.6f})"


def _source_frame_durations(probe: VideoProbe) -> list[float]:
    timestamps = probe.frame_timestamps
    if not timestamps:
        return []

    durations: list[float] = []
    for current, next_timestamp in zip(timestamps, timestamps[1:]):
        durations.append(max(0.0, next_timestamp - current))

    fallback = median(durations) if durations else 0.0
    if probe.duration_seconds is not None and len(timestamps) > 1:
        last_duration = probe.duration_seconds - timestamps[-1]
        if last_duration <= 0:
            last_duration = fallback
    else:
        last_duration = fallback
    durations.append(max(0.0, last_duration))
    return durations


def _format_video_summary(probe: VideoProbe) -> str:
    parts = video_parts(
        width=probe.width,
        height=probe.height,
        pixel_format=probe.pixel_format,
        bit_depth=probe.bit_depth,
    )
    parts.append(_format_matrix(probe))
    return join_parts(parts)


def _format_matrix(probe: VideoProbe) -> str:
    expected = probe.expected_color_matrix or "unknown expected"
    if probe.color_matrix is not None:
        if probe.color_space:
            return f"{probe.color_matrix} ({probe.color_space}), expected {expected}"
        return f"{probe.color_matrix}, expected {expected}"
    if probe.color_space:
        return f"{probe.color_space}, expected {expected}"
    return f"unknown, expected {expected}"


def _matrix_code_for_source(probe: VideoProbe) -> int:
    return _matrix_code(probe.color_matrix or probe.expected_color_matrix)


def _matrix_code(matrix: str | None) -> int:
    if matrix == "Rec709":
        return 1
    if matrix == "Rec2020":
        return 9
    if matrix == "Rec601":
        return 6
    return 2


def _default_final_clip(probe: VideoProbe) -> str:
    source_matrix = probe.color_matrix
    expected_matrix = probe.expected_color_matrix
    if source_matrix is None or expected_matrix is None or source_matrix == expected_matrix:
        return "last"
    return "MBT_CorrectMatrix(last) # source metadata differs from resolution default"


def _format_frame_count(probe: VideoProbe) -> str:
    return format_frame_count(probe.stream_frame_count, len(probe.frame_timestamps))


def _avs_path(path: str) -> str:
    return path.replace("\\", "/").replace('"', '\\"')


def helper_script_path() -> Path:
    return Path(__file__).resolve().with_name(HELPER_SCRIPT_NAME)