import os
import sys
import shutil
import sqlite3
import json
import time
import hashlib
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, "last_db_hash": None}

def get_dir_stats():
    """Returns total file count across data directories to detect new uploads."""
    count = 0
    for d in ['uploads', 'processed', 'output']:
        path = os.path.join(DATA_DIR, d)
        if os.path.exists(path):
            count += len(os.listdir(path))
    return count

def get_file_hash(path):
    if not os.path.exists(path): return None
    hasher = hashlib.md5()
    with open(path, 'rb') as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hasher.update(chunk)
    return hasher.hexdigest()

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

def safe_db_backup():
    if not os.path.exists(DB_FILE): return None
    
    # We backup to a temp file to get a consistent hash/file
    backup_db = DB_FILE + ".bak"
    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()
        return backup_db
    except Exception as e:
        print(f"Database backup failed: {e}")
        return None

def download():
    if not REPO_ID: return
    print(f"Downloading data from {REPO_ID}...")
    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", get_file_hash(DB_FILE))
    except Exception as e:
        print(f"Download failed: {e}")

def upload():
    if not REPO_ID or not HF_TOKEN: return
    if os.path.exists(LOCK_FILE):
        if time.time() - os.path.getmtime(LOCK_FILE) < 600: return

    try:
        with open(LOCK_FILE, 'w') as f: f.write(str(os.getpid()))
        
        meta = get_metadata()
        backup_path = safe_db_backup()
        if not backup_path: return

        current_hash = get_file_hash(backup_path)
        current_count = get_dir_stats()

        # ONLY UPLOAD IF CHANGED
        if current_hash == meta.get("last_db_hash") and current_count == meta.get("file_count", 0):
            print("No changes detected (DB hash and file count match). skipping upload.")
            os.remove(backup_path)
            return

        print(f"Changes detected. Syncing to {REPO_ID}...")
        # Replace the real db with the consistent backup for upload
        shutil.move(backup_path, DB_FILE)
        update_metadata("backup", current_hash)
        
        api = HfApi(token=HF_TOKEN)
        api.upload_folder(
            folder_path=DATA_DIR,
            repo_id=REPO_ID,
            repo_type="dataset",
            commit_message=f"Auto-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():
    for d in ['output', 'processed', 'uploads']: os.makedirs(f"{DATA_DIR}/{d}", 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()
    else: print("Usage: python hf_sync.py [download|upload|init]")