PageBot/pagebot.py

#!/usr/bin/python3
import asyncio
import discord
import os
import subprocess
import time
import secrets_file
import threading
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
import shutil
import speech_recognition as sr

def wait_for_file_ready(filepath, stable_seconds=1, timeout=60):
    elapsed = 0
    last_size = -1
    stable_time = 0
    while elapsed < timeout:
        try:
            size = os.path.getsize(filepath)
        except OSError:
            return False
        if size == last_size:
            stable_time += 1
            if stable_time >= stable_seconds:
                return True
        else:
            stable_time = 0
        last_size = size
        time.sleep(1)
        elapsed += 1
    return False

last_file_time = time.time()
_processing_lock = threading.Lock()
_processing_files = set()

class MyHandler(FileSystemEventHandler):
    def on_created(self, event):
        global last_file_time
        if event.is_directory:
            return
        filepath = event.src_path
        filename, file_extension = os.path.splitext(filepath)
        if filename.endswith('_stt'):
            return
        last_file_time = time.time()
        if file_extension.lower() == '.amr':
            with _processing_lock:
                if filepath in _processing_files:
                    return
                _processing_files.add(filepath)
            try:
                wait_for_file_ready(filepath)
                os.remove(filepath)
                print("Removing AMR.")
            finally:
                with _processing_lock:
                    _processing_files.discard(filepath)
        if file_extension.lower() == '.mp3':
            with _processing_lock:
                if filepath in _processing_files:
                    return
                _processing_files.add(filepath)
            try:
                print("New MP3!")
                wait_for_file_ready(filepath)
                if secrets_file.speech_to_text:
                    mp3_copy = f"{filename}_stt.mp3"
                    shutil.copy2(filepath, mp3_copy)
                print("Converting to MP4")
                mp4_file = convert_to_mp4(filepath)
                print("Sending to Discord")
                asyncio.run_coroutine_threadsafe(upload_to_discord(mp4_file), client.loop)
                if secrets_file.speech_to_text:
                    print("Transcribing in background")
                    mp4_basename = os.path.basename(mp4_file) if mp4_file else os.path.basename(filepath)
                    def do_transcribe():
                        text = convert_to_text(mp3_copy, filename)
                        try:
                            os.remove(mp3_copy)
                        except OSError:
                            pass
                        if text:
                            asyncio.run_coroutine_threadsafe(
                                send_transcript(mp4_basename, text), client.loop
                            )
                    threading.Thread(target=do_transcribe, daemon=True).start()
            finally:
                with _processing_lock:
                    _processing_files.discard(filepath)

def convert_to_mp4(mp3_file):
    try:
        mp4_file = os.path.splitext(mp3_file)[0] + '.mp4'
        subprocess.run([
            'ffmpeg', '-y', '-loop', '1',
            '-i', secrets_file.image_path,
            '-i', mp3_file,
            '-c:a', 'aac', '-b:a', '192k',
            '-c:v', 'libx264', '-preset', 'ultrafast', '-tune', 'stillimage', '-pix_fmt', 'yuv420p',
            '-shortest', mp4_file
        ])
        os.remove(mp3_file)
        return mp4_file
    except Exception as e:
        print(f"Error during conversion: {e}")
        return None

def convert_to_text(mp3_path, mp3_name):
    print(mp3_path)
    try:
        wav_file = f"{mp3_name}.wav"
        subprocess.run(['ffmpeg', '-y', '-i', mp3_path, wav_file])
        r = sr.Recognizer()
        with sr.AudioFile(wav_file) as source:
            data = r.record(source)
        text = r.recognize_google(data)
        os.remove(wav_file)
        return text
    except Exception as e:
        print(f"Error during conversion: {e}")
        return ""

def parse_dispatch_filename(filename):
    name = os.path.splitext(filename)[0]
    parts = name.split('-', 1)
    station_number = parts[0].strip()
    remainder = parts[1].strip() if len(parts) > 1 else ""
    segments = remainder.split('_')
    station_name = segments[0] if segments else ""
    timestamp = ""
    if len(segments) >= 6:
        year, month, day = segments[1], segments[2], segments[3]
        hour, minute, second = segments[4], segments[5], segments[6] if len(segments) > 6 else "00"
        timestamp = f"{month}/{day}/{year} {hour}:{minute}:{second}"
    return station_number, station_name, timestamp

async def upload_to_discord(mp4_file):
    if mp4_file is None:
        print("Conversion failed. Skipping upload.")
        return
    channel = client.get_channel(secrets_file.channel_id)
    if channel:
        filename = os.path.basename(mp4_file)
        station_number, station_name, timestamp = parse_dispatch_filename(filename)

        embed = discord.Embed(
            title=f"Station {station_number} - {station_name}",
            description=timestamp,
            color=0xFF0000
        )

        role_name = station_number
        role = discord.utils.get(channel.guild.roles, name=role_name)
        mention = f"{role.mention}" if role else ""

        with open(mp4_file, 'rb') as f:
            await channel.send(
                content=f"{mention} {secrets_file.notify_text}".strip() if mention else None,
                embed=embed,
                file=discord.File(f)
            )

        if not role:
            print(f"Role '{role_name}' not found.")
        if secrets_file.delete_after_upload:
            os.remove(mp4_file)
    else:
        print(f"Could not find channel with ID {secrets_file.channel_id}")

async def send_transcript(filename, text):
    channel = client.get_channel(secrets_file.channel_id)
    if channel:
        station_number, station_name, timestamp = parse_dispatch_filename(filename)
        embed = discord.Embed(
            title=f"Transcript - Station {station_number} - {station_name}",
            description=text,
            color=0x3498db
        )
        await channel.send(embed=embed)

TTD_TIMEOUT = getattr(secrets_file, 'ttd_timeout', 900)

def kill_process_tree(pid):
    try:
        subprocess.run(['taskkill', '/F', '/T', '/PID', str(pid)],
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    except Exception as e:
        print(f"Error killing process tree: {e}")

def launch_and_watch(program_path):
    global last_file_time
    program_directory = os.path.dirname(program_path)
    while True:
        last_file_time = time.time()
        process = subprocess.Popen(program_path, cwd=program_directory)
        while True:
            try:
                process.wait(timeout=30)
                break
            except subprocess.TimeoutExpired:
                if time.time() - last_file_time > TTD_TIMEOUT:
                    print("TTD appears to have locked up (no new files in 15 minutes). Restarting...")
                    kill_process_tree(process.pid)
                    break
        if process.returncode is None or process.returncode != 0:
            print("TTD has crashed. Relaunching...")
            kill_process_tree(process.pid)
        else:
            print("TTD has exited normally.")
            break
        time.sleep(10)

if __name__ == "__main__":
    observer = Observer()

    intents = discord.Intents.default()
    intents.message_content = True
    client = discord.Client(intents=intents)

    @client.event
    async def on_ready():
        print(f'We have logged in as {client.user}')

        event_handler = MyHandler()
        observer.schedule(event_handler, secrets_file.watch_folder, recursive=False)
        observer.start()

        if secrets_file.ttd_path != "":
            watchdog_thread = threading.Thread(target=launch_and_watch, args=(secrets_file.ttd_path,), daemon=True)
            watchdog_thread.start()

    try:
        client.run(secrets_file.key)
    finally:
        observer.stop()
        observer.join()