# Copyright (c) 2026 DeepBeepMeep - part of WanGP, WanGP Community License 2.0 (see LICENSE.txt)
from __future__ import annotations

import os
from dataclasses import dataclass

from . import media_io as media
from . import output_paths
from . import video_buffers as video


@dataclass
class MuxSession:
    mux_process: object | None = None
    output_path_for_write: str = ""
    mux_output_path: str = ""
    reserved_metadata_path: str | None = None
    initial_metadata: dict | None = None
    stopped: bool = False
    mux_finished: bool = False

    def set_output_path(self, output_path: str) -> None:
        self.output_path_for_write = output_path
        self.mux_output_path = ""

    def set_initial_metadata(self, metadata: dict) -> None:
        self.initial_metadata = metadata

    def ensure_started(
        self,
        *,
        server_config: dict,
        ffmpeg_path: str,
        process_is_hdr: bool,
        use_live_av_mux: bool,
        output_container: str,
        source_path: str,
        exact_start_seconds: float,
        selected_audio_track: int | None,
        resolved_width: int,
        resolved_height: int,
        fps_float: float,
        source_audio_duration_seconds: float | None = None,
    ) -> None:
        if self.mux_process is not None:
            return
        self.mux_output_path = media.reserve_working_output_path(self.output_path_for_write)
        self.reserved_metadata_path = media.create_reserved_metadata_file(self.mux_output_path, self.initial_metadata)
        if process_is_hdr:
            hdr_video_crf = server_config.get("hdr_video_crf", 8)
            self.mux_process = media.start_hdr_av_mux_process(ffmpeg_path, self.mux_output_path, resolved_width, resolved_height, fps_float, hdr_video_crf, output_container, source_path, exact_start_seconds, selected_audio_track, self.reserved_metadata_path, source_audio_duration_seconds) if use_live_av_mux else media.start_hdr_video_mux_process(ffmpeg_path, self.mux_output_path, resolved_width, resolved_height, fps_float, hdr_video_crf, output_container, self.reserved_metadata_path)
            return
        video_codec = server_config.get("video_output_codec", "libx264_8")
        self.mux_process = media.start_av_mux_process(ffmpeg_path, self.mux_output_path, resolved_width, resolved_height, fps_float, video_codec, output_container, source_path, exact_start_seconds, selected_audio_track, self.reserved_metadata_path, source_audio_duration_seconds) if use_live_av_mux else media.start_video_mux_process(ffmpeg_path, self.mux_output_path, resolved_width, resolved_height, fps_float, video_codec, output_container, self.reserved_metadata_path)

    def write_chunk(self, *, process_is_hdr: bool, video_tensor_hdr, video_tensor_uint8, start_frame: int, frame_count: int):
        if process_is_hdr and video_tensor_hdr is not None:
            return video.write_hdr_video_chunk(self.mux_process, video_tensor_hdr, start_frame=start_frame, frame_count=frame_count)
        return video.write_video_chunk(self.mux_process, video_tensor_uint8, start_frame=start_frame, frame_count=frame_count)

    def finalize(self, *, timeout_seconds: float | None = None, label: str = "") -> tuple[int, str, bool]:
        kwargs = {"label": label}
        if timeout_seconds is not None:
            kwargs["timeout_seconds"] = timeout_seconds
        return_code, stderr, forced_termination = media.finalize_mux_process(self.mux_process, **kwargs)
        self.mux_finished = True
        return return_code, stderr, forced_termination

    def promote_output(self, *, notify=None) -> tuple[str, str]:
        source_path = self.mux_output_path
        promoted_path = output_paths.promote_output_file(source_path, self.output_path_for_write, notify=notify)
        self.output_path_for_write = promoted_path
        self.mux_output_path = promoted_path
        return source_path, promoted_path

    def cleanup_partial_outputs(self) -> None:
        if self.mux_process is not None and not self.mux_finished and self.mux_process.poll() is None:
            try:
                self.finalize()
            except Exception:
                pass
        if self.mux_output_path and os.path.isfile(self.mux_output_path) and os.path.getsize(self.mux_output_path) <= 0:
            media.delete_file_if_exists(self.mux_output_path, label="empty working output")
        media.delete_file_if_exists(self.reserved_metadata_path, label="reserved metadata file")