import yt_dlp
import os
import subprocess
import sys
import json
import time
import asyncio
import websockets

# WebSocket 서버 주소
# WEBSOCKET_SERVER_URL = "wss://progress.it7.kr"
WEBSOCKET_SERVER_URL = "ws://progress-ws:3002"

# 설정
TARGET_FILESIZE_MB = 50
TARGET_FILESIZE_BYTES = TARGET_FILESIZE_MB * 1024 * 1024
DOWNLOAD_DIR = "/var/www/it7/html/downloads/"
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

# 인자 받기
if len(sys.argv) < 5:
    print("❌ Not enough arguments.")
    sys.exit(2)

try:
    url = sys.argv[1]
    speed_options = json.loads(sys.argv[2])
    original_filename = sys.argv[3]
    session_id = sys.argv[4]
    base_filename = original_filename.replace('.mp4', '')
    output_path = os.path.join(DOWNLOAD_DIR, original_filename)
except Exception as e:
    print(f"❌ Argument parsing error: {e}")
    sys.exit(2)

# 공통: WebSocket에 메시지 전송
async def send_ws_message(ws, session_id, progress, message=None):
    try:
        payload = {
            "type": "progress",
            "session_id": session_id,
            "progress": progress
        }
        if message:
            payload["message"] = message
        await ws.send(json.dumps(payload))
    except Exception as e:
        print(f"❌ WebSocket send error: {e}")

# 다운로드 함수
def download_video(url, output_path):
    ydl_opts = {
        'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]',
        'outtmpl': output_path,
        'merge_output_format': 'mp4',
        'external_downloader': 'aria2c',
        'external_downloader_args': ['-x', '16', '-k', '1M']
    }
    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url])
        print(f"✅ Downloaded: {output_path}")
        return output_path
    except Exception as e:
        print(f"❌ Download error: {e}")
        return None

# 비디오 길이 가져오기
def get_video_duration(file_path):
    try:
        result = subprocess.run([
            "ffprobe", "-v", "error", "-show_entries",
            "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", file_path
        ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        return float(result.stdout.strip())
    except Exception as e:
        print(f"❌ Duration error: {e}")
        return None

# atempo 체이닝
def build_atempo_chain(speed):
    chain = []
    while speed > 2.0:
        chain.append("atempo=2.0")
        speed /= 2.0
    while speed < 0.5:
        chain.append("atempo=0.5")
        speed *= 2.0
    chain.append(f"atempo={round(speed, 2)}")
    return ",".join(chain)

# 개별 변환
async def convert_video(ws, input_file, output_file, speed, session_id):
    duration = get_video_duration(input_file)
    if not duration:
        print("❌ Cannot get video duration.")
        return None

    new_duration = duration / speed
    target_total_kbps = (TARGET_FILESIZE_BYTES * 8) / new_duration / 1000
    audio_kbps = 128
    video_kbps = int(target_total_kbps - audio_kbps)
    if video_kbps <= 0:
        print("❌ Bitrate too low.")
        return None

    atempo_chain = build_atempo_chain(speed)
    setpts_value = 1 / speed

    cmd = [
        "ffmpeg", "-y", "-hwaccel", "cuda", "-i", input_file,
        "-filter_complex", f"[0:v]setpts={setpts_value}*PTS[v];[0:a]{atempo_chain}[a]",
        "-map", "[v]", "-map", "[a]",
        "-c:v", "h264_nvenc", "-preset", "p1", "-rc", "vbr",
        "-b:v", f"{video_kbps}k", "-maxrate", f"{video_kbps}k", "-bufsize", f"{video_kbps*2}k",
        "-c:a", "aac", "-b:a", f"{audio_kbps}k",
        "-progress", "pipe:1", output_file
    ]

    print(f"✨ Running: {' '.join(cmd)}")

    await send_ws_message(ws, session_id, 2)  # 변환 시작 알림 (2%)

    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    last_progress = 2

    for line in process.stdout:
        if line.startswith("out_time_ms="):
            try:
                out_time_ms = int(line.strip().split('=')[1])
                current_sec = out_time_ms / 1_000_000
                progress = int((current_sec / new_duration) * 100)
                progress = min(progress, 100)
                if progress != last_progress:
                    await send_ws_message(ws, session_id, progress)
                    print(f"🚀 {progress}%")
                    last_progress = progress
            except Exception as e:
                print(f"❌ Parsing error: {e}")

    process.wait()

    await send_ws_message(ws, session_id, 100)  # 변환 완료

    if process.returncode == 0:
        print(f"✅ {speed}x Conversion Done: {output_file}")
        return output_file
    else:
        print(f"❌ ffmpeg Error. Return code: {process.returncode}")
        return None

# 메인 처리
async def main():
    # 오래된 파일 정리
    now = time.time()
    for f in os.listdir(DOWNLOAD_DIR):
        path = os.path.join(DOWNLOAD_DIR, f)
        if os.path.isfile(path) and now - os.path.getmtime(path) > 300:
            os.remove(path)
            print(f"🗑 Deleted {path}")

    async with websockets.connect(WEBSOCKET_SERVER_URL) as ws:
        # ✅ 다운로드 시작 알림
        await send_ws_message(ws, session_id, 0, message="다운로드 중...")

        # 다운로드
        downloaded = download_video(url, output_path)
        if not downloaded or not os.path.exists(downloaded):
            print("❌ Download failed.")
            sys.exit(2)

        converted_files = []
        for speed in speed_options:
            output_file = os.path.join(DOWNLOAD_DIR, f"{base_filename}_{speed}x.mp4")
            result = await convert_video(ws, downloaded, output_file, float(speed), session_id)
            if result and os.path.exists(output_file):
                converted_files.append(output_file)

        # 원본 삭제
        if os.path.exists(downloaded):
            os.remove(downloaded)
            print(f"🗑 Removed original file.")

        # 결과 출력
        print(json.dumps({"converted_files": converted_files}))

if __name__ == "__main__":
    asyncio.run(main())