import os
import sys
import re
import time
import shutil
import subprocess
from pathlib import Path
import cv2

if sys.platform == "win32":
    try:
        sys.stdout.reconfigure(encoding="utf-8")
        sys.stderr.reconfigure(encoding="utf-8")
    except Exception:
        pass

BASE_DIR = Path(__file__).resolve().parent.parent
VIDEOS_DIR = BASE_DIR / "04_Assets" / "videos"
KEYFRAMES_DIR = BASE_DIR / "04_Assets" / "keyframes"
DOWNLOADS_DIR = Path.home() / "Downloads"

def monitor_shot(shot_id: str, timeout_sec: int = 300):
    print(f"\n=======================================================")
    print(f"🎬 GIÁM SÁT TIẾN TRÌNH RENDER VÀ TẢI XUỐNG: {shot_id}")
    print(f"=======================================================")
    
    # Ghi nhận các file hiện có trong Downloads
    initial_downloads = set(DOWNLOADS_DIR.glob("*.mp4"))
    target_video_path = VIDEOS_DIR / f"{shot_id}_10s.mp4"
    
    # Ghi nhận các nút tải xuống đã có từ trước
    pre_snap = subprocess.run("agent-browser --session muse snapshot -i", shell=True, capture_output=True, text=True, encoding="utf-8").stdout
    pre_dl_refs = set(re.findall(r'button "Tải video xuống" \[ref=([^\]]+)\]', pre_snap))
    
    video_saved = False
    start_time = time.time()
    
    for attempt in range(timeout_sec // 5):
        snap = subprocess.run("agent-browser --session muse snapshot -i", shell=True, capture_output=True, text=True, encoding="utf-8").stdout
        is_busy = "Dừng nhiệm vụ" in snap or 'button "Dừng"' in snap
        current_dl_refs = re.findall(r'button "Tải video xuống" \[ref=([^\]]+)\]', snap)
        new_dl_refs = [r for r in current_dl_refs if r not in pre_dl_refs]
        
        elapsed = int(time.time() - start_time)
        status_msg = "Đang render..." if is_busy else ("Đã xong!" if new_dl_refs else "Đang xử lý...")
        print(f"      [{elapsed}s] Trạng thái: {status_msg}", flush=True)
        
        if (not is_busy) and new_dl_refs:
            last_dl_ref = new_dl_refs[-1]
            print(f"\n      [✓] Phát hiện video mới hoàn tất! Bấm tải xuống (@{last_dl_ref})...", flush=True)
            subprocess.run(f'agent-browser --session muse click "@{last_dl_ref}"', shell=True, capture_output=True, text=True)
            time.sleep(6)
            
            # Tìm file mới trong Downloads
            current_downloads = set(DOWNLOADS_DIR.glob("*.mp4"))
            new_files = current_downloads - initial_downloads
            if not new_files:
                recent = sorted(DOWNLOADS_DIR.glob("*.mp4"), key=lambda f: f.stat().st_mtime, reverse=True)
                if recent and (time.time() - recent[0].stat().st_mtime < 180):
                    new_files = [recent[0]]
            
            if new_files:
                downloaded = list(new_files)[0]
                shutil.copy2(downloaded, target_video_path)
                print(f"      ✓ Đã lưu video vào: {target_video_path}")
                video_saved = True
                break
        time.sleep(5)
        
    if not video_saved:
        print(f"[!] Hết thời gian chờ hoặc không tải được video.")
        return False
        
    # Trích xuất Tail Frame (Frame 239)
    print(f"\n[4/5] Đang trích xuất Tail Frame bằng OpenCV...")
    cap = cv2.VideoCapture(str(target_video_path))
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.set(cv2.CAP_PROP_POS_FRAMES, max(0, total_frames - 1))
    ret, tail_frame = cap.read()
    cap.release()
    
    shot_kf_dir = KEYFRAMES_DIR / shot_id
    shot_kf_dir.mkdir(parents=True, exist_ok=True)
    tail_frame_path = shot_kf_dir / "clean_frame_239.jpg"
    if ret and tail_frame is not None:
        cv2.imwrite(str(tail_frame_path), tail_frame)
        print(f"      ✓ Đã lưu Tail Frame: {tail_frame_path}")
        
    print(f"\n[✓] HOÀN TẤT THÀNH CÔNG: {shot_id}!\n")
    return True

if __name__ == "__main__":
    shot_id = sys.argv[1] if len(sys.argv) > 1 else "ep02_scene01_shot01"
    monitor_shot(shot_id)
