import sys
import time
import shutil
import cv2
from pathlib import Path
sys.path.append(str(Path("05_Production_Pipeline").resolve()))
from invisible_playwright import InvisiblePlaywright
from muse_invpw_driver import sync_worker_profile, DEFAULT_PROFILE

def download_completed_video():
    sync_worker_profile()
    target_video = Path("04_Assets/videos/ep01_scene05_shot07_10s_v1.mp4").resolve()
    target_kf = Path("04_Assets/keyframes/ep01_scene05_shot07/clean_frame_239.jpg").resolve()
    target_kf.parent.mkdir(parents=True, exist_ok=True)
    
    print(f"Connecting to Muse.ai with profile: {DEFAULT_PROFILE}...")
    with InvisiblePlaywright(profile_dir=DEFAULT_PROFILE, headless=True) as ctx:
        page = ctx.pages[0] if ctx.pages else ctx.new_page()
        page.goto("https://muse.ai", wait_until="domcontentloaded", timeout=30000)
        time.sleep(4)
        
        # Locate download button using aria-label
        dl_btn = page.locator('button[aria-label*="Tải video xuống"], button[aria-label*="Download"], button:has-text("Tải video xuống")').last
        print(f"Found download button: count={dl_btn.count()}")
        if dl_btn.count() == 0:
            print("No download button found!")
            return False
            
        print("Clicking download button and intercepting download...")
        with page.expect_download(timeout=60000) as dl_info:
            dl_btn.click()
            
        dl = dl_info.value
        temp_file = Path("04_Assets/temp_downloads/raw_shot07.mp4").resolve()
        temp_file.parent.mkdir(parents=True, exist_ok=True)
        dl.save_as(str(temp_file))
        print(f"Downloaded temp file size: {temp_file.stat().st_size} bytes")
        
        if temp_file.stat().st_size > 10000:
            shutil.move(str(temp_file), str(target_video))
            print(f"Moved to: {target_video}")
            
            # Extract tail frame 239
            cap = cv2.VideoCapture(str(target_video))
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            print(f"Total video frames: {total_frames}")
            target_idx = max(0, total_frames - 1)
            cap.set(cv2.CAP_PROP_POS_FRAMES, target_idx)
            ret, frame = cap.read()
            if ret:
                cv2.imwrite(str(target_kf), frame, [cv2.IMWRITE_JPEG_QUALITY, 95])
                print(f"Extracted tail frame 239: {target_kf}")
            cap.release()
            return True
        else:
            print("Downloaded file is too small or corrupt.")
            return False

if __name__ == "__main__":
    download_completed_video()