import asyncio
from collections.abc import AsyncIterator

from loguru import logger

from restreamer.constants import FFMPEG_READ_CHUNK_SIZE, FFMPEG_TERMINATE_TIMEOUT_SECONDS
from restreamer.resolver import ResolvedStream

_RECONNECT_INPUT_ARGS = ["-reconnect", "1", "-reconnect_streamed", "1", "-reconnect_delay_max", "5"]


def build_ffmpeg_args(stream: ResolvedStream) -> list[str]:
    args = ["ffmpeg", "-loglevel", "warning", "-nostdin"]

    header_lines = "".join(f"{key}: {value}\r\n" for key, value in stream.http_headers.items())

    if header_lines:
        args += ["-headers", header_lines]
    args += [*_RECONNECT_INPUT_ARGS, "-i", stream.video_url]

    if stream.audio_url:
        if header_lines:
            args += ["-headers", header_lines]
        args += [*_RECONNECT_INPUT_ARGS, "-i", stream.audio_url]
        args += ["-map", "0:v:0", "-map", "1:a:0"]

    args += [
        "-c",
        "copy",
        "-f",
        "mpegts",
        "-mpegts_flags",
        "+resend_headers",
        "pipe:1",
    ]
    return args


class FfmpegStreamProcess:
    """Wraps a running ffmpeg process that copies a resolved stream to stdout as MPEG-TS."""

    def __init__(self, stream: ResolvedStream, channel_id: str) -> None:
        self._stream = stream
        self._channel_id = channel_id
        self._process: asyncio.subprocess.Process | None = None
        self._stderr_task: asyncio.Task[None] | None = None

    async def start(self) -> None:
        args = build_ffmpeg_args(self._stream)
        logger.debug("Starting ffmpeg for channel '{}'", self._channel_id)
        self._process = await asyncio.create_subprocess_exec(
            *args,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )
        self._stderr_task = asyncio.create_task(self._drain_stderr())

    async def iter_chunks(self) -> AsyncIterator[bytes]:
        if self._process is None or self._process.stdout is None:
            raise RuntimeError("ffmpeg process has not been started")

        while True:
            chunk = await self._process.stdout.read(FFMPEG_READ_CHUNK_SIZE)
            if not chunk:
                return
            yield chunk

    async def wait(self) -> int:
        if self._process is None:
            raise RuntimeError("ffmpeg process has not been started")
        return await self._process.wait()

    async def stop(self) -> None:
        if self._process is None or self._process.returncode is not None:
            return

        # wait() only returns once every pipe hits EOF, so unread stdout left over
        # from the consumer breaking out of iter_chunks() must be drained, or it hangs forever.
        stdout_drain_task = asyncio.create_task(self._drain_stdout())

        self._process.terminate()
        try:
            await asyncio.wait_for(self._process.wait(), timeout=FFMPEG_TERMINATE_TIMEOUT_SECONDS)
        except TimeoutError:
            logger.warning("ffmpeg for channel '{}' did not exit in time, killing it", self._channel_id)
            self._process.kill()
            await self._process.wait()

        await stdout_drain_task
        if self._stderr_task is not None:
            await self._stderr_task

    async def _drain_stdout(self) -> None:
        if self._process is None or self._process.stdout is None:
            return
        while await self._process.stdout.read(FFMPEG_READ_CHUNK_SIZE):
            pass

    async def _drain_stderr(self) -> None:
        if self._process is None or self._process.stderr is None:
            return
        async for line in self._process.stderr:
            logger.debug("[ffmpeg:{}] {}", self._channel_id, line.decode(errors="replace").rstrip())