"""FFmpeg video decoding, encoding, audio extraction, and streaming pipe I/O."""

from __future__ import annotations

import json
import subprocess
import threading
from dataclasses import dataclass
from pathlib import Path
from typing import Generator

import numpy as np
from loguru import logger


@dataclass
class VideoInfo:
    width: int
    height: int
    fps: float
    total_frames: int
    duration: float
    codec: str
    has_audio: bool
    audio_codec: str | None
    file_size_bytes: int


class FFmpegError(RuntimeError):
    pass


class FFmpegHandler:
    def __init__(self, ffmpeg_path: str = "ffmpeg", ffprobe_path: str = "ffprobe") -> None:
        self.ffmpeg = ffmpeg_path
        self.ffprobe = ffprobe_path

    def probe(self, video_path: str) -> VideoInfo:
        """Get video metadata via ffprobe."""
        cmd = [
            self.ffprobe, "-v", "quiet",
            "-print_format", "json",
            "-show_format", "-show_streams",
            video_path,
        ]
        proc = self._run(cmd, capture_stdout=True)
        data = json.loads(proc.stdout)

        video_stream = None
        audio_stream = None
        for s in data.get("streams", []):
            if s["codec_type"] == "video" and video_stream is None:
                video_stream = s
            elif s["codec_type"] == "audio" and audio_stream is None:
                audio_stream = s

        if video_stream is None:
            raise FFmpegError(f"No video stream found in {video_path}")

        fmt = data.get("format", {})

        # Parse fps from r_frame_rate (e.g. "30/1", "30000/1001", or "30")
        fps_str = video_stream.get("r_frame_rate", "30/1")
        if "/" in fps_str:
            num, den = fps_str.split("/")
            fps = float(num) / float(den) if float(den) != 0 else 30.0
        else:
            try:
                fps = float(fps_str)
            except ValueError:
                fps = 30.0

        total_frames = int(video_stream.get("nb_frames", 0))
        duration = float(fmt.get("duration", video_stream.get("duration", 0)))
        if total_frames == 0 and duration > 0:
            total_frames = int(duration * fps)

        return VideoInfo(
            width=int(video_stream["width"]),
            height=int(video_stream["height"]),
            fps=fps,
            total_frames=total_frames,
            duration=duration,
            codec=video_stream.get("codec_name", "unknown"),
            has_audio=audio_stream is not None,
            audio_codec=audio_stream.get("codec_name") if audio_stream else None,
            file_size_bytes=int(fmt.get("size", 0)),
        )

    def extract_frames(
        self,
        video_path: str,
        output_dir: str,
        start_frame: int = 0,
        end_frame: int | None = None,
        scale: tuple[int, int] | None = None,
    ) -> list[str]:
        """Extract frames to PNG files. Returns list of file paths."""
        output = Path(output_dir)
        output.mkdir(parents=True, exist_ok=True)

        info = self.probe(video_path)
        start_time = start_frame / info.fps if start_frame > 0 else 0

        cmd = [self.ffmpeg, "-y"]
        if start_time > 0:
            cmd += ["-ss", f"{start_time:.6f}"]
        cmd += ["-i", video_path]

        if end_frame is not None:
            count = end_frame - start_frame
            cmd += ["-frames:v", str(count)]

        if scale:
            cmd += ["-vf", f"scale={scale[0]}:{scale[1]}"]

        pattern = str(output / "%06d.png")
        cmd += ["-start_number", "0", pattern]

        self._run(cmd)

        frames = sorted(output.glob("*.png"))
        return [str(f) for f in frames]

    def frames_to_numpy(
        self,
        video_path: str,
        start_frame: int = 0,
        count: int | None = None,
        scale: tuple[int, int] | None = None,
        video_info: VideoInfo | None = None,
    ) -> np.ndarray:
        """Decode frames directly to numpy array (N, H, W, 3) uint8 RGB via pipe.

        Pass video_info to skip redundant ffprobe calls.
        """
        info = video_info or self.probe(video_path)
        w = scale[0] if scale else info.width
        h = scale[1] if scale else info.height
        start_time = start_frame / info.fps if start_frame > 0 else 0

        cmd = [self.ffmpeg]
        if start_time > 0:
            cmd += ["-ss", f"{start_time:.6f}"]
        cmd += ["-i", video_path]

        if count is not None:
            cmd += ["-frames:v", str(count)]

        vf_filters = []
        if scale:
            vf_filters.append(f"scale={w}:{h}")
        if vf_filters:
            cmd += ["-vf", ",".join(vf_filters)]

        cmd += [
            "-f", "rawvideo",
            "-pix_fmt", "rgb24",
            "pipe:1",
        ]

        proc = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        stderr_lines: list[str] = []
        stderr_thread = threading.Thread(
            target=self._drain_stderr, args=(proc, stderr_lines), daemon=True,
        )
        stderr_thread.start()

        frame_size = w * h * 3
        frames: list[np.ndarray] = []
        while True:
            raw = proc.stdout.read(frame_size)
            if len(raw) < frame_size:
                break
            frame = np.frombuffer(raw, dtype=np.uint8).reshape(h, w, 3)
            frames.append(frame)

        proc.stdout.close()
        try:
            proc.wait(timeout=120)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait(timeout=5)
            raise FFmpegError(f"frames_to_numpy timed out for {video_path}")
        stderr_thread.join(timeout=5)

        if proc.returncode != 0:
            err_msg = "\n".join(stderr_lines[-10:])
            raise FFmpegError(f"frames_to_numpy failed (rc={proc.returncode}): {err_msg}")

        if not frames:
            raise FFmpegError(f"No frames decoded from {video_path}")

        if count is not None and len(frames) != count:
            logger.warning(f"Expected {count} frames, decoded {len(frames)} from {video_path}")

        return np.stack(frames)

    def extract_audio(self, video_path: str, output_path: str) -> bool:
        """Extract audio track. Returns False if no audio."""
        info = self.probe(video_path)
        if not info.has_audio:
            return False

        cmd = [
            self.ffmpeg, "-y",
            "-i", video_path,
            "-vn", "-acodec", "copy",
            output_path,
        ]
        self._run(cmd)
        return True

    def encode_video(
        self,
        frames_dir: str | None = None,
        frames_array: np.ndarray | None = None,
        output_path: str = "output.mp4",
        fps: float = 30.0,
        codec: str = "libx264",
        crf: int = 18,
        audio_path: str | None = None,
        pixel_format: str = "yuv420p",
    ) -> str:
        """Encode video from frame directory or numpy array."""
        if frames_array is not None:
            return self._encode_from_array(frames_array, output_path, fps, codec, crf, audio_path, pixel_format)
        if frames_dir is not None:
            return self._encode_from_dir(frames_dir, output_path, fps, codec, crf, audio_path, pixel_format)
        raise ValueError("Either frames_dir or frames_array must be provided")

    def merge_audio(self, video_path: str, audio_path: str, output_path: str) -> str:
        """Merge audio track into video (copy both streams)."""
        cmd = [
            self.ffmpeg, "-y",
            "-i", video_path,
            "-i", audio_path,
            "-c:v", "copy", "-c:a", "copy",
            "-shortest",
            output_path,
        ]
        self._run(cmd)
        return output_path

    # --- Streaming Encoder ---

    def create_streaming_encoder(
        self,
        output_path: str,
        width: int,
        height: int,
        fps: float = 30.0,
        codec: str = "libx264",
        crf: int = 18,
        pixel_format: str = "yuv420p",
        hw_accel: bool = False,
    ) -> StreamingEncoder:
        """Create a StreamingEncoder context manager for pipe-based encoding.

        If hw_accel=True, uses VideoToolbox hardware encoder (C3).
        """
        if hw_accel:
            codec = "h264_videotoolbox"
            # VideoToolbox uses -b:v instead of -crf; ~10Mbps ≈ crf 18 quality
            crf = -1  # signal to StreamingEncoder to use bitrate mode

        return StreamingEncoder(
            ffmpeg_path=self.ffmpeg,
            output_path=output_path,
            width=width,
            height=height,
            fps=fps,
            codec=codec,
            crf=crf,
            pixel_format=pixel_format,
        )

    # --- Private helpers ---

    def _encode_from_array(
        self, frames: np.ndarray, output_path: str, fps: float,
        codec: str, crf: int, audio_path: str | None, pixel_format: str,
    ) -> str:
        n, h, w, _ = frames.shape
        cmd = [
            self.ffmpeg, "-y",
            "-f", "rawvideo", "-pix_fmt", "rgb24",
            "-s", f"{w}x{h}", "-r", str(fps),
            "-i", "pipe:0",
        ]

        if audio_path:
            cmd += ["-i", audio_path, "-c:a", "copy"]

        cmd += [
            "-c:v", codec, "-crf", str(crf),
            "-pix_fmt", pixel_format,
            output_path,
        ]

        proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE)

        stderr_lines: list[str] = []
        t = threading.Thread(target=self._drain_stderr, args=(proc, stderr_lines), daemon=True)
        t.start()

        for i in range(n):
            proc.stdin.write(frames[i].tobytes())

        proc.stdin.close()
        proc.wait()
        t.join(timeout=5)

        if proc.returncode != 0:
            err_msg = "\n".join(stderr_lines[-10:])
            raise FFmpegError(f"encode_video failed (rc={proc.returncode}): {err_msg}")

        return output_path

    def _encode_from_dir(
        self, frames_dir: str, output_path: str, fps: float,
        codec: str, crf: int, audio_path: str | None, pixel_format: str,
    ) -> str:
        pattern = str(Path(frames_dir) / "%06d.png")
        cmd = [
            self.ffmpeg, "-y",
            "-framerate", str(fps),
            "-i", pattern,
        ]

        if audio_path:
            cmd += ["-i", audio_path, "-c:a", "copy"]

        cmd += [
            "-c:v", codec, "-crf", str(crf),
            "-pix_fmt", pixel_format,
            output_path,
        ]

        self._run(cmd)
        return output_path

    def _run(
        self, cmd: list[str], capture_stdout: bool = False, timeout: int = 300,
    ) -> subprocess.CompletedProcess:
        """Run a subprocess with timeout."""
        try:
            proc = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=timeout,
            )
        except subprocess.TimeoutExpired:
            raise FFmpegError(f"Command timed out ({timeout}s): {' '.join(cmd[:5])}...")
        if proc.returncode != 0:
            stderr_tail = proc.stderr.strip().split("\n")[-10:]
            raise FFmpegError(
                f"Command failed (rc={proc.returncode}): {' '.join(cmd[:5])}...\n"
                + "\n".join(stderr_tail)
            )
        return proc

    @staticmethod
    def _drain_stderr(proc: subprocess.Popen, lines: list[str]) -> None:
        """Read stderr in background thread to prevent pipe deadlock (S2-R1-014)."""
        try:
            for raw_line in proc.stderr:
                line = raw_line.decode("utf-8", errors="replace").rstrip() if isinstance(raw_line, bytes) else raw_line.rstrip()
                lines.append(line)
                if "error" in line.lower() or "warning" in line.lower():
                    logger.debug(f"ffmpeg: {line}")
        except Exception:
            pass


class StreamingEncoder:
    """Context manager for streaming frame-by-frame encoding via FFmpeg pipe.

    Solves S2-R1-003: no np.stack of all frames in memory.

    Usage:
        with handler.create_streaming_encoder("out.mp4", 1920, 1080, 30.0) as enc:
            for chunk_frames in process_chunks():
                enc.write_frames(chunk_frames)
    """

    def __init__(
        self,
        ffmpeg_path: str,
        output_path: str,
        width: int,
        height: int,
        fps: float = 30.0,
        codec: str = "libx264",
        crf: int = 18,
        pixel_format: str = "yuv420p",
    ) -> None:
        self.ffmpeg_path = ffmpeg_path
        self.output_path = output_path
        self.width = width
        self.height = height
        self.fps = fps
        self.codec = codec
        self.crf = crf
        self.pixel_format = pixel_format
        self._proc: subprocess.Popen | None = None
        self._stderr_lines: list[str] = []
        self._stderr_thread: threading.Thread | None = None
        self._frames_written = 0

    def __enter__(self) -> StreamingEncoder:
        cmd = [
            self.ffmpeg_path, "-y",
            "-f", "rawvideo", "-pix_fmt", "rgb24",
            "-s", f"{self.width}x{self.height}",
            "-r", str(self.fps),
            "-i", "pipe:0",
            "-c:v", self.codec,
        ]

        # VideoToolbox uses bitrate, not CRF (C3)
        if self.crf >= 0:
            cmd += ["-crf", str(self.crf)]
        else:
            cmd += ["-b:v", "10M"]

        cmd += ["-pix_fmt", self.pixel_format, self.output_path]

        self._proc = subprocess.Popen(
            cmd,
            stdin=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        self._stderr_thread = threading.Thread(
            target=FFmpegHandler._drain_stderr,
            args=(self._proc, self._stderr_lines),
            daemon=True,
        )
        self._stderr_thread.start()
        return self

    def write_frames(self, frames: list[np.ndarray] | np.ndarray) -> None:
        """Write frames to the FFmpeg pipe."""
        if self._proc is None or self._proc.stdin is None:
            raise FFmpegError("StreamingEncoder not started")

        if isinstance(frames, np.ndarray) and frames.ndim == 4:
            for i in range(frames.shape[0]):
                self._proc.stdin.write(frames[i].tobytes())
                self._frames_written += 1
        else:
            for frame in frames:
                self._proc.stdin.write(frame.tobytes())
                self._frames_written += 1

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        if self._proc is None:
            return

        try:
            self._proc.stdin.close()
        except Exception:
            pass

        try:
            self._proc.wait(timeout=60)
        except subprocess.TimeoutExpired:
            logger.error("StreamingEncoder: ffmpeg timed out, killing process")
            self._proc.kill()
            self._proc.wait(timeout=5)

        if self._stderr_thread:
            self._stderr_thread.join(timeout=5)

        if exc_type is None and self._proc.returncode != 0:
            err_msg = "\n".join(self._stderr_lines[-10:])
            raise FFmpegError(
                f"StreamingEncoder failed (rc={self._proc.returncode}): {err_msg}"
            )

        logger.info(f"Streaming encode complete: {self._frames_written} frames → {self.output_path}")

    @property
    def frames_written(self) -> int:
        return self._frames_written
