import os
import sqlite3
import argparse
from pathlib import Path
from huggingface_hub import HfApi, login

def get_db_path():
    # Priority list of where the DB might be
    paths = [
        "/data/tracker.db",
        "/home/adhyansh/datasets/yt/hf_space_host/db/tracker.db",
        "/home/adhyansh/datasets/yt/scripts/db/tracker.db"
    ]
    for p in paths:
        if os.path.exists(p):
            return p
    return None

def load_hf_token():
    token = os.environ.get("HF_TOKEN", "")
    if not token:
        p = Path.home() / ".cache" / "huggingface" / "token"
        if p.exists():
            token = p.read_text().strip()
    return token

def main():
    parser = argparse.ArgumentParser(description="List files in an HF Repo and mark matching videos as 'done' in the local database.")
    parser.add_argument("--repo-id", default="AdhyanshVerma/YT", help="Hugging Face Repository ID")
    parser.add_argument("--db-path", default=None, help="Explicit path to the tracker.db SQLite database")
    args = parser.parse_args()

    # 1. Connect to Database
    db_path = args.db_path or get_db_path()
    if not db_path:
        print("❌ Error: Could not find tracker.db in default locations. Please specify with --db-path")
        return
        
    print(f"✅ Using Database: {db_path}")
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

    # Fetch all video_ids that are NOT done
    try:
        cursor.execute("SELECT video_id FROM videos WHERE status != 'done'")
        pending_videos = [row[0] for row in cursor.fetchall()]
        print(f"📊 Found {len(pending_videos)} videos in DB that are not marked as 'done'.")
    except Exception as e:
        print(f"❌ Database error: {e}")
        return

    if not pending_videos:
        print("✅ All videos are already marked as done. Exiting.")
        return

    # 2. Authenticate with Hugging Face
    token = load_hf_token()
    if token:
        login(token=token)
        print("✅ Authenticated with Hugging Face token.")
    else:
        print("⚠️ Warning: No HF_TOKEN found. If the repo is private, this might fail.")

    api = HfApi()
    
    # 3. List all files from the target repository
    repo_types_to_try = ["dataset", "space", "model"]
    files = []
    
    print(f"🔍 Searching for files in repository '{args.repo_id}'...")
    for repo_type in repo_types_to_try:
        try:
            # We use recursive=True to list ALL files without downloading them
            files = api.list_repo_files(repo_id=args.repo_id, repo_type=repo_type)
            print(f"✅ Successfully listed {len(files)} files as a '{repo_type}' repo.")
            break
        except Exception as e:
            # Continue checking other repo types
            pass
            
    if not files:
        print(f"❌ Error: Could not find or list files in repository '{args.repo_id}'. Make sure it exists and you have access.")
        return
        
    # 4. Check files against database and modify DB
    print("\n🔄 Analyzing files against the database...")
    
    marked_done_count = 0

    for vid in pending_videos:
        match_found = False
        
        # Search for the video ID anywhere in the full file paths (this catches directory names like staging/.../Jb9Jx0KkJl8/...)
        for f in files:
            if vid in f:
                match_found = True
                break
                
        if match_found:
            print(f"  🟢 Match found for video: {vid} - Marking as 'done'!")
            cursor.execute("UPDATE videos SET status='done' WHERE video_id=?", (vid,))
            marked_done_count += 1
            
    # 5. Save changes
    if marked_done_count > 0:
        conn.commit()
        print(f"\n🎉 Success! Updated {marked_done_count} videos to 'done' in the database.")
    else:
        print("\n⚠️ No new matching videos found in the repository. No database changes made.")
        
    conn.close()

if __name__ == "__main__":
    main()