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 pull_latest():
    sync_worker_profile()
    downloads_dir = Path.home() / "Downloads"
    existing = set(downloads_dir.glob("*.mp4"))
    
    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)
        
        dl_selector = 'button[aria-label*="Tải video xuống"], button[aria-label*="Download"], button:has-text("Tải video xuống")'
        buttons = page.locator(dl_selector)
        count = buttons.count()
        print(f"Total download buttons: {count}")
        if count == 0:
            print("No buttons found!")
            return False
            
        print("Clicking the latest download button...")
        btn = buttons.last
        btn.scroll_into_view_if_needed()
        btn.click(force=True)
        
        found = None
        for _ in range(15):
            time.sleep(1)
            new_files = [f for f in downloads_dir.glob("*.mp4") if f not in existing]
            if new_files:
                new_files.sort(key=lambda x: x.stat().st_mtime, reverse=True)
                if new_files[0].stat().st_size > 50000:
                    found = new_files[0]
                    break
                    
        if found:
            print(f"Downloaded: {found.name} ({found.stat().st_size} bytes)")
            # Let's save as shot08
            target_v = Path("04_Assets/videos/ep01_scene05_shot08_10s_v1.mp4").resolve()
            shutil.copy2(str(found), str(target_v))
            print(f"Saved to: {target_v}")
            
            # Extract tail frame
            kf_dir = Path("04_Assets/keyframes/ep01_scene05_shot08").resolve()
            kf_dir.mkdir(parents=True, exist_ok=True)
            kf_path = kf_dir / "clean_frame_239.jpg"
            cap = cv2.VideoCapture(str(target_v))
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, total_frames - 1))
            ret, frame = cap.read()
            if ret:
                cv2.imwrite(str(kf_path), frame, [cv2.IMWRITE_JPEG_QUALITY, 95])
                print(f"Extracted tail frame 239: {kf_path}")
            cap.release()
            return True
        else:
            print("Download file not found.")
            return False

if __name__ == "__main__":
    pull_latest()