"""VDO.Ninja integration — stream DJ HUD output via WebRTC through the browser.

Architecture:
    DJ HUD  →  WebSocket (JPEG frames)  →  Local HTML page  →  WHIP → VDO.Ninja
                                            ↑ served by tiny HTTP server

The HTML page:
 1. Receives JPEG frames over a local WebSocket from DJ HUD
 2. Draws them onto a <canvas>
 3. Captures the canvas as a MediaStream via captureStream()
 4. Pushes the stream directly to VDO.Ninja via the WHIP protocol
    (standard WebRTC HTTP Ingest Protocol — RFC 9725)
 5. Viewer opens: https://vdo.ninja/?whip=STREAM_ID

This approach requires zero extra Python dependencies — the browser handles all
WebRTC/WHIP complexity using its built-in RTCPeerConnection.
Works on Chrome, Edge, Firefox, Safari (macOS and Windows).
"""

import io
import hashlib
import struct
import socket
import threading
import time
import webbrowser
from http.server import HTTPServer, BaseHTTPRequestHandler
from socketserver import ThreadingMixIn
from typing import Optional

try:
    from PIL import Image
except ImportError:
    Image = None

# ─── Constants ────────────────────────────────────────────────────────────────

DEFAULT_PORT = 7088
JPEG_QUALITY = 75
MAX_FPS = 30

# ─── WebSocket helpers (RFC 6455, minimal server implementation) ──────────────

def _ws_accept_key(key: str) -> str:
    import hashlib as _hashlib, base64 as _b64
    magic = b"258EAFA5-E914-47DA-95CA-5AB5FE842B39"
    return _b64.b64encode(_hashlib.sha1(key.encode() + magic).digest()).decode()


def _ws_encode_frame(data: bytes, opcode: int = 0x02) -> bytes:
    frame = bytearray()
    frame.append(0x80 | opcode)
    length = len(data)
    if length < 126:
        frame.append(length)
    elif length < 65536:
        frame.append(126)
        frame.extend(struct.pack(">H", length))
    else:
        frame.append(127)
        frame.extend(struct.pack(">Q", length))
    frame.extend(data)
    return bytes(frame)


def _ws_decode_frame(data: bytes):
    if len(data) < 2:
        return None
    b0, b1 = data[0], data[1]
    opcode = b0 & 0x0F
    masked = bool(b1 & 0x80)
    length = b1 & 0x7F
    offset = 2
    if length == 126:
        if len(data) < 4: return None
        length = struct.unpack(">H", data[2:4])[0]; offset = 4
    elif length == 127:
        if len(data) < 10: return None
        length = struct.unpack(">Q", data[2:10])[0]; offset = 10
    if masked:
        if len(data) < offset + 4: return None
        mask = data[offset:offset + 4]; offset += 4
    if len(data) < offset + length:
        return None
    payload = bytearray(data[offset:offset + length])
    if masked:
        for i in range(length):
            payload[i] ^= mask[i % 4]
    return opcode, bytes(payload), offset + length


# ─── HTML page — WHIP direct push ────────────────────────────────────────────

def _build_html(ws_port: int, stream_id: str, fps: int) -> str:
    return f"""<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<title>DJ HUD — VDO.Ninja Stream</title>
<style>
  * {{ margin: 0; padding: 0; box-sizing: border-box; }}
  body {{ background: #0f0f1a; color: #e0e0e0; font-family: system-ui, -apple-system, sans-serif;
         display: flex; flex-direction: column; min-height: 100vh; }}

  .header {{ padding: 12px 20px; background: #16213e; border-bottom: 1px solid #2a3a5e;
             display: flex; align-items: center; gap: 16px; flex-shrink: 0; }}
  .header h1 {{ font-size: 15px; font-weight: 600; color: #7ec8ff; white-space: nowrap; }}
  .badge {{ font-size: 12px; padding: 3px 10px; border-radius: 10px; font-weight: 500; }}
  .badge-off {{ background: #3d1414; color: #f87171; }}
  .badge-ready {{ background: #143d1f; color: #4ade80; }}
  .badge-live {{ background: #14293d; color: #60a5fa; animation: pulse 2s infinite; }}
  @keyframes pulse {{ 0%,100% {{ opacity:1; }} 50% {{ opacity:0.7; }} }}

  .toolbar {{ padding: 10px 20px; background: #121225; border-bottom: 1px solid #222;
              display: flex; gap: 12px; align-items: center; flex-wrap: wrap; flex-shrink: 0; }}
  .toolbar label {{ font-size: 13px; color: #999; }}
  .toolbar input {{ background: #1a1a2e; color: #eee; border: 1px solid #333;
                    padding: 5px 10px; border-radius: 4px; font-size: 13px; width: 170px; }}
  .toolbar input:focus {{ border-color: #4a7fff; outline: none; }}

  .btn {{ border: none; padding: 7px 18px; border-radius: 5px; cursor: pointer;
          font-size: 13px; font-weight: 500; transition: all 0.15s; white-space: nowrap; }}
  .btn-go {{ background: #16a34a; color: white; }}
  .btn-go:hover {{ background: #15803d; }}
  .btn-go:disabled {{ background: #333; color: #666; cursor: default; }}
  .btn-stop {{ background: #dc2626; color: white; }}
  .btn-stop:hover {{ background: #b91c1c; }}
  .btn-sm {{ padding: 4px 12px; font-size: 12px; background: #333; color: #ccc; }}
  .btn-sm:hover {{ background: #444; }}

  .viewer-bar {{ padding: 8px 20px; background: #0d1b2a; border-bottom: 1px solid #1a2a3e;
                 display: none; align-items: center; gap: 10px; font-size: 13px; flex-shrink: 0; }}
  .viewer-bar a {{ color: #60a5fa; text-decoration: none; word-break: break-all; }}
  .viewer-bar a:hover {{ text-decoration: underline; }}

  .canvas-wrap {{ flex: 1; display: flex; align-items: center; justify-content: center;
                  padding: 16px; min-height: 0; }}
  canvas {{ border: 1px solid #2a2a3e; border-radius: 4px; max-width: 100%;
            max-height: 100%; background: #000; object-fit: contain; }}

  .footer {{ padding: 8px 20px; font-size: 12px; color: #555; border-top: 1px solid #1a1a2e;
             display: flex; gap: 16px; flex-shrink: 0; }}

  .log {{ padding: 6px 20px; font-size: 11px; color: #666; max-height: 60px;
          overflow-y: auto; border-top: 1px solid #1a1a2e; flex-shrink: 0; }}
  .log .err {{ color: #f87171; }}
  .log .ok {{ color: #4ade80; }}
</style>
</head>
<body>

<div class="header">
  <h1>DJ HUD &rarr; VDO.Ninja</h1>
  <span class="badge badge-off" id="wsBadge">DJ HUD: connecting...</span>
  <span class="badge badge-off" id="whipBadge">WHIP: idle</span>
</div>

<div class="toolbar">
  <label>Stream ID:
    <input type="text" id="streamId" value="{stream_id}" placeholder="pick a unique name"
           spellcheck="false" autocomplete="off">
  </label>
  <button class="btn btn-go" id="startBtn" onclick="toggleWhip()" disabled>Start streaming</button>
</div>

<div class="viewer-bar" id="viewerBar">
  <span>Viewer link:</span>
  <a id="viewerLink" href="#" target="_blank"></a>
  <button class="btn btn-sm" onclick="copyLink()">Copy</button>
</div>

<div class="canvas-wrap">
  <canvas id="c" width="1080" height="80"></canvas>
</div>

<div class="footer">
  <span id="fpsInfo">&mdash; fps</span>
  <span id="sizeInfo">&mdash;</span>
  <span id="bitrateInfo"></span>
</div>
<div class="log" id="log"></div>

<script>
// ═══════════════════════════════════════════════════════════════════
// CONFIG
// ═══════════════════════════════════════════════════════════════════
const WS_URL  = "ws://127.0.0.1:{ws_port}/ws";
const WHIP_BASE = "https://whip.vdo.ninja/";
const VIEW_BASE = "https://vdo.ninja/?whip=";
const CANVAS_FPS = {fps};
const STUN = {{ urls: "stun:stun4.l.google.com:19302" }};

// ═══════════════════════════════════════════════════════════════════
// STATE
// ═══════════════════════════════════════════════════════════════════
const canvas = document.getElementById("c");
const ctx    = canvas.getContext("2d");
let ws = null, wsOk = false;
let pc = null, whipSession = null, streaming = false;
let frameCount = 0, fpsTimer = Date.now();

// ═══════════════════════════════════════════════════════════════════
// LOGGING
// ═══════════════════════════════════════════════════════════════════
function log(msg, cls) {{
  const el = document.getElementById("log");
  const d = document.createElement("div");
  if (cls) d.className = cls;
  d.textContent = new Date().toLocaleTimeString() + "  " + msg;
  el.appendChild(d);
  el.scrollTop = el.scrollHeight;
  while (el.children.length > 50) el.removeChild(el.firstChild);
}}

// ═══════════════════════════════════════════════════════════════════
// WEBSOCKET — receive frames from DJ HUD
// ═══════════════════════════════════════════════════════════════════
function connectWs() {{
  ws = new WebSocket(WS_URL);
  ws.binaryType = "arraybuffer";

  ws.onopen = () => {{
    wsOk = true;
    document.getElementById("wsBadge").textContent = "DJ HUD: connected";
    document.getElementById("wsBadge").className = "badge badge-ready";
    document.getElementById("startBtn").disabled = false;
    log("Connected to DJ HUD", "ok");
  }};

  ws.onmessage = (evt) => {{
    const blob = new Blob([evt.data], {{ type: "image/jpeg" }});
    const url  = URL.createObjectURL(blob);
    const img  = new Image();
    img.onload = () => {{
      if (canvas.width !== img.width || canvas.height !== img.height) {{
        canvas.width  = img.width;
        canvas.height = img.height;
        document.getElementById("sizeInfo").textContent = img.width + " x " + img.height;
      }}
      ctx.drawImage(img, 0, 0);
      URL.revokeObjectURL(url);
      frameCount++;
      const now = Date.now();
      if (now - fpsTimer >= 1000) {{
        document.getElementById("fpsInfo").textContent = frameCount + " fps";
        frameCount = 0;
        fpsTimer = now;
      }}
    }};
    img.src = url;
  }};

  ws.onclose = () => {{
    wsOk = false;
    document.getElementById("wsBadge").textContent = "DJ HUD: disconnected";
    document.getElementById("wsBadge").className = "badge badge-off";
    if (!streaming) document.getElementById("startBtn").disabled = true;
    log("Disconnected from DJ HUD — retrying in 2s");
    setTimeout(connectWs, 2000);
  }};

  ws.onerror = () => {{ try {{ ws.close(); }} catch(_) {{}} }};
}}

// ═══════════════════════════════════════════════════════════════════
// WHIP — push canvas stream to VDO.Ninja
// ═══════════════════════════════════════════════════════════════════

async function startWhip() {{
  const sid = document.getElementById("streamId").value.trim();
  if (!sid) {{
    alert("Enter a Stream ID first.\\n\\nViewers will open:\\n" + VIEW_BASE + "YOUR_ID");
    return;
  }}

  log("Starting WHIP push to " + WHIP_BASE + sid + " ...");

  // Capture canvas as MediaStream
  const stream = canvas.captureStream(CANVAS_FPS);

  // Create peer connection
  pc = new RTCPeerConnection({{ iceServers: [STUN] }});

  // Add video track (send-only)
  for (const track of stream.getVideoTracks()) {{
    pc.addTrack(track, stream);
  }}

  // Prefer H264 for best VDO.Ninja compatibility
  try {{
    const transceivers = pc.getTransceivers();
    for (const t of transceivers) {{
      if (t.sender && t.sender.track && t.sender.track.kind === "video") {{
        const codecs = RTCRtpSender.getCapabilities("video").codecs;
        const h264 = codecs.filter(c => c.mimeType === "video/H264");
        if (h264.length > 0) {{
          t.setCodecPreferences([...h264, ...codecs.filter(c => c.mimeType !== "video/H264")]);
        }}
      }}
    }}
  }} catch(e) {{ /* setCodecPreferences not in all browsers */ }}

  // Monitor connection state
  pc.oniceconnectionstatechange = () => {{
    log("ICE: " + pc.iceConnectionState);
    if (pc.iceConnectionState === "connected" || pc.iceConnectionState === "completed") {{
      document.getElementById("whipBadge").textContent = "WHIP: live";
      document.getElementById("whipBadge").className = "badge badge-live";
    }} else if (pc.iceConnectionState === "failed" || pc.iceConnectionState === "disconnected") {{
      document.getElementById("whipBadge").textContent = "WHIP: " + pc.iceConnectionState;
      document.getElementById("whipBadge").className = "badge badge-off";
      if (pc.iceConnectionState === "failed") {{
        log("ICE connection failed", "err");
      }}
    }}
  }};

  try {{
    // Create and set local offer
    const offer = await pc.createOffer();
    await pc.setLocalDescription(offer);

    // Wait for ICE gathering (WHIP needs complete candidate list)
    await waitForIce(pc, 5000);
    log("ICE gathering complete — sending offer...");

    // WHIP: POST SDP offer to endpoint
    const resp = await fetch(WHIP_BASE + encodeURIComponent(sid), {{
      method: "POST",
      headers: {{ "Content-Type": "application/sdp" }},
      body: pc.localDescription.sdp,
    }});

    if (!resp.ok) {{
      const errBody = await resp.text();
      throw new Error("WHIP " + resp.status + ": " + errBody.slice(0, 200));
    }}

    // Set remote answer
    const answerSdp = await resp.text();
    whipSession = resp.headers.get("Location");
    await pc.setRemoteDescription({{ type: "answer", sdp: answerSdp }});
    log("WHIP connected — stream is live!", "ok");

    // Update UI
    streaming = true;
    document.getElementById("startBtn").textContent = "Stop streaming";
    document.getElementById("startBtn").className = "btn btn-stop";
    document.getElementById("whipBadge").textContent = "WHIP: connecting...";
    document.getElementById("whipBadge").className = "badge badge-ready";

    const viewUrl = VIEW_BASE + encodeURIComponent(sid);
    document.getElementById("viewerLink").href = viewUrl;
    document.getElementById("viewerLink").textContent = viewUrl;
    document.getElementById("viewerBar").style.display = "flex";
    log("Viewer link: " + viewUrl, "ok");

    startBitrateMonitor();

  }} catch(e) {{
    log("WHIP error: " + e.message, "err");
    if (e.message && e.message.includes("Failed to fetch")) {{
      log("Cannot reach whip.vdo.ninja — check your internet connection", "err");
    }}
    cleanupPc();
  }}
}}

function waitForIce(pc, ms) {{
  return new Promise(resolve => {{
    if (pc.iceGatheringState === "complete") {{ resolve(); return; }}
    const t = setTimeout(resolve, ms);
    pc.onicegatheringstatechange = () => {{
      if (pc.iceGatheringState === "complete") {{ clearTimeout(t); resolve(); }}
    }};
  }});
}}

async function stopWhip() {{
  log("Stopping stream...");
  if (whipSession) {{
    try {{ await fetch(whipSession, {{ method: "DELETE" }}); }} catch(_) {{}}
    whipSession = null;
  }}
  cleanupPc();
  streaming = false;
  document.getElementById("startBtn").textContent = "Start streaming";
  document.getElementById("startBtn").className = "btn btn-go";
  document.getElementById("whipBadge").textContent = "WHIP: idle";
  document.getElementById("whipBadge").className = "badge badge-off";
  document.getElementById("viewerBar").style.display = "none";
  document.getElementById("bitrateInfo").textContent = "";
  log("Stopped");
}}

function cleanupPc() {{
  if (pc) {{ try {{ pc.close(); }} catch(_) {{}} pc = null; }}
}}

function toggleWhip() {{
  if (streaming) stopWhip(); else startWhip();
}}

// ═══════════════════════════════════════════════════════════════════
// BITRATE MONITOR
// ═══════════════════════════════════════════════════════════════════
let brInterval = null, prevBytes = 0, prevTs = 0;

function startBitrateMonitor() {{
  if (brInterval) clearInterval(brInterval);
  prevBytes = 0; prevTs = 0;
  brInterval = setInterval(async () => {{
    if (!pc || !streaming) {{ clearInterval(brInterval); return; }}
    try {{
      const stats = await pc.getStats();
      stats.forEach(r => {{
        if (r.type === "outbound-rtp" && r.kind === "video") {{
          if (prevTs > 0) {{
            const kbps = Math.round(((r.bytesSent - prevBytes) * 8) / ((r.timestamp - prevTs) / 1000) / 1000);
            document.getElementById("bitrateInfo").textContent = kbps + " kbps";
          }}
          prevBytes = r.bytesSent;
          prevTs = r.timestamp;
        }}
      }});
    }} catch(_) {{}}
  }}, 2000);
}}

// ═══════════════════════════════════════════════════════════════════
// COPY LINK
// ═══════════════════════════════════════════════════════════════════
function copyLink() {{
  const url = document.getElementById("viewerLink").textContent;
  navigator.clipboard.writeText(url).then(() => {{
    const b = event.target;
    b.textContent = "Copied!";
    setTimeout(() => b.textContent = "Copy", 1500);
  }});
}}

// ═══════════════════════════════════════════════════════════════════
// INIT
// ═══════════════════════════════════════════════════════════════════
connectWs();
log("Ready — enter a Stream ID and click Start streaming");
</script>
</body>
</html>"""


# ═══════════════════════════════════════════════════════════════════════════════
# HTTP + WebSocket Server
# ═══════════════════════════════════════════════════════════════════════════════

class _VDOHandler(BaseHTTPRequestHandler):

    def log_message(self, format, *args):
        pass

    def setup(self):
        super().setup()
        self._is_websocket = False

    def finish(self):
        """Override to prevent closing the socket for WebSocket connections.
        The WS read loop manages the connection lifetime instead."""
        if self._is_websocket:
            # Don't let StreamRequestHandler close rfile/wfile —
            # the WebSocket loop is still using the raw socket.
            return
        super().finish()

    def do_GET(self):
        if self.path == "/ws":
            self._handle_ws_upgrade()
        elif self.path in ("/", "/index.html"):
            self._serve_html()
        elif self.path == "/health":
            self._respond(200, b"ok", "text/plain")
        else:
            self.send_error(404)

    def _respond(self, code, body, content_type):
        self.send_response(code)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _serve_html(self):
        srv = self.server
        html = _build_html(
            ws_port=srv.server_address[1],
            stream_id=srv.stream_id,
            fps=min(MAX_FPS, srv.target_fps),
        ).encode("utf-8")
        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.send_header("Content-Length", str(len(html)))
        self.send_header("Cache-Control", "no-cache")
        self.end_headers()
        self.wfile.write(html)

    def _handle_ws_upgrade(self):
        key = self.headers.get("Sec-WebSocket-Key", "")
        if not key:
            self.send_error(400); return

        # Mark this connection as WebSocket so finish() won't close it
        self._is_websocket = True

        # Write the 101 response directly to the raw socket to avoid
        # buffering issues with BaseHTTPRequestHandler's wfile on Windows.
        # The wfile buffered writer sometimes doesn't flush before we
        # switch to raw WebSocket mode.
        accept = _ws_accept_key(key)
        response = (
            "HTTP/1.1 101 Switching Protocols\r\n"
            "Upgrade: websocket\r\n"
            "Connection: Upgrade\r\n"
            f"Sec-WebSocket-Accept: {accept}\r\n"
            "\r\n"
        )
        try:
            self.request.sendall(response.encode("ascii"))
        except Exception as e:
            print(f"[VDO.Ninja] WS handshake send failed: {e}")
            return

        # Prevent BaseHTTPRequestHandler from trying to handle more HTTP
        # requests on this connection — it's now a WebSocket.
        self.close_connection = True

        print(f"[VDO.Ninja] WebSocket client connected from {self.client_address}")

        srv = self.server
        conn = self.request
        conn.settimeout(1.0)
        with srv._ws_lock:
            srv._ws_clients.add(conn)
        try:
            buf = b""
            while srv._running:
                try:
                    data = conn.recv(4096)
                    if not data: break
                    buf += data
                    while buf:
                        result = _ws_decode_frame(buf)
                        if result is None: break
                        opcode, payload, consumed = result
                        buf = buf[consumed:]
                        if opcode == 0x08: return
                        if opcode == 0x09:
                            conn.sendall(_ws_encode_frame(payload, 0x0A))
                except socket.timeout:
                    continue
                except (ConnectionResetError, ConnectionAbortedError, BrokenPipeError):
                    break
                except OSError:
                    break
        finally:
            with srv._ws_lock:
                srv._ws_clients.discard(conn)
            print(f"[VDO.Ninja] WebSocket client disconnected")


class _ThreadedHTTPServer(ThreadingMixIn, HTTPServer):
    """HTTPServer that handles each request in a new thread."""
    daemon_threads = True

    def shutdown_request(self, request):
        """Override to prevent closing WebSocket connections.
        Normal HTTP requests get closed as usual. WebSocket connections
        are tracked in _ws_clients and stay open."""
        # Check if this socket is an active WebSocket connection
        if hasattr(self, '_ws_clients') and hasattr(self, '_ws_lock'):
            with self._ws_lock:
                if request in self._ws_clients:
                    # Don't close — the WS loop is still using this socket
                    return
        super().shutdown_request(request)


class VDONinjaServer(_ThreadedHTTPServer):
    """Local HTTP + WebSocket server that bridges DJ HUD frames to the browser."""

    def __init__(self, port: int = DEFAULT_PORT, stream_id: str = "",
                 target_fps: int = 15):
        self.stream_id = stream_id
        self.target_fps = target_fps
        self._ws_clients: set = set()
        self._ws_lock = threading.Lock()
        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._last_frame_time = 0.0
        self.allow_reuse_address = True
        super().__init__(("127.0.0.1", port), _VDOHandler)

    def start(self):
        if self._running: return
        self._running = True
        self._thread = threading.Thread(target=self._serve_loop, daemon=True)
        self._thread.start()

    def _serve_loop(self):
        try:
            self.serve_forever(poll_interval=0.5)
        except Exception:
            pass

    def stop(self):
        self._running = False
        self.shutdown()  # stops serve_forever()
        with self._ws_lock:
            for conn in list(self._ws_clients):
                try:
                    conn.sendall(_ws_encode_frame(b"", 0x08))
                    conn.close()
                except Exception:
                    pass
            self._ws_clients.clear()
        try: self.server_close()
        except Exception: pass

    def push_frame(self, pil_image: "Image.Image", quality: int = JPEG_QUALITY):
        with self._ws_lock:
            if not self._ws_clients or Image is None:
                return
        now = time.monotonic()
        if now - self._last_frame_time < 1.0 / max(1, self.target_fps):
            return
        self._last_frame_time = now

        buf = io.BytesIO()
        pil_image.save(buf, format="JPEG", quality=quality, optimize=False)
        frame = _ws_encode_frame(buf.getvalue(), 0x02)
        dead = set()
        with self._ws_lock:
            clients = list(self._ws_clients)
        for conn in clients:
            try:
                conn.sendall(frame)
            except Exception:
                dead.add(conn)
        if dead:
            with self._ws_lock:
                self._ws_clients -= dead

    @property
    def is_running(self):
        return self._running

    @property
    def client_count(self):
        with self._ws_lock:
            return len(self._ws_clients)

    @property
    def url(self):
        return f"http://127.0.0.1:{self.server_address[1]}"

    def open_browser(self):
        webbrowser.open(self.url)
