Blame for ingest.py

"""Media ingest helpers. Video path is RTMPS; OAuth is optional metadata only.

Never log the stream key or the full RTMP URL.
"""
import os, threading, time
from datetime import datetime
from zoneinfo import ZoneInfo

APP_VERSION = 'relay-v2.45-probed-canvas'


def env_rtmp_destination():
    """Build rtmps://…/key from secrets. Empty string if not configured."""
    url = (os.getenv('YOUTUBE_RTMP_URL') or os.getenv('YOUTUBE_RTMPS_URL') or '').strip()
    key = (os.getenv('YOUTUBE_STREAM_KEY') or os.getenv('YOUTUBE_RTMP_KEY') or '').strip()
    if url.startswith(('rtmp://', 'rtmps://')) and key:
        if url.rstrip('/').endswith('/' + key):
            return url
        return url.rstrip('/') + '/' + key
    if url.startswith(('rtmp://', 'rtmps://')):
        return url
    if key:
        return 'rtmps://a.rtmps.youtube.com:443/live2/' + key
    return ''


def rtmp_configured():
    return bool(env_rtmp_destination())


def max_concurrent_relays():
    """Cap on simultaneous FFmpeg publishers. Not a quality guarantee.

    Default 2: cpu-basic has carried 1 comfortably and 2 at high quota; 3+ is
    not a pass. Operators may raise MAX_CONCURRENT_RELAYS; health reports the
    value actually in force. Ceiling 8 is only a fork-bomb guard.
    """
    raw = os.getenv('MAX_CONCURRENT_RELAYS')
    if raw is None or str(raw).strip() == '':
        raw = os.getenv('YOUTUBE_MAX_EVENTS') or '0'
    try:
        n = int(str(raw).strip())
    except (TypeError, ValueError):
        n = 0
    # 0 = every LIVE account may copy (one FFmpeg per user). No hidden ceiling.
    return n if n > 0 else 0


def light_max():
    """Operator opt-in ceiling for legacy x264 fallback; zero disables video re-encoding.

    The default is deliberately zero: AVC and Enhanced-RTMP HEVC are copied, including decoder
    configuration replacement. Historical CBR/x264 modes remain available only by explicitly
    raising this limit.
    """
    try:
        n = int(str(os.getenv('YOUTUBE_LIGHT_MAX', '0')).strip() or '0')
    except (TypeError, ValueError):
        n = 0
    return max(0, min(n, 8))


def format_session_title(broadcaster, when=None):
    """Baghdad clock of this TikTok LIVE session — not process start time."""
    when = when or datetime.now(ZoneInfo('Asia/Baghdad'))
    values = {
        'broadcaster': broadcaster or '',
        'name': broadcaster or '',
        'nickname': broadcaster or '',
        'date': when.strftime('%Y-%m-%d'),
        'time': when.strftime('%H:%M'),
        'title': broadcaster or '',
        'username': '',
        'hashtags': '',
        'room_id': '',
        'live_title': '',
    }
    tpl = (os.getenv('YOUTUBE_TITLE_TEMPLATE') or '{broadcaster} | {date} {time}').strip()
    try:
        title = tpl.format(**values).strip()[:100]
    except (KeyError, ValueError, IndexError, AttributeError):
        title = '%s | %s %s' % (values['broadcaster'], values['date'], values['time'])
    return title, values['date'], values['time']


class RelayHub:
    """At most one FFmpeg relay per broadcaster. Owners: 'oauth' or 'direct'."""

    def __init__(self):
        self.lock = threading.RLock()
        self.slots = {}
        self.restarts = {}
        # Atomic per-broadcaster reservations prevent simultaneous format fallbacks from
        # exceeding YOUTUBE_LIGHT_MAX while their existing FFmpeg children are still exiting.
        self.light_reservations = set()

    @staticmethod
    def _ffmpeg_alive(cur):
        snd = (cur or {}).get('sender')
        try:
            return bool(snd and snd.alive())
        except Exception:
            return False

    @staticmethod
    def _supervisor_alive(cur):
        if not cur:
            return False
        th = cur.get('thread')
        if th and th.is_alive():
            return True
        return RelayHub._ffmpeg_alive(cur)

    def occupy(self, user, owner, sender=None, thread=None, room_id='', title='', dest=''):
        with self.lock:
            cur = self.slots.get(user)
            if self._supervisor_alive(cur):
                return False
            if cur:
                self.restarts[user] = int(self.restarts.get(user) or 0) + 1
            self.slots[user] = {
                'owner': owner, 'sender': sender, 'thread': thread,
                'room_id': room_id, 'title': title, 'started': time.time(),
                'stop': None,
                'dest': dest or ('env' if owner == 'direct' else 'oauth'),
            }
            return True

    def update(self, user, **kw):
        with self.lock:
            if user in self.slots:
                self.slots[user].update(kw)

    def release(self, user, owner=None):
        with self.lock:
            cur = self.slots.get(user)
            if not cur:
                self.light_reservations.discard(user)
                return
            if owner and cur.get('owner') != owner:
                return
            self.slots.pop(user, None)
            self.light_reservations.discard(user)

    def busy(self, user):
        with self.lock:
            return self._supervisor_alive(self.slots.get(user))

    def owner(self, user):
        with self.lock:
            return (self.slots.get(user) or {}).get('owner')

    def snapshot(self, user=None):
        with self.lock:
            if user is not None:
                return dict(self.slots.get(user) or {})
            return {u: dict(v) for u, v in self.slots.items()}

    def active_count(self):
        """FFmpeg processes actually pushing — not idle supervisor threads."""
        with self.lock:
            return sum(1 for v in self.slots.values() if self._ffmpeg_alive(v))

    def _light_count_locked(self, except_user=None):
        """Count active encoders plus reserved slots, without counting either twice."""
        active=set();n=0
        for u,v in self.slots.items():
            if except_user and u==except_user:continue
            if u in self.light_reservations:
                active.add(u);n+=1
                continue
            if not self._ffmpeg_alive(v):continue
            snd=v.get('sender');mode=''
            try:
                mode=str(getattr(snd,'sticky_mode',None) or '')
                if not mode and snd:
                    snap=snd.snapshot() or {};mode=str(snap.get('send_mode') or '')
            except Exception:pass
            if mode in ('matchgop','fixed720','fixed1080'):
                active.add(u);n+=1
        n+=sum(1 for u in self.light_reservations if u!=except_user and u not in active)
        return n

    def light_count(self, except_user=None):
        """Publishers and atomic reservations using matchgop/fixed720/fixed1080."""
        with self.lock:return self._light_count_locked(except_user)

    def light_available(self, except_user=None):
        return self.light_count(except_user) < light_max()

    def reserve_light(self,user):
        """Atomically claim one per-user encoder slot before changing a live publisher."""
        user=str(user or '')
        if not user:return False
        with self.lock:
            if user in self.light_reservations:return True
            limit=light_max()
            if limit<=0 or self._light_count_locked(except_user=user)>=limit:return False
            self.light_reservations.add(user)
            return True

    def release_light(self,user):
        with self.lock:self.light_reservations.discard(str(user or ''))

    def restart_count(self, user):
        with self.lock:
            return int(self.restarts.get(user) or 0)

    def at_capacity(self, except_user=None):
        with self.lock:
            n = sum(1 for u, v in self.slots.items()
                    if u != except_user and self._ffmpeg_alive(v))
            cap = max_concurrent_relays()
            if cap <= 0:
                return False
            return n >= cap

    def env_key_busy(self, except_user=None):
        """A single env stream key can only accept one publisher (any owner)."""
        with self.lock:
            for u, v in self.slots.items():
                if except_user and u == except_user:
                    continue
                if v.get('dest') == 'env' and self._supervisor_alive(v):
                    return True
            return False


hub = RelayHub()