import queue
import subprocess
import threading
import time
import os
import signal
from copy import deepcopy

from django.conf import settings
from insta360.rtmp import Client

from whitebox.db import get_model
from whitebox.plugins import get_class
from whitebox.utils import get_plugin_logger
from utils.daemons import BaseDaemonCommand

from whitebox_plugin_device_insta360.devices import Insta360Base


logger = get_plugin_logger(__name__)


DeviceConnection = get_model("device.DeviceConnection")
device_manager = get_class("device.DeviceManager")


class FrameBuffer(queue.Queue):
    """
    Queue implementation that can auto-evict oldest items when trying to put in
    a new item while the queue is currently full.

    It is reliable ONLY when there is a SINGLE producer (i.e. thread pushing
    frames into it).
    """

    def put_safe(self, item):
        if self.full():
            try:
                # If full, evict the oldest item
                self.get_nowait()
            except queue.Empty:
                # Even if it was full, guard against empty queue in case of race
                # condition where consumers already cleared it by now
                pass

        return self.put_nowait(item)


class FFmpegManager:
    def __init__(
        self,
        frame_source: queue.Queue,
        stream_target: str,
        register_subprocess,
    ):
        self._frame_source = frame_source
        self._stream_target = stream_target
        self._register_subprocess = register_subprocess

        self._is_running_event = threading.Event()
        self._thread = None

    def _check_is_running(self):
        return self._is_running_event.is_set()

    def get_ffmpeg_command(self, output_target):
        # prevent ruff from formatting this into an unreadable mess
        # fmt: off
        return [
            "ffmpeg",

            # "-use_wallclock_as_timestamps", "1",

            # Decoder settings
            "-fflags", "+genpts+discardcorrupt",
            "-flags", "low_delay",

            "-f", "h264",
            "-i", "pipe:0",  # input from stdin

            "-r", "30",

            # Filter settings
            "-filter_complex",
            f"[0:v]scale=1440:720,v360=dfisheye:e:ih_fov=193:iv_fov=193[out]",
            "-map", "[out]",

            # Encoder settings
            "-c:v", "libx264",
            "-preset", "ultrafast",
            "-tune", "zerolatency",
            "-pix_fmt", "yuv420p",
            "-profile:v", "main",
            "-level", "3.1",

            # Bitrate settings
            # This ensures that keyframes are pushed frequently enough to keep the
            # streaming latency low
            "-g", "30",  # keyframe interval
            "-keyint_min", "30",  # minimum keyframe interval
            "-sc_threshold", "0",  # disable scene change detection

            "-muxpreload", "0",
            "-muxdelay", "0",

            "-f", "flv",  # output format
            "-rtmp_live", "live",

            # e.g.: rtmp://srs-dev/live/livestream
            output_target,
        ]
        # fmt: on

    def _start_process(self):
        proc = subprocess.Popen(
            self.get_ffmpeg_command(self._stream_target),
            stdin=subprocess.PIPE,
            # Keep stdout discarded; keep stderr for diagnostics.
            # stdout=subprocess.DEVNULL,
            # stderr=subprocess.DEVNULL,
            # Ensure that stdin is not buffered for lowest latency
            bufsize=0,
            # Start a new session so we can reliably kill the whole process
            # group when needed
            start_new_session=True,
        )
        self._register_subprocess(proc)

        return proc

    def _feed_frames_to_ffmpeg(self, proc):
        # Feed frames from the source queue into FFmpeg stdin
        # This method will return when either:
        #   - stop event is set
        #   - FFmpeg process exits
        #   - stdin breaks (BrokenPipeError)
        stdin = proc.stdin
        if stdin is None:
            return

        buffer = bytearray()
        last_flush = time.monotonic()
        flush_threshold_bytes = 256 * 1024  # 256KB
        flush_threshold_seconds = 10 / 1000  # 10ms

        while self._check_is_running():
            # If FFmpeg already exited, stop feeding.
            if proc.poll() is not None:
                return

            try:
                # This effectively sleeps for up to 0.5s and avoids busy-waiting
                # when the queue is empty
                chunk = self._frame_source.get(timeout=0.5)
                buffer.extend(chunk)
            except queue.Empty:
                # In case of a timeout, we'll just flush below
                chunk = None

            should_flush = buffer and (
                # Flush every time threshold
                time.monotonic() - last_flush >= flush_threshold_seconds
                # or when buffer is full
                or len(buffer) >= flush_threshold_bytes
                # or when we got nothing from the stream but there's data
                # in the buffer
                or chunk is None
            )
            if not should_flush:
                continue

            try:
                stdin.write(buffer)
            except BrokenPipeError:
                logger.error("FFmpeg stdin - broken pipe")
                return
            except Exception as e:
                # Any unexpected stdin error: stop feeding; the supervisor loop
                # will restart the process.
                logger.exception(f"FFmpeg stdin - unexpected error: {e}")
                return

            buffer.clear()
            last_flush = time.monotonic()

    def _clear_frame_queue(self):
        drop_count = 0

        while True:
            try:
                self._frame_source.get_nowait()
                drop_count += 1
            except queue.Empty:
                break

        if drop_count:
            logger.warning(f"Dropped {drop_count} frames from the queue")

    def _terminate_process(self, proc, timeout=2.0):
        # In case process is already out, abort
        if proc.poll() is not None:
            return

        # Close stdin to signal EOF to ffmpeg.
        try:
            if proc.stdin:
                proc.stdin.close()
        except Exception as e:
            logger.exception(f"Failed to close stdin: {e}")

        # Try to stop process gracefully
        try:
            proc.terminate()
        except Exception as e:
            logger.exception(f"Failed to terminate proc: {e}")

        # Wait for termination
        try:
            proc.wait(timeout=timeout)
            return
        except Exception as e:
            logger.exception(f"Failed to wait for proc: {e}")

        # In case the process does not want to terminate easily, after the wait
        # timeout, try to kill the whole process group (created via start_new_session)
        try:
            os.killpg(proc.pid, signal.SIGKILL)
        except Exception as e:
            logger.exception(f"Failed to kill process group: {e}")
            try:
                proc.kill()
            except Exception as e2:
                logger.exception(f"Failed to kill proc: {e2}")

        # Wait for the kill to take effect
        try:
            proc.wait(timeout=timeout)
        except Exception as e:
            logger.exception(f"FFmpeg is still alive after kill: {e}")

    def run_until_stopped(self):
        # Ensure FFmpeg is running until stop event is set - if it dies, restart
        # it.
        while self._check_is_running():
            try:
                proc = self._start_process()
            except Exception as e:
                logger.exception("Error starting ffmpeg:", e)
                time.sleep(0.5)
                continue

            # Clear queue right before we start feeding frames to ensure we do
            # not process stale frames for minimum latency

            # This is commented out while FFmpeg stalling investigation is in
            # progress to ensure that we do not lose any frames in the meantime,
            # as it might've caused some errors
            #
            # self._clear_frame_queue()

            # Feed frames until we are stopping, ffmpeg exits, or stdin breaks
            self._feed_frames_to_ffmpeg(proc)

            self._terminate_process(proc)

    def start(self):
        if self._check_is_running():
            return

        self._is_running_event.set()
        self._thread = threading.Thread(
            target=self.run_until_stopped,
            daemon=True,
        )
        self._thread.name = f"ffmpeg-{self._stream_target}"
        self._thread.start()

    def stop(self, join=True):
        if not self._is_running_event.is_set():
            return

        self._is_running_event.clear()
        if join:
            self._thread.join()


class CameraSession:
    def __init__(
        self,
        device_connection,
        register_subprocess,
        queue_max_size=1000,
    ):
        self._device_connection = device_connection
        self._register_subprocess = register_subprocess
        self._is_connected = False
        self._lock = threading.Lock()

        self._client = None
        self._should_record = False
        self._recording_status = None

        self._ffmpeg = None

        self._run_thread = None
        self._is_running = threading.Event()

        # Add queue's max size to avoid going OOM in case there's a hiccup
        # somewhere in the pipeline.
        self._processing_queue = FrameBuffer(maxsize=queue_max_size)

    # region video processing

    async def on_video_stream_callback(self, data: bytes, **kwargs):
        # Async callback for Insta360 client that pushes frames to the
        # processing queue. In case a queue becomes full for whatever reason,
        # oldest frames are dropped
        self._processing_queue.put_safe(data)

    # endregion video processing

    # region video recording

    def _get_capture_status(self):
        response = self._client.get_capture_current_status()
        # When not recording, `status` is an empty dict. In those cases,
        # explicitly return `None` instead
        return response.get("status") or None

    def _ensure_video_recording(self):
        if self._recording_status:
            logger.info("Recording start requested, but already recording")
            return

        # In case the camera is already in video mode, this operation will be
        # super quick, but in case it's changing the mode, it might take like 1s
        self._client.set_camera_mode("VIDEO_NORMAL")

        self._client.start_capture()
        self._recording_status = True

    def _ensure_video_not_recording(self):
        if not self._recording_status:
            logger.info("Recording stop requested, but already not recording")
            return

        self._client.stop_capture()
        self._recording_status = False

    # endregion video recording

    # region notifying

    def on_device_connected(self):
        with self._lock:
            if self._is_connected:
                return

            self._is_connected = True

            # Open the client and start streaming again
            self._client.open()
            self._recording_status = self._get_capture_status()

            self._client.start_preview_stream()
            # Start FFmpeg to stream
            self._ffmpeg.start()

    def on_device_disconnected(self):
        with self._lock:
            if not self._is_connected:
                return

            self._is_connected = False

            # Close the client so everything is cleaned up
            self._client.close()
            # Close FFmpeg
            self._ffmpeg.stop()

    def notify_should_record(self):
        with self._lock:
            self._should_record = True
            self._ensure_video_recording()

    def notify_should_not_record(self):
        with self._lock:
            self._should_record = False
            self._ensure_video_not_recording()

    # endregion notifying

    async def on_disconnect_callback(self, **kwargs):
        logger.warning("Received `on_disconnect` event, closing client")
        self.on_device_disconnected()

    def _init_client(self):
        c = Client()
        c.on_video_stream(wait=True)(self.on_video_stream_callback)
        c.on_disconnect(wait=True)(self.on_disconnect_callback)
        return c

    def run(self):
        if self._is_running.is_set():
            raise Exception

        self._is_running.set()
        dc = self._device_connection

        stream_name = f"device_connection_{dc.pk}"
        stream_target = f"rtmp://{settings.SRS_HOST}/live/{stream_name}"
        self._client = self._init_client()

        self._ffmpeg = FFmpegManager(
            frame_source=self._processing_queue,
            stream_target=stream_target,
            register_subprocess=self._register_subprocess,
        )

        self._is_connected = dc.connection_status == dc.ConnectionStatus.CONNECTED

        if not self._is_connected:
            return

        self._client.open()
        self._recording_status = self._get_capture_status()

        self._client.start_preview_stream()
        self._ffmpeg.start()

    def shutdown(self):
        self._is_running.clear()
        self._client.close()
        self._ffmpeg.stop()


class Command(BaseDaemonCommand):
    help = "Insta360 Camera Daemon"

    events_to_subscribe = [
        "observation.flight.start",
        "observation.flight.end",
        "observation.device.connection_status.update",
    ]

    def _get_all_insta360_device_connections(self):
        targets = [
            codename
            for codename, klass in device_manager.get_device_classes().items()
            if issubclass(klass, Insta360Base)
        ]

        return DeviceConnection.objects.filter(
            codename__in=targets,
        )

    def _refresh_camera_sessions(self):
        # Checks if there's a camera session for every single eligible
        # (Insta360-linked) DeviceConnection object - check for new ones so that
        # their sessions can be started, and check for the ones that don't exist
        # anymore (e.g. deleted) so they can be stopped
        device_connections = self._get_all_insta360_device_connections()
        handled = set()

        # For each one of them, start a streaming thread
        for dc in device_connections:
            handled.add(dc.pk)
            if dc.pk in self.camera_sessions:
                continue

            self.stdout.write(
                f"Starting camera session for {dc.codename} (ID: {dc.pk})...",
            )
            cs = CameraSession(
                dc,
                register_subprocess=self.register_subprocess,
            )
            self.camera_sessions[dc.id] = cs
            cs.run()
            self.stdout.write("  Done")

        unhandled = set(self.camera_sessions.keys()) - handled
        for pk in unhandled:
            self.stdout.write(
                f"Stopping camera session for {pk}...",
            )
            cs = self.camera_sessions[pk]
            cs.shutdown()
            del self.camera_sessions[pk]
            self.stdout.write("  Done")

    # region event processing

    def process_observation_flight_start(self, payload):
        # When a flight is in progress, all cameras should be recording
        for cs in self.camera_sessions.values():
            cs.notify_should_record()

    def process_observation_flight_end(self, payload):
        # When no flight in progress, no cameras should be recording
        for cs in self.camera_sessions.values():
            cs.notify_should_not_record()

    def process_device_connection_status_update(self, payload):
        data = payload["data"]
        pk = data["id"]
        cs = self.camera_sessions.get(pk)

        if not cs:
            return

        status = data["connection_status"]
        # Only final states are taken into account - transition states, like
        # connecting/disconnecting are not
        if status == "connected":
            cs.on_device_connected()
        elif status in ["disconnected", "failed"]:
            cs.on_device_disconnected()

    def _process_incoming_message(self, message):
        self.stdout.write(f"Received message: {message}")

        data = deepcopy(message)
        type_ = data.pop("type", None)

        if type_ == "observation.flight.start":
            handler = self.process_observation_flight_start
        elif type_ == "observation.flight.end":
            handler = self.process_observation_flight_end
        elif type_ == "observation.device.connection_status.update":
            handler = self.process_device_connection_status_update
        else:
            self.stderr.write(f"Unknown message type: {type_}")
            return

        try:
            handler(data)
        except Exception as e:
            logger.exception(f"Error processing message: {str(e)}")

    def _process_message_queue(self):
        while True:
            try:
                message = self.events_queue.get(block=False)
            except queue.Empty:
                break

            try:
                self._process_incoming_message(message)
            except Exception as e:
                logger.exception("Failed to process incoming message")

    # endregion event processing

    def run_daemon(self, *args, **options):
        self.camera_sessions = {}
        last_camera_session_refresh = 0

        while not self.daemon_shutdown_event.is_set():
            # In case there are incoming messages in the queue, process them
            # immediately until the queue is empty
            self._process_message_queue()

            if time.monotonic() - last_camera_session_refresh > 5:
                last_camera_session_refresh = time.monotonic()
                self._refresh_camera_sessions()

            # Effectively a sleep(1) with shutdown event check
            self.daemon_shutdown_event.wait(1)

        # Ensure all camera sessions are stopped
        for cs in self.camera_sessions.values():
            cs.shutdown()
