local_worker/worker.py

import os
import sys

# Setup dynamic CUDA paths before importing Whisper
nvidia_subdirs = [
    ("nvidia", "cublas", "lib"),
    ("nvidia", "cudnn", "lib"),
    ("nvidia", "cuda_nvrtc", "lib")
]
added_paths = []
for p in sys.path:
    if "site-packages" in p:
        for subdir in nvidia_subdirs:
            full_path = os.path.join(p, *subdir)
            if os.path.exists(full_path) and full_path not in added_paths:
                added_paths.append(full_path)
if added_paths:
    existing = os.environ.get("LD_LIBRARY_PATH", "")
    if existing:
        os.environ["LD_LIBRARY_PATH"] = ":".join(added_paths) + ":" + existing
    else:
        os.environ["LD_LIBRARY_PATH"] = ":".join(added_paths)

import time
import socket
import select
import threading
import queue
import subprocess
import random
import requests
import numpy as np
from dotenv import load_dotenv
from faster_whisper import WhisperModel
from vod_backfiller import download_vod_chat, transcribe_vod_audio

# Load configurations
load_dotenv()

API_URL = os.getenv("API_URL", "http://localhost:3000")
API_KEY = os.getenv("API_KEY", "")
TWITCH_CHANNEL = os.getenv("TWITCH_CHANNEL", "winx_prinx").lower()
CAPTURE_METHOD = os.getenv("CAPTURE_METHOD", "stream").lower()
WHISPER_MODEL_SIZE = os.getenv("WHISPER_MODEL", "base")
WHISPER_DEVICE = os.getenv("WHISPER_DEVICE", "cpu")
WHISPER_COMPUTE_TYPE = os.getenv("WHISPER_COMPUTE_TYPE", "int8")

# Verify essential secrets
if not API_KEY:
    print("[Error] API_KEY is missing in .env! Local worker cannot push data.")
    sys.exit(1)

# Thread-safe message queue for Twitch chat messages
chat_queue = queue.Queue()
# Flag to signal thread termination
stop_flag = threading.Event()

# =========================================================================
# TWITCH CHAT LOGGER (IRC CLIENT)
# =========================================================================

def parse_irc_tags(tags_str):
    """Parse IRC v3 tags into a dictionary"""
    tags = {}
    if not tags_str:
        return tags
    
    parts = tags_str.split(";")
    for part in parts:
        if "=" in part:
            k, v = part.split("=", 1)
            tags[k] = v
    return tags

def twitch_chat_listener():
    """Background thread to connect to Twitch IRC and read chat messages"""
    server = "irc.chat.twitch.tv"
    port = 6667
    
    # Generate a random anonymous nickname
    anon_nick = f"justinfan{random.randint(10000, 99999)}"
    
    print(f"[Twitch IRC] Connecting anonymously as {anon_nick}...")
    
    while not stop_flag.is_set():
        try:
            irc_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            irc_sock.settimeout(10.0)
            irc_sock.connect((server, port))
            
            # Authenticate anonymously
            irc_sock.send(f"PASS oauth:anonymous\r\n".encode("utf-8"))
            irc_sock.send(f"NICK {anon_nick}\r\n".encode("utf-8"))
            
            # Request tags and commands capability to see badges/sub status
            irc_sock.send("CAP REQ :twitch.tv/tags twitch.tv/commands twitch.tv/membership\r\n".encode("utf-8"))
            
            # Join target channel
            irc_sock.send(f"JOIN #{TWITCH_CHANNEL}\r\n".encode("utf-8"))
            print(f"[Twitch IRC] Joined channel #{TWITCH_CHANNEL}")
            
            buffer = ""
            irc_sock.setblocking(False)
            
            while not stop_flag.is_set():
                # Use select for non-blocking read with timeout to allow exit checks
                ready = select.select([irc_sock], [], [], 1.0)
                if not ready[0]:
                    continue
                
                try:
                    data = irc_sock.recv(4096).decode("utf-8", errors="ignore")
                except socket.timeout:
                    continue
                
                if not data:
                    print("[Twitch IRC] Connection closed by remote host.")
                    break
                
                buffer += data
                while "\r\n" in buffer:
                    line, buffer = buffer.split("\r\n", 1)
                    
                    # 1. Handle Ping-Pong
                    if line.startswith("PING"):
                        irc_sock.send("PONG :tmi.twitch.tv\r\n".encode("utf-8"))
                        continue
                    
                    # 2. Parse PRIVMSG (chat message)
                    # Twitch format: @tags :user!user@user.tmi.twitch.tv PRIVMSG #channel :message
                    if "PRIVMSG" in line:
                        tags = {}
                        tags_str = ""
                        
                        # Extract tags if they exist
                        if line.startswith("@"):
                            tags_str, remainder = line[1:].split(" ", 1)
                            tags = parse_irc_tags(tags_str)
                            line = remainder
                        
                        parts = line.split(" PRIVMSG ", 1)
                        if len(parts) < 2:
                            continue
                            
                        prefix, msg_parts = parts
                        user = prefix.split("!", 1)[0].replace(":", "")
                        
                        channel_part, message_text = msg_parts.split(" :", 1)
                        
                        # Extract user metadata
                        msg_id = tags.get("id", f"local-{time.time_ns()}")
                        display_name = tags.get("display-name", user)
                        badges = tags.get("badges", "")
                        
                        is_streamer = (user.lower() == TWITCH_CHANNEL)
                        is_mod = "moderator" in badges or "broadcaster" in badges
                        is_sub = "subscriber" in badges or "founder" in badges
                        
                        parsed_msg = {
                            "id": msg_id,
                            "username": user,
                            "displayName": display_name,
                            "message": message_text,
                            "timestamp": new_iso_timestamp(),
                            "isStreamer": is_streamer,
                            "isMod": is_mod,
                            "isSub": is_sub
                        }
                        
                        # Add to queue
                        chat_queue.put(parsed_msg)
                        
        except Exception as e:
            print(f"[Twitch IRC] Connection error: {e}. Retrying in 10 seconds...")
            time.sleep(10)
        finally:
            try:
                irc_sock.close()
            except:
                pass

def new_iso_timestamp():
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())

def chat_sender():
    """Periodically sends accumulated chat messages to backend API"""
    print("[Chat Sender] Started background queue sender.")
    headers = {"x-api-key": API_KEY, "Content-Type": "application/json"}
    
    while not stop_flag.is_set():
        messages = []
        # Drain the queue
        while not chat_queue.empty():
            try:
                messages.append(chat_queue.get_nowait())
            except queue.Empty:
                break
        
        if messages:
            try:
                url = f"{API_URL}/api/log/messages"
                response = requests.post(url, json={"messages": messages}, headers=headers, timeout=5)
                if response.status_code == 200:
                    pass # Success
                else:
                    print(f"[Chat Sender] Failed to sync messages. API returned status {response.status_code}")
            except Exception as e:
                print(f"[Chat Sender] Network error sending messages: {e}")
                
        # Wait 5 seconds before next sync
        time.sleep(5)

# =========================================================================
# AUDIO TRANSCRIBER (WHISPER)
# =========================================================================

def init_whisper_model():
    """Load Faster-Whisper model into memory"""
    print(f"[Whisper] Loading model '{WHISPER_MODEL_SIZE}' on {WHISPER_DEVICE} ({WHISPER_COMPUTE_TYPE})...")
    # This might take a few minutes on first run as the model downloads
    model = WhisperModel(
        WHISPER_MODEL_SIZE, 
        device=WHISPER_DEVICE, 
        compute_type=WHISPER_COMPUTE_TYPE
    )
    print("[Whisper] Model loaded successfully.")
    return model

def transcribe_audio_segment(model, pcm_bytes):
    """Run Whisper transcription on PCM 16kHz 16-bit mono audio bytes"""
    if len(pcm_bytes) == 0:
        return []
    
    # Convert raw 16-bit PCM bytes to float32 numpy array normalized to [-1.0, 1.0]
    audio_np = np.frombuffer(pcm_bytes, dtype=np.int16).astype(np.float32) / 32768.0
    
    # Force language="ru" for better Russian accuracy and faster execution
    segments, info = model.transcribe(
        audio_np, 
        beam_size=5, 
        language="ru",
        vad_filter=True, # Voice Activity Detection filters out silence
        vad_parameters=dict(min_silence_duration_ms=500)
    )
    
    words_list = []
    for segment in segments:
        text = segment.text.strip()
        if text:
            print(f"[Whisper Speech] Transcribed: \"{text}\"")
            # Split segment into clean words
            words = text.split()
            words_list.extend(words)
            
    return words_list

def send_voice_words(words):
    """Post transcribed words to backend server"""
    if not words:
        return
    
    headers = {"x-api-key": API_KEY, "Content-Type": "application/json"}
    url = f"{API_URL}/api/log/voice"
    
    try:
        response = requests.post(url, json={"words": words, "timestamp": new_iso_timestamp()}, headers=headers, timeout=5)
        if response.status_code == 200:
            print(f"[Sync] Sent {len(words)} spoken words to cloud.")
        else:
            print(f"[Sync] Failed to send words. Status code: {response.status_code}")
    except Exception as e:
        print(f"[Sync] Error sending voice words to cloud: {e}")

def notify_stream_end():
    """Notify backend that the stream has ended"""
    headers = {"x-api-key": API_KEY}
    url = f"{API_URL}/api/log/stream-end"
    try:
        requests.post(url, headers=headers, timeout=5)
        print("[Sync] Sent stream-end signal.")
    except Exception as e:
        print(f"[Sync] Failed to send stream-end: {e}")

# =========================================================================
# CAPTURE METHODS
# =========================================================================

def check_stream_live():
    """Check if the Twitch channel is currently streaming using streamlink"""
    try:
        # Runs streamlink check without downloading stream
        result = subprocess.run(
            ["streamlink", f"twitch.tv/{TWITCH_CHANNEL}", "--stream-url"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        # If streamlink returns a stream URL, it's live!
        return result.returncode == 0 and result.stdout.strip().startswith("http")
    except FileNotFoundError:
        print("[Error] 'streamlink' utility is not installed or not in system PATH! Please install it.")
        time.sleep(10)
        return False

def run_stream_capture(model):
    """Capture Twitch stream audio using streamlink + ffmpeg and transcribe"""
    print(f"[Capture] Monitoring Twitch channel '{TWITCH_CHANNEL}'...")
    
    while not stop_flag.is_set():
        if not check_stream_live():
            print(f"[Capture] Channel '{TWITCH_CHANNEL}' is offline. Checking for pending VOD backfill tasks...")
            try:
                headers_get = {"x-api-key": API_KEY}
                res = requests.get(f"{API_URL}/api/streams/pending-backfill", headers=headers_get, timeout=10)
                if res.status_code == 200:
                    data = res.json()
                    pending_streams = data.get("streams", [])
                    if pending_streams:
                        print(f"[Backfill] Found {len(pending_streams)} pending VOD backfill(s). Starting automated processing...")
                        
                        target = pending_streams[0]
                        stream_id = target.get("id")
                        vod_id = target.get("twitch_vod_id")
                        title = target.get("title", "Unknown Archive")
                        start_time = target.get("start_time")
                        
                        min_msg_time = target.get("min_message_time")
                        max_msg_time = target.get("max_message_time")
                        min_voice_time = target.get("min_voice_time")
                        max_voice_time = target.get("max_voice_time")

                        # Parse ISO strings to epoch timestamps
                        def parse_time(t_str):
                            if not t_str:
                                return None
                            try:
                                clean_t = t_str.split(".")[0].replace("Z", "").replace("+00:00", "")
                                return time.mktime(time.strptime(clean_t, "%Y-%m-%dT%H:%M:%S"))
                            except Exception as parse_err:
                                print(f"[Backfill] Error parsing time '{t_str}': {parse_err}")
                                return None

                        start_epoch = parse_time(start_time)
                        
                        # -- CHAT BACKFILL TASKS --
                        chat_tasks = []
                        if not min_msg_time:
                            # 0 messages, download everything
                            chat_tasks.append((0, None))
                        else:
                            min_msg_epoch = parse_time(min_msg_time)
                            max_msg_epoch = parse_time(max_msg_time)
                            
                            min_msg_offset = int(min_msg_epoch - start_epoch) if (min_msg_epoch and start_epoch) else 0
                            max_msg_offset = int(max_msg_epoch - start_epoch) if (max_msg_epoch and start_epoch) else 0
                            
                            # If first message is > 60s from start, download the beginning gap
                            if min_msg_offset > 60:
                                chat_tasks.append((0, min_msg_offset))
                            # Download from last message to end
                            chat_tasks.append((max_msg_offset, None))

                        # -- VOICE BACKFILL TASKS --
                        voice_tasks = []
                        if not min_voice_time:
                            # 0 words, transcribe everything
                            voice_tasks.append((0, None))
                        else:
                            min_voice_epoch = parse_time(min_voice_time)
                            max_voice_epoch = parse_time(max_voice_time)
                            
                            min_voice_offset = int(min_voice_epoch - start_epoch) if (min_voice_epoch and start_epoch) else 0
                            max_voice_offset = int(max_voice_epoch - start_epoch) if (max_voice_epoch and start_epoch) else 0
                            
                            # If first voice word is > 60s from start, transcribe the beginning gap
                            if min_voice_offset > 60:
                                voice_tasks.append((0, min_voice_offset))
                            # Transcribe from last word (with 10s safety overlap) to end
                            voice_tasks.append((max(0, max_voice_offset - 10), None))

                        print(f"\n=========================================================")
                        print(f"[Backfill] Processing: {title}")
                        print(f"[Backfill] VOD ID:       {vod_id}")
                        print(f"[Backfill] Stream ID:   {stream_id}")
                        print(f"[Backfill] Chat Tasks:  {chat_tasks}")
                        print(f"[Backfill] Voice Tasks: {voice_tasks}")
                        print(f"=========================================================")
                        
                        # 1. Download and upload chat segments
                        for start_off, end_off in chat_tasks:
                            print(f"[Backfill] Downloading chat from {start_off}s to {'end' if end_off is None else str(end_off) + 's'}...")
                            chat_comments = download_vod_chat(vod_id, start_offset=start_off, end_offset=end_off)
                            if chat_comments:
                                print(f"[Backfill] Uploading {len(chat_comments)} chat comments...")
                                batch_size = 100
                                headers_post = {"x-api-key": API_KEY, "Content-Type": "application/json"}
                                for i in range(0, len(chat_comments), batch_size):
                                    batch = chat_comments[i:i+batch_size]
                                    try:
                                        requests.post(f"{API_URL}/api/log/messages", json={"messages": batch}, headers=headers_post, timeout=10)
                                    except Exception as e:
                                        print(f"[Backfill] Chat batch upload error: {e}")
                        
                        # 2. Download and transcribe audio segments
                        print(f"[Backfill] Starting Whisper audio transcription...")
                        for start_off, end_off in voice_tasks:
                            print(f"[Backfill] Transcribing audio from {start_off}s to {'end' if end_off is None else str(end_off) + 's'}...")
                            try:
                                start_time_iso = start_time.replace("+00:00", "").replace("Z", "") + "Z"
                                transcribe_vod_audio(vod_id, start_time_iso, model, start_offset=start_off, end_offset=end_off)
                            except Exception as e:
                                print(f"[Backfill] Transcription error: {e}")
                        
                        # 3. Mark completed
                        print(f"[Backfill] Marking stream {stream_id} as backfilled...")
                        try:
                            res_mark = requests.post(f"{API_URL}/api/streams/mark-backfilled", json={"streamId": stream_id}, headers=headers_get, timeout=10)
                            if res_mark.status_code == 200:
                                print(f"[Backfill] Successfully marked stream {stream_id} as backfilled!")
                            else:
                                print(f"[Backfill] Failed to mark stream: {res_mark.status_code} - {res_mark.text}")
                        except Exception as e:
                            print(f"[Backfill] Network error marking stream: {e}")
                        
                        # Loop again immediately to process next VOD or check if live
                        continue
                    else:
                        print("[Capture] No pending VOD backfills found.")
                else:
                    print(f"[Backfill] Failed to fetch pending backfills: {res.status_code}")
            except Exception as err:
                print(f"[Backfill] Error checking pending backfills: {err}")
                
            print(f"[Capture] Retrying stream check in 60 seconds...")
            time.sleep(60)
            continue
        
        print(f"[Capture] Stream is LIVE! Starting audio pipeline...")
        
        # Start streamlink piping audio to ffmpeg, converting it to raw 16kHz mono 16-bit PCM
        streamlink_cmd = ["streamlink", f"twitch.tv/{TWITCH_CHANNEL}", "audio,worst", "-O"]
        ffmpeg_cmd = ["ffmpeg", "-i", "pipe:0", "-ac", "1", "-ar", "16000", "-f", "s16le", "-"]
        
        p_streamlink = None
        p_ffmpeg = None
        
        try:
            p_streamlink = subprocess.Popen(streamlink_cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
            p_ffmpeg = subprocess.Popen(ffmpeg_cmd, stdin=p_streamlink.stdout, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
            
            # Read in 10-second chunks
            # 16000Hz * 16-bit(2 bytes) * 10 seconds = 320,000 bytes
            chunk_size = 16000 * 2 * 10 
            
            while not stop_flag.is_set():
                pcm_data = p_ffmpeg.stdout.read(chunk_size)
                if not pcm_data:
                    print("[Capture] Stream audio stopped or disconnected.")
                    break
                
                # Transcribe in a background thread to prevent blocking audio read
                threading.Thread(
                    target=lambda data=pcm_data: send_voice_words(transcribe_audio_segment(model, data))
                ).start()
                
        except Exception as err:
            print(f"[Capture] Audio pipeline crashed: {err}")
        finally:
            # Clean up processes
            for p in [p_ffmpeg, p_streamlink]:
                if p:
                    try:
                        p.terminate()
                        p.wait(timeout=2)
                    except:
                        try: p.kill() 
                        except: pass
            
            notify_stream_end()
            print("[Capture] Pipeline stopped. Waiting 30 seconds before checking stream status...")
            time.sleep(30)

def run_microphone_capture(model):
    """Capture local microphone/system audio and transcribe"""
    try:
        import sounddevice as sd
    except ImportError:
        print("[Error] 'sounddevice' library is missing! Install it via pip install sounddevice.")
        sys.exit(1)
        
    print("[Capture] Starting microphone capture. Listening...")
    
    sample_rate = 16000
    duration = 10 # Transcribe every 10 seconds
    
    while not stop_flag.is_set():
        try:
            # Record float32 directly from default input device
            recording = sd.rec(
                int(duration * sample_rate), 
                samplerate=sample_rate, 
                channels=1, 
                dtype='float32'
            )
            sd.wait() # Wait until recording is finished
            
            # sounddevice records float32 directly, no need for raw conversion
            audio_np = recording.flatten()
            
            # Transcribe and send words
            def process_mic(audio=audio_np):
                segments, info = model.transcribe(audio, beam_size=5, language="ru")
                words_list = []
                for segment in segments:
                    text = segment.text.strip()
                    if text:
                        print(f"[Whisper Mic] Transcribed: \"{text}\"")
                        words_list.extend(text.split())
                send_voice_words(words_list)
                
            threading.Thread(target=process_mic).start()
            
        except Exception as e:
            print(f"[Capture] Microphone error: {e}")
            time.sleep(5)

# =========================================================================
# MAIN EXECUTION
# =========================================================================

if __name__ == "__main__":
    print("=========================================================")
    print("      Twitch Analytics Local Logger & Voice Sync         ")
    print("=========================================================")
    print(f"Target Channel:  {TWITCH_CHANNEL}")
    print(f"Capture Method:  {CAPTURE_METHOD}")
    print(f"Whisper Model:   {WHISPER_MODEL_SIZE}")
    print(f"API Backend URL: {API_URL}")
    print("=========================================================")
    
    # 1. Initialize Whisper
    whisper_model = init_whisper_model()
    
    # 2. Start Twitch Chat background threads
    t_chat = threading.Thread(target=twitch_chat_listener, daemon=True)
    t_sender = threading.Thread(target=chat_sender, daemon=True)
    
    t_chat.start()
    t_sender.start()
    
    # 3. Start Audio Capture (Blocks main thread)
    try:
        if CAPTURE_METHOD == "stream":
            run_stream_capture(whisper_model)
        elif CAPTURE_METHOD == "microphone":
            run_microphone_capture(whisper_model)
        else:
            print(f"[Error] Unknown CAPTURE_METHOD '{CAPTURE_METHOD}'. Choose 'stream' or 'microphone'.")
    except KeyboardInterrupt:
        print("\n[Shutting Down] Gracefully stopping threads...")
        stop_flag.set()
        time.sleep(1)
        print("[Shutting Down] Done.")