#!/usr/bin/env python3
"""Avvia/ferma ffmpeg sul nodo stream (loopback MediaMTX → YouTube).
Il control plane chiama questo agent all'avvio della diretta; ffmpeg resta sul CPX.
"""
from __future__ import annotations

import json
import os
import signal
import subprocess
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse

SECRET = os.environ.get("STREAM_NODE_AGENT_SECRET", "")
LISTEN = os.environ.get("STREAM_NODE_AGENT_LISTEN", "0.0.0.0:9100")
HLS_BASE = os.environ.get("STREAM_NODE_LOCAL_HLS_URL", "http://127.0.0.1:8888").rstrip("/")
LOG_DIR = os.environ.get("STREAM_NODE_RELAY_LOG_DIR", "/var/log")

_lock = threading.Lock()
# session_id -> {"pid": int, "path": str, "proc": Popen}
_relays: dict[str, dict] = {}


def _authorized(handler: BaseHTTPRequestHandler) -> bool:
    if not SECRET:
        return True
    hdr = handler.headers.get("Authorization", "")
    return hdr == f"Bearer {SECRET}"


def _ffmpeg_cmd(path: str, rtmps: str) -> list[str]:
    return [
        "ffmpeg",
        "-nostdin",
        "-hide_banner",
        "-loglevel",
        "warning",
        "-fflags",
        "+genpts+discardcorrupt",
        "-reconnect",
        "1",
        "-reconnect_streamed",
        "1",
        "-reconnect_on_network_error",
        "1",
        "-reconnect_delay_max",
        "2",
        "-rw_timeout",
        "15000000",
        "-live_start_index",
        "-1",
        "-i",
        f"{HLS_BASE}/{path}/index.m3u8",
        "-c:v",
        "copy",
        "-c:a",
        "copy",
        "-bsf:a",
        "aac_adtstoasc",
        "-f",
        "flv",
        rtmps,
    ]


def _alive(pid: int) -> bool:
    try:
        os.kill(pid, 0)
        return True
    except OSError:
        return False


def _stop_session(session_id: str) -> None:
    info = _relays.pop(session_id, None)
    if not info:
        return
    proc = info.get("proc")
    pid = int(info.get("pid") or 0)
    if proc and proc.poll() is None:
        try:
            os.killpg(proc.pid, signal.SIGTERM)
        except OSError:
            pass
        try:
            proc.wait(timeout=2)
        except Exception:
            try:
                os.killpg(proc.pid, signal.SIGKILL)
            except OSError:
                pass
    elif pid:
        try:
            os.kill(pid, signal.SIGTERM)
        except OSError:
            pass


def _start_session(session_id: str, path: str, rtmps: str) -> int:
    existing = _relays.get(session_id)
    if existing:
        proc = existing.get("proc")
        pid = int(existing.get("pid") or 0)
        if proc is not None and proc.poll() is None and existing.get("path") == path:
            return pid
    _stop_session(session_id)
    os.makedirs(LOG_DIR, exist_ok=True)
    safe = path.replace("/", "_")
    log_path = os.path.join(LOG_DIR, f"youtube-relay-{safe}.log")
    log = open(log_path, "ab", buffering=0)
    proc = subprocess.Popen(
        _ffmpeg_cmd(path, rtmps),
        stdout=log,
        stderr=log,
        start_new_session=True,
    )
    _relays[session_id] = {"pid": proc.pid, "path": path, "proc": proc}
    return proc.pid


class Handler(BaseHTTPRequestHandler):
    def log_message(self, fmt: str, *args) -> None:
        print(f"[relay-agent] {self.address_string()} {fmt % args}")

    def _json(self, code: int, payload: dict) -> None:
        body = json.dumps(payload).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self) -> None:
        parsed = urlparse(self.path)
        if parsed.path in ("/health", "/"):
            self._json(200, {"ok": True, "relays": len(_relays)})
            return
        if not _authorized(self):
            self._json(401, {"error": "unauthorized"})
            return
        if parsed.path.startswith("/relays/"):
            session_id = parsed.path.split("/relays/", 1)[1].strip("/")
            with _lock:
                info = _relays.get(session_id)
                running = bool(info and _alive(int(info["pid"])))
                if info and not running:
                    _relays.pop(session_id, None)
            if not running:
                self._json(404, {"running": False, "session_id": session_id})
                return
            self._json(200, {"running": True, "session_id": session_id, "pid": info["pid"], "path": info["path"]})
            return
        self._json(404, {"error": "not found"})

    def do_POST(self) -> None:
        if not _authorized(self):
            self._json(401, {"error": "unauthorized"})
            return
        if urlparse(self.path).path != "/relays":
            self._json(404, {"error": "not found"})
            return
        length = int(self.headers.get("Content-Length") or 0)
        try:
            data = json.loads(self.rfile.read(length) or b"{}")
        except json.JSONDecodeError:
            self._json(400, {"error": "invalid json"})
            return
        session_id = str(data.get("session_id") or "").strip()
        path = str(data.get("path") or "").strip()
        rtmps = str(data.get("rtmps") or "").strip()
        if not session_id or not path or not rtmps.startswith("rtmps://"):
            self._json(400, {"error": "session_id, path, rtmps required"})
            return
        with _lock:
            pid = _start_session(session_id, path, rtmps)
        self._json(201, {"pid": pid, "session_id": session_id, "path": path})

    def do_DELETE(self) -> None:
        if not _authorized(self):
            self._json(401, {"error": "unauthorized"})
            return
        parsed = urlparse(self.path)
        if not parsed.path.startswith("/relays/"):
            self._json(404, {"error": "not found"})
            return
        session_id = parsed.path.split("/relays/", 1)[1].strip("/")
        with _lock:
            _stop_session(session_id)
        self._json(200, {"stopped": True, "session_id": session_id})


def main() -> None:
    host, port_s = LISTEN.rsplit(":", 1)
    httpd = ThreadingHTTPServer((host, int(port_s)), Handler)
    print(f"[relay-agent] listen {LISTEN} hls={HLS_BASE}")
    httpd.serve_forever()


if __name__ == "__main__":
    main()