import subprocess
import os
import sys
import time
import signal
import io
import threading
from queue import Queue, Empty
from PIL import Image
from selenium import webdriver
from selenium.webdriver.chrome.options import Options as ChromeOptions
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
from webdriver_manager.core.os_manager import ChromeType

class HeadlessWebStreamer:
    def __init__(self, website_url, music_url, stream_key, fps=20):
        self.website_url = website_url
        self.music_url = music_url
        self.stream_key = stream_key or os.getenv("YT_STREAM_KEY")
        self.youtube_url = "rtmp://a.rtmp.youtube.com/live2"
        self.fps = fps
        self.is_running = False
        self.driver = None
        self.process = None
        self.frame_queue = Queue(maxsize=5)
        
        signal.signal(signal.SIGINT, self._exit_gracefully)
        signal.signal(signal.SIGTERM, self._exit_gracefully)

    def _exit_gracefully(self, sig, frame):
        self.stop_stream()
        sys.exit(0)

    def init_browser(self):
        print("🌐 Инициализация Chromium...")
        try:
            options = ChromeOptions()
            options.add_argument('--headless=new')
            options.add_argument('--no-sandbox')
            options.add_argument('--disable-dev-shm-usage')
            options.add_argument('--disable-gpu')
            # Задаем "длинное" окно, чтобы захватить низ сайта
            options.add_argument('--window-size=1280,1200') 
            options.binary_location = "/usr/bin/chromium"
            
            driver_path = ChromeDriverManager(chrome_type=ChromeType.CHROMIUM).install()
            self.driver = webdriver.Chrome(service=Service(driver_path), options=options)
            self.driver.get(self.website_url)
            
            time.sleep(5)
            return True
        except Exception as e:
            print(f"❌ Ошибка браузера: {e}")
            return False

    def start_stream(self):
        if not self.stream_key:
            print("❌ Ошибка: Ключ трансляции не найден!")
            return

        if not self.init_browser(): 
            return

        # ИСПРАВЛЕННЫЙ ФИЛЬТР FFMPEG:
        # 1. scale=-1:720 — уменьшаем высоту до 720, сохраняя пропорции ширины
        # 2. pad=1280:720 — помещаем результат на холст 1280x720 (добавляем черные поля по бокам)
        video_filter = (
            "scale=-1:720,"
            "pad=1280:720:(1280-iw)/2:0:black,"
            "setsar=1:1,setdar=16/9,format=yuv420p"
        )

        ffmpeg_cmd = [
            'ffmpeg', '-y',
            '-f', 'image2pipe', '-vcodec', 'mjpeg', '-r', str(self.fps), '-i', '-',
            '-re', '-stream_loop', '-1', '-i', self.music_url,
            '-vf', video_filter,
            '-c:v', 'libx264', '-preset', 'ultrafast',
            '-tune', 'zerolatency', '-g', str(self.fps*2), '-b:v', '2500k',
            '-c:a', 'aac', '-b:a', '128k', '-ar', '44100',
            '-map', '0:v:0', '-map', '1:a:0',
            '-f', 'flv', f"{self.youtube_url}/{self.stream_key}"
        ]

        self.process = subprocess.Popen(ffmpeg_cmd, stdin=subprocess.PIPE, stdout=sys.stdout, stderr=sys.stdout)
        self.is_running = True

        def capture_worker():
            while self.is_running:
                try:
                    png = self.driver.get_screenshot_as_png()
                    img = Image.open(io.BytesIO(png)).convert('RGB')
                    if self.frame_queue.full():
                        try: self.frame_queue.get_nowait()
                        except: pass
                    self.frame_queue.put(img)
                except: 
                    break
                time.sleep(1/self.fps)

        threading.Thread(target=capture_worker, daemon=True).start()
        print(f"🚀 Стрим запущен. Формат: Авто-вписывание в 16:9.")

        while self.is_running:
            if self.process.poll() is not None:
                break
            try:
                img = self.frame_queue.get(timeout=10)
                img_byte_arr = io.BytesIO()
                img.save(img_byte_arr, format='JPEG', quality=85)
                self.process.stdin.write(img_byte_arr.getvalue())
                self.process.stdin.flush()
            except (BrokenPipeError, OSError):
                break
            except Empty:
                continue

        self.stop_stream()

    def stop_stream(self):
        self.is_running = False
        if self.driver: self.driver.quit()
        if self.process:
            try:
                self.process.stdin.close()
                self.process.terminate()
            except: pass
        print("⏹️ Трансляция остановлена.")

if __name__ == "__main__":
    MY_MUSIC = "https://dash.rid3usercontent.run.place/cdn/music/funk.mp3"
    WEBSITE = "https://status-lin.web.app/"
    
    target_music = os.getenv("MUSIC_URL", MY_MUSIC)
    stream_key = os.getenv("YT_STREAM_KEY")
    
    streamer = HeadlessWebStreamer(WEBSITE, target_music, stream_key)
    streamer.start_stream()