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 check_and_download_shot11():
    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(3)
        
        # Click the top chat item in sidebar: 'Tạo video Thúy Kiều tại mộ' or first thread
        thread_link = page.locator('button:has-text("Tạo video Thúy Kiều"), a:has-text("Tạo video Thúy Kiều")').first
        if thread_link.count() > 0:
            print("Clicking top thread:", thread_link.inner_text())
            thread_link.click()
            time.sleep(3)
            
        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"Download buttons count in this thread: {count}")
        
        if count > 0:
            btn = buttons.last
            btn.scroll_into_view_if_needed()
            btn.click(force=True)
            print("Clicked download button!")
            
            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)")
                target_v = Path("04_Assets/videos/ep01_scene05_shot11_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_shot11").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("No download button found in this thread.")
            return False

if __name__ == "__main__":
    check_and_download_shot11()