import asyncio
import contextlib
import json
import logging
import os
import socket
import uuid
from typing import Any

from iqpilot.common.params import Params
from iqpilot.cereal import messaging

from iqpilot.system.webrtc.webrtcd import CerealIncomingMessageProxy, CerealOutgoingMessageProxy, CerealProxyRunner, DynamicPubMaster


def _default_route_ip() -> str | None:
  sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
  try:
    sock.connect(("8.8.8.8", 53))
    return sock.getsockname()[0]
  except OSError:
    return None
  finally:
    sock.close()


class LivestreamBitrateController:
  bitrates = [500_000, 1_500_000, int(os.environ.get("STREAM_BITRATE", 5_000_000))]
  label_to_bitrate = {"low": bitrates[0], "med": bitrates[1], "high": bitrates[2]}
  sample_interval = 0.2
  high_loss = 0.10
  sustained_loss = 0.05
  down_samples = 5

  def __init__(self, get_stats, enabled: bool = True):
    self._get_stats = get_stats
    self._enabled = enabled
    self._auto = True
    self._level = 1
    self._counter = 0
    self._up_samples = 5
    self._previous: tuple[Any, ...] | None = None
    self._task: asyncio.Task | None = None
    self.current_bitrate = self.bitrates[self._level]
    self._publish(self.current_bitrate)

  def start(self) -> None:
    if self._task is None:
      self._task = asyncio.create_task(self.run())

  def stop(self) -> None:
    if self._task is not None:
      self._task.cancel()
    self._task = None

  def enable(self, enabled: bool) -> None:
    self._enabled = enabled

  def set_quality(self, quality: str) -> None:
    if quality in self.label_to_bitrate:
      self._auto = False
      self._publish(self.label_to_bitrate[quality])
    elif quality == "auto":
      self._auto = True

  def _publish(self, bitrate: int) -> None:
    self.current_bitrate = int(bitrate)
    Params().put("LivestreamEncoderBitrate", self.current_bitrate)

  def _sample_loss(self) -> float | None:
    report = next(iter(self._get_stats().values()), None)
    if report is None:
      return None
    current = (report.ssrc, report.fraction_lost, report.packets_lost, report.highest_seq_no, report.jitter, report.lsr, report.dlsr)
    if current == self._previous:
      return None
    self._previous = current
    return report.fraction_lost / 256.0

  async def run(self) -> None:
    while True:
      await asyncio.sleep(self.sample_interval)
      if not self._enabled or not self._auto:
        continue
      loss = self._sample_loss()
      if loss is None:
        continue
      if loss >= self.sustained_loss and self._level > 0:
        self._counter += 1
        if loss >= self.high_loss or self._counter >= self.down_samples:
          self._level -= 1
          self._up_samples *= 2
          self._counter = 0
          self._publish(self.bitrates[self._level])
      elif loss <= 0 and self._level < len(self.bitrates) - 1:
        self._counter -= 1
        if -self._counter >= self._up_samples:
          self._level += 1
          self._counter = 0
          self._publish(self.bitrates[self._level])


class StreamSession:
  shared_pub_master: DynamicPubMaster | None = None

  def __init__(self, sdp: str, cameras: list[str], incoming_services: list[str], outgoing_services: list[str],
               ice_servers: list[dict[str, Any]] | None = None, debug_mode: bool = False, ui_stream: bool = False):
    from iqpilot.system.webrtc.device.native_audio import AudioInputOpusProducer, DebugAudioOpusProducer
    from iqpilot.system.webrtc.device.native_video import DebugVideoStreamTrack, LiveStreamVideoStreamTrack
    from iqpilot.system.webrtc.rtc.builder import WebRTCAnswerBuilder
    from iqpilot.system.webrtc.rtc.info import parse_info_from_offer

    config = parse_info_from_offer(sdp)
    if len(cameras) != config.n_expected_camera_tracks:
      raise ValueError("Incoming stream has misconfigured number of video tracks")
    builder = WebRTCAnswerBuilder(sdp, bind_address=_default_route_ip(), ice_servers=ice_servers or [])
    video_track_type = DebugVideoStreamTrack if debug_mode else LiveStreamVideoStreamTrack
    self.video_tracks = [video_track_type(camera) for camera in cameras]
    for camera, track in zip(cameras, self.video_tracks, strict=True):
      builder.add_video_stream(camera, track)

    audio_track_type = DebugAudioOpusProducer if debug_mode else AudioInputOpusProducer
    self.audio_output = audio_track_type() if config.expected_audio_track else None
    if self.audio_output is not None:
      builder.add_audio_stream(self.audio_output)
    if config.incoming_audio_track:
      builder.offer_to_receive_audio_stream()
    self.stream = builder.stream()

    self.identifier = str(uuid.uuid4())
    self.incoming_bridge_services = incoming_services
    if incoming_services and self.shared_pub_master is None:
      StreamSession.shared_pub_master = DynamicPubMaster([])
    self.incoming_bridge = CerealIncomingMessageProxy(self.shared_pub_master) if self.shared_pub_master is not None and incoming_services else None
    self.outgoing_bridge = CerealOutgoingMessageProxy(messaging.SubMaster(outgoing_services)) if outgoing_services else None
    self.outgoing_bridge_runner = CerealProxyRunner(self.outgoing_bridge) if self.outgoing_bridge is not None else None
    self.ui_stream_requested = ui_stream
    self.ui_stream_runner: CerealProxyRunner | None = None
    self.audio_input_proxy = None
    self.audio_recv_requested = config.incoming_audio_track
    self.bitrate_controller = LivestreamBitrateController(self.stream.get_receiver_report_stats)
    self.run_task: asyncio.Task | None = None
    self._cleanup_lock = asyncio.Lock()
    self._cleanup_done = False
    self.logger = logging.getLogger("webrtcd")

  def start(self) -> None:
    self.run_task = asyncio.create_task(self.run())

  async def get_answer(self):
    return await self.stream.start()

  async def stop_async(self) -> None:
    if self.run_task is not None and not self.run_task.done() and self.run_task is not asyncio.current_task():
      self.run_task.cancel()
      with contextlib.suppress(asyncio.CancelledError):
        await self.run_task
    self.run_task = None
    await self.post_run_cleanup()

  def add_ice_candidate(self, candidate: Any) -> None:
    self.stream.add_ice_candidate(candidate)

  def message_handler(self, message: bytes | str) -> None:
    try:
      payload = json.loads(message)
    except (TypeError, ValueError):
      return
    if not isinstance(payload, dict):
      return
    message_type = payload.get("type")
    if message_type == "timingSei":
      for track in self.video_tracks:
        track.timing_sei_enabled = bool(payload.get("enabled", False))
    elif message_type == "setQuality":
      self.bitrate_controller.set_quality(str(payload.get("quality", "auto")))
    elif message_type == "setAudioEnabled" and self.audio_output is not None:
      self.audio_output.enable(bool(payload.get("enabled", True)))
    elif message_type == "switchCamera":
      for track in self.video_tracks:
        track.switch_camera(str(payload.get("camera", "")))
    elif message_type == "setUiStream":
      self.set_ui_stream(bool(payload.get("enabled", False)))
    elif self.incoming_bridge is not None:
      try:
        self.incoming_bridge.send(message)
      except Exception:
        self.logger.exception("Cereal incoming proxy failure")

  def set_ui_stream(self, enabled: bool) -> None:
    if enabled:
      if self.ui_stream_runner is not None or not self.stream.has_messaging_channel():
        return
      from iqpilot.system.webrtc.ui_stream import UIStreamMessageProxy
      proxy = UIStreamMessageProxy(bitrate_getter=lambda: self.bitrate_controller.current_bitrate)
      proxy.add_channel(self.stream.get_messaging_channel())
      self.ui_stream_runner = CerealProxyRunner(proxy)
      self.ui_stream_runner.start()
    elif self.ui_stream_runner is not None:
      self.ui_stream_runner.stop()
      self.ui_stream_runner = None

  async def run(self) -> None:
    try:
      await self.stream.wait_for_connection()
      if self.stream.has_messaging_channel():
        self.stream.set_message_handler(self.message_handler)
        if self.incoming_bridge is not None:
          await self.shared_pub_master.add_services_if_needed(self.incoming_bridge_services)
        if self.outgoing_bridge_runner is not None:
          self.outgoing_bridge_runner.proxy.add_channel(self.stream.get_messaging_channel())
          self.outgoing_bridge_runner.start()
        if self.ui_stream_requested:
          self.set_ui_stream(True)
      if self.audio_recv_requested and self.stream.has_incoming_audio_track():
        from iqpilot.system.webrtc.device.native_audio import IncomingOpusCerealProxy
        self.audio_input_proxy = IncomingOpusCerealProxy(self.stream.get_incoming_audio_track())
        self.audio_input_proxy.start()
      self.bitrate_controller.start()
      await self.stream.wait_for_disconnection()
    except Exception:
      self.logger.exception("libdatachannel stream session failure")
    finally:
      await self.post_run_cleanup()

  async def post_run_cleanup(self) -> None:
    async with self._cleanup_lock:
      if self._cleanup_done:
        return
      self._cleanup_done = True
      self.bitrate_controller.stop()
      if self.ui_stream_runner is not None:
        self.ui_stream_runner.stop()
        self.ui_stream_runner = None
      if self.outgoing_bridge_runner is not None:
        self.outgoing_bridge_runner.stop()
      if self.audio_input_proxy is not None:
        await self.audio_input_proxy.stop()
      for track in self.video_tracks:
        track.stop()
      await self.stream.stop()