#!/usr/bin/env python3
"""
CHT Sender: muxed video + audio over TCP/mpegts
Source: Wayland (kmsgrab) + desktop audio + webcam mic
Usage: sudo python3 stream_av.py [RECEIVER_IP] [PORT]

Requires: sudo for kmsgrab, PulseAudio for audio capture
Auto-restarts on stall or fatal kmsgrab errors (DRM plane format change).
"""

import os
import re
import resource
import signal
import subprocess
import sys
import tempfile
import threading
import time

RECEIVER_IP   = sys.argv[1] if len(sys.argv) > 1 else "mcrndeb"
PORT          = sys.argv[2] if len(sys.argv) > 2 else "4444"
STALL_TIMEOUT = 10

REAL_UID = os.environ.get("SUDO_UID") or str(os.getuid())
os.environ["PULSE_SERVER"] = f"unix:/run/user/{REAL_UID}/pulse/native"

FATAL_PATTERNS = ["framebuffer format changed", "Error during demuxing"]


def get_monitor_source():
    try:
        out = subprocess.check_output(["pactl", "info"], stderr=subprocess.DEVNULL).decode()
        for line in out.splitlines():
            if "Default Sink:" in line:
                return line.split()[-1] + ".monitor"
    except Exception:
        pass
    return None


def get_webcam_mic():
    try:
        out = subprocess.check_output(
            ["pactl", "list", "short", "sources"], stderr=subprocess.DEVNULL
        ).decode()
        for line in out.splitlines():
            if "C922" in line:
                return line.split()[1]
    except Exception:
        pass
    return None


def build_args(monitor, webcam_mic, progress_file):
    args = [
        "ffmpeg",
        "-init_hw_device", "drm=drm:/dev/dri/card0",
        "-init_hw_device", "vaapi=va@drm",
        "-thread_queue_size", "64",
        "-device", "/dev/dri/card0", "-f", "kmsgrab", "-framerate", "30", "-i", "-",
        "-thread_queue_size", "1024", "-f", "pulse", "-i", monitor,
    ]
    if webcam_mic:
        args += [
            "-thread_queue_size", "1024", "-f", "pulse", "-i", webcam_mic,
            "-filter_complex", "[1:a][2:a]amix=inputs=2:duration=longest[aout]",
            "-map", "0:v", "-map", "[aout]",
        ]
    args += [
        "-vf", "hwmap=derive_device=vaapi,scale_vaapi=w=1920:h=1080:format=nv12,fps=30",
        "-c:v", "h264_vaapi", "-qp", "20", "-g", "30", "-keyint_min", "30", "-bf", "0",
        "-c:a", "aac", "-b:a", "128k",
        "-max_muxing_queue_size", "64",
        "-flush_packets", "1", "-fflags", "nobuffer", "-muxdelay", "0", "-muxpreload", "0",
        "-f", "mpegts", f"tcp://{RECEIVER_IP}:{PORT}",
        "-hide_banner", "-progress", progress_file,
    ]
    return args


def read_progress(progress_file, key):
    try:
        with open(progress_file) as f:
            content = f.read()
        matches = re.findall(rf"^{key}=(\d+)", content, re.MULTILINE)
        return int(matches[-1]) if matches else 0
    except Exception:
        return 0


def run_once(args, progress_file):
    """Run the pipeline, monitor it. Returns True if should restart."""
    fatal = threading.Event()
    print(f"[cmd] {' '.join(args)}", flush=True)

    proc = subprocess.Popen(args, stderr=subprocess.PIPE, start_new_session=True)

    def watch_stderr():
        for raw in proc.stderr:
            line = raw.decode("utf-8", errors="replace").rstrip()
            if line:
                print(f"[ffmpeg] {line}", flush=True)
                if any(p in line for p in FATAL_PATTERNS):
                    print("Fatal kmsgrab error — restarting immediately", flush=True)
                    fatal.set()

    threading.Thread(target=watch_stderr, daemon=True).start()

    last_bytes = last_frame = 0
    stall_since = time.monotonic()

    def kill_proc():
        try:
            os.killpg(os.getpgid(proc.pid), signal.SIGINT)
        except Exception:
            proc.kill()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()

    try:
        while proc.poll() is None:
            time.sleep(1)

            if fatal.is_set():
                kill_proc()
                break

            cur_bytes = read_progress(progress_file, "total_size")
            cur_frame = read_progress(progress_file, "frame")

            if cur_bytes > last_bytes or cur_frame > last_frame:
                last_bytes, last_frame = cur_bytes, cur_frame
                stall_since = time.monotonic()

            if time.monotonic() - stall_since > STALL_TIMEOUT:
                print(f"Stalled at frame {last_frame} / {last_bytes}B — restarting", flush=True)
                kill_proc()
                break
    except KeyboardInterrupt:
        kill_proc()
        raise

    rc = proc.wait()
    print(f"ffmpeg exited rc={rc}", flush=True)
    return True  # always restart; only KeyboardInterrupt stops the loop


def main():
    try:
        resource.setrlimit(resource.RLIMIT_NOFILE, (65536, 65536))
    except Exception:
        pass

    monitor   = get_monitor_source()
    webcam_mic = get_webcam_mic()

    print(f"Monitor source : {monitor}")
    print(f"Webcam mic     : {webcam_mic or 'not found'}")
    print(f"Streaming to   : {RECEIVER_IP}:{PORT}")

    if not monitor:
        sys.exit("ERROR: could not find monitor source")

    with tempfile.NamedTemporaryFile(delete=False, suffix=".progress") as f:
        progress_file = f.name

    try:
        while True:
            print(f"--- Starting sender {time.strftime('%c')} ---", flush=True)
            open(progress_file, "w").close()

            args = build_args(monitor, webcam_mic, progress_file)
            run_once(args, progress_file)
            print("Restarting in 2s...", flush=True)
            time.sleep(2)
    except KeyboardInterrupt:
        print("\nStopped.", flush=True)
    finally:
        try:
            os.unlink(progress_file)
        except FileNotFoundError:
            pass


if __name__ == "__main__":
    main()