import os
import sys
import shutil
import sqlite3
from huggingface_hub import snapshot_download, HfApi

# Configuration
REPO_ID = os.environ.get("DATASET_REPO_ID")
HF_TOKEN = os.environ.get("HF_TOKEN")

def verify_data():
    db_path = "data_repo/database.db"
    if not os.path.exists(db_path):
        print(f"VERIFICATION FAILED: {db_path} does not exist.")
        return False
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        cursor.execute("SELECT username FROM users WHERE username = ?", ("akshit",))
        row = cursor.fetchone()
        conn.close()
        
        if row:
            print("VERIFICATION SUCCESS: User 'akshit' found in database.")
            return True
        else:
            print("VERIFICATION FAILED: User 'akshit' NOT found in database.")
            return False
    except Exception as e:
        print(f"VERIFICATION ERROR: {e}")
        return False

def download():
    if not REPO_ID:
        print("DATASET_REPO_ID not set, skipping download.")
        return
    
    print(f"Downloading data from {REPO_ID}...")
    try:
        # snapshot_download is more efficient for many files than the CLI
        snapshot_download(
            repo_id=REPO_ID,
            repo_type="dataset",
            local_dir="data_repo",
            token=HF_TOKEN,
            max_workers=8
        )
        print("Download successful.")
        verify_data()
    except Exception as e:
        print(f"Download failed: {e}")

def upload():
    if not REPO_ID:
        print("DATASET_REPO_ID not set, skipping upload.")
        return
    if not HF_TOKEN:
        print("HF_TOKEN not set, skipping upload.")
        return

    print(f"Uploading data to {REPO_ID}...")
    try:
        api = HfApi(token=HF_TOKEN)
        api.upload_folder(
            folder_path="data_repo",
            repo_id=REPO_ID,
            repo_type="dataset",
            delete_patterns="*", # Optional: sync deletion if needed
        )
        print("Upload successful.")
    except Exception as e:
        print(f"Upload failed: {e}")

def init_local():
    """Ensure data_repo has the necessary structure if download failed or it's new."""
    os.makedirs("data_repo/output", exist_ok=True)
    os.makedirs("data_repo/processed", exist_ok=True)
    os.makedirs("data_repo/uploads", exist_ok=True)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python hf_sync.py [download|upload|init|verify]")
        sys.exit(1)
    
    action = sys.argv[1]
    if action == "download":
        download()
    elif action == "upload":
        upload()
    elif action == "init":
        init_local()
    elif action == "verify":
        verify_data()
    else:
        print(f"Unknown action: {action}")