import os
import sys
import shutil
import sqlite3
import json
import time
from datetime import datetime
from huggingface_hub import snapshot_download, HfApi

# Configuration
REPO_ID = os.environ.get("DATASET_REPO_ID")
HF_TOKEN = os.environ.get("HF_TOKEN")
DATA_DIR = "data_repo"
DB_FILE = os.path.join(DATA_DIR, "database.db")
METADATA_FILE = os.path.join(DATA_DIR, "sync_metadata.json")
LOCK_FILE = "/tmp/hf_sync.lock"

def get_metadata():
    if os.path.exists(METADATA_FILE):
        try:
            with open(METADATA_FILE, 'r') as f:
                return json.load(f)
        except:
            pass
    return {"version": 0, "last_sync": None, "source": None}

def update_metadata(action):
    meta = get_metadata()
    meta["version"] += 1
    meta["last_sync"] = datetime.now().isoformat()
    meta["source"] = os.environ.get("HOSTNAME", "local")
    meta["last_action"] = action
    with open(METADATA_FILE, 'w') as f:
        json.dump(meta, f, indent=2)

def safe_db_backup():
    """Create a consistent snapshot of the SQLite database."""
    if not os.path.exists(DB_FILE):
        return
    
    print("Creating consistent database snapshot...")
    backup_db = DB_FILE + ".tmp"
    try:
        source_conn = sqlite3.connect(DB_FILE)
        dest_conn = sqlite3.connect(backup_db)
        with dest_conn:
            source_conn.backup(dest_conn)
        source_conn.close()
        dest_conn.close()
        shutil.move(backup_db, DB_FILE)
    except Exception as e:
        print(f"Database backup failed: {e}")
        if os.path.exists(backup_db):
            os.remove(backup_db)

def verify_data():
    if not os.path.exists(DB_FILE):
        print(f"VERIFICATION FAILED: {DB_FILE} does not exist.")
        return False
    
    try:
        conn = sqlite3.connect(DB_FILE)
        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.")
            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} (Restore)...")
    try:
        snapshot_download(
            repo_id=REPO_ID,
            repo_type="dataset",
            local_dir=DATA_DIR,
            token=HF_TOKEN,
            max_workers=8
        )
        print("Download successful.")
        update_metadata("restore")
        verify_data()
    except Exception as e:
        print(f"Download failed: {e}")

def upload():
    if not REPO_ID or not HF_TOKEN:
        return

    # Simple lock check
    if os.path.exists(LOCK_FILE):
        # If lock is older than 10 mins, assume it's stale
        if time.time() - os.path.getmtime(LOCK_FILE) < 600:
            print("Upload already in progress, skipping...")
            return
    
    try:
        with open(LOCK_FILE, 'w') as f: f.write(str(os.getpid()))
        
        print(f"Starting scheduled backup to {REPO_ID}...")
        safe_db_backup()
        update_metadata("backup")
        
        api = HfApi(token=HF_TOKEN)
        api.upload_folder(
            folder_path=DATA_DIR,
            repo_id=REPO_ID,
            repo_type="dataset",
            commit_message=f"Automated backup v{get_metadata()['version']}"
        )
        print("Upload successful.")
    except Exception as e:
        print(f"Upload failed: {e}")
    finally:
        if os.path.exists(LOCK_FILE):
            os.remove(LOCK_FILE)

def init_local():
    os.makedirs(f"{DATA_DIR}/output", exist_ok=True)
    os.makedirs(f"{DATA_DIR}/processed", exist_ok=True)
    os.makedirs(f"{DATA_DIR}/uploads", exist_ok=True)

if __name__ == "__main__":
    action = sys.argv[1] if len(sys.argv) > 1 else "help"
    if action == "download": download()
    elif action == "upload": upload()
    elif action == "init": init_local()
    elif action == "verify": verify_data()
    else:
        print("Usage: python hf_sync.py [download|upload|init|verify]")