#!/usr/bin/env python3
import socket, subprocess, sys, time, signal, threading, select, math
from zeroconf import ServiceInfo, Zeroconf

# Configuration
CONF = {
    'video_port': 5000, 'audio_port': 5001,
    'video_dev': '/dev/video20', 'audio_sink': 'AirLink_Sink', 'audio_src': 'AirLink_Mic'
}

class AirLinkServer:
    def __init__(self):
        self.ip = self.get_lan_ip()
        self.running = False
        self.procs = {'video': None, 'audio': None, 'ffmpeg_audio': None}
        self.sockets = {'video': None, 'audio': None}
        self.stats = {'v_bytes': 0, 'a_bytes': 0, 'v_pkts': 0, 'a_pkts': 0, 'last_time': time.time()}
        self.zc = None
        self.lock = threading.Lock()

    def get_lan_ip(self):
        """Finds LAN IP by prioritizing physical interfaces over VPN/Tun interfaces."""
        try:
            # List global IPv4 addresses, prioritize eth/wlan
            cmd = "ip -4 -o addr show scope global | awk '{print $2,$4}'"
            out = subprocess.check_output(cmd, shell=True).decode().strip().split('\n')
            candidates = []
            for line in out:
                if not line: continue
                iface, ip = line.split()[0], line.split()[1].split('/')[0]
                if iface.startswith(('e', 'w')): return ip # Found physical interface
                candidates.append(ip)
            return candidates[0] if candidates else '127.0.0.1'
        except:
            # Fallback to standard socket connect
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            try: s.connect(("8.8.8.8", 80)); return s.getsockname()[0]
            except: return '127.0.0.1'
            finally: s.close()

    def setup_system(self):
        print(f"[INIT] Configuring system on {self.ip}...")
        try:
            subprocess.run("sudo modprobe -r v4l2loopback; sudo modprobe v4l2loopback video_nr=20 card_label=AirLink_Cam exclusive_caps=1", shell=True, check=True, stderr=subprocess.DEVNULL)
            print(f"[INFO] Video Device: {CONF['video_dev']}")
            
            subprocess.run(f"pactl load-module module-null-sink sink_name={CONF['audio_sink']} sink_properties=device.description='AirLink_Audio' rate=48000 channels=1", shell=True, stderr=subprocess.DEVNULL)
            subprocess.run(f"pactl load-module module-remap-source source_name={CONF['audio_src']} master={CONF['audio_sink']}.monitor source_properties=device.description='AirLink_Mic' channels=1", shell=True, stderr=subprocess.DEVNULL)
            print(f"[INFO] Audio Source: {CONF['audio_src']}")
            return True
        except Exception as e:
            print(f"[ERR] Setup failed: {e}")
            return False

    def paint_black(self):
        with self.lock:
            if self.procs['video']:
                self.procs['video'].terminate()
                self.procs['video'] = None
            try:
                subprocess.run(['ffmpeg', '-hide_banner', '-loglevel', 'error', '-f', 'lavfi', '-i', 'color=black:s=1280x720', '-pix_fmt', 'yuyv422', '-frames:v', '5', '-f', 'v4l2', CONF['video_dev']], check=False)
                print("[INFO] Video stream stopped (Black Screen)")
            except: pass

    def start_pipeline(self, p_type):
        with self.lock:
            if p_type == 'video' and (not self.procs['video'] or self.procs['video'].poll() is not None):
                self.procs['video'] = subprocess.Popen(
                    ['ffmpeg', '-fflags', 'nobuffer', '-flags', 'low_delay', '-analyzeduration', '0', '-probesize', '32',
                     '-hide_banner', '-loglevel', 'fatal', '-f', 'h264', '-i', 'pipe:0',
                     '-vf', 'scale=1280:720:force_original_aspect_ratio=decrease,pad=1280:720:(ow-iw)/2:(oh-ih)/2',
                     '-vsync', '0', '-copyts', '-pix_fmt', 'yuyv422', '-f', 'v4l2', CONF['video_dev']],
                    stdin=subprocess.PIPE, stderr=subprocess.DEVNULL, bufsize=0)
            
            elif p_type == 'audio' and (not self.procs['audio'] or self.procs['audio'].poll() is not None):
                self.procs['ffmpeg_audio'] = subprocess.Popen(
                    ['ffmpeg', '-hide_banner', '-loglevel', 'fatal', '-fflags', 'nobuffer', '-f', 's16le', '-ar', '48000', '-ac', '1', '-i', 'pipe:0', '-f', 's16le', '-ar', '48000', '-ac', '1', 'pipe:1'],
                    stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, bufsize=0)
                self.procs['audio'] = subprocess.Popen(
                    ['pacat', '--playback', f"--device={CONF['audio_sink']}", '--rate=48000', '--channels=1', '--format=s16le', '--latency-msec=20'],
                    stdin=self.procs['ffmpeg_audio'].stdout, stderr=subprocess.DEVNULL, bufsize=0)

    def log_stats(self):
        now = time.time()
        delta = now - self.stats['last_time']
        if delta >= 1.0:
            v_mbps = (self.stats['v_bytes'] * 8) / (1000 * 1000) / delta
            a_kbps = (self.stats['a_bytes'] * 8) / 1000 / delta
            print(f"[STATS] Video: {v_mbps:6.2f} Mbps ({self.stats['v_pkts']} pps) | Audio: {a_kbps:6.2f} kbps ({self.stats['a_pkts']} pps)")
            self.stats = {'v_bytes': 0, 'a_bytes': 0, 'v_pkts': 0, 'a_pkts': 0, 'last_time': now}

    def run(self):
        if not self.setup_system(): return
        self.running = True
        
        # Discovery
        info = ServiceInfo("_procam._udp.local.", "AirLink._procam._udp.local.", addresses=[socket.inet_aton(self.ip)], port=CONF['video_port'], properties={'version': '2.0'})
        self.zc = Zeroconf(); self.zc.register_service(info)

        # Sockets
        for k, port, buf in [('video', CONF['video_port'], 4*1024*1024), ('audio', CONF['audio_port'], 512*1024)]:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, buf)
            s.bind(("0.0.0.0", port))
            s.setblocking(0)
            self.sockets[k] = s

        self.start_pipeline('video'); self.start_pipeline('audio')
        print(f"\n[READY] Server listening on {self.ip}:{CONF['video_port']}\n")

        while self.running:
            self.log_stats()
            # Ensure audio is alive
            if self.procs['audio'] and self.procs['audio'].poll() is not None: self.start_pipeline('audio')

            ready, _, _ = select.select(list(self.sockets.values()), [], [], 0.1)
            for sock in ready:
                try:
                    data, addr = sock.recvfrom(65535)
                    
                    # Control Signals
                    if len(data) < 16:
                        if data == b'PING': sock.sendto(b'PONG', addr); continue
                        if data == b'HEARTBEAT': sock.sendto(b'ACK', addr); continue
                        if data in [b'PAUSE', b'BYE']: self.paint_black(); continue

                    # Data Routing
                    if sock == self.sockets['video']:
                        self.stats['v_bytes'] += len(data); self.stats['v_pkts'] += 1
                        if not self.procs['video'] or self.procs['video'].poll() is not None: self.start_pipeline('video')
                        if self.procs['video']: self.procs['video'].stdin.write(data); self.procs['video'].stdin.flush()
                    
                    elif sock == self.sockets['audio']:
                        self.stats['a_bytes'] += len(data); self.stats['a_pkts'] += 1
                        if self.procs['ffmpeg_audio']: self.procs['ffmpeg_audio'].stdin.write(data); self.procs['ffmpeg_audio'].stdin.flush()
                except Exception: pass

    def cleanup(self):
        self.running = False
        self.paint_black()
        if self.zc: self.zc.close()
        for mod in [CONF['audio_sink'], CONF['audio_src']]: subprocess.run(f"pactl unload-module {mod}", shell=True, stderr=subprocess.DEVNULL)
        print("\n[INFO] Cleanup complete. Exiting.")

if __name__ == "__main__":
    server = AirLinkServer()
    signal.signal(signal.SIGINT, lambda s,f: server.cleanup() or sys.exit(0))
    server.run()