core/backup.py

import os
import threading
import time
import shutil
import logging
from pathlib import Path

logger = logging.getLogger("HF_Backup")

# Constants
BASE_DIR = Path(__file__).resolve().parent.parent
DATA_DIR = BASE_DIR / "data"
DB_PATH = DATA_DIR / "quantforge_v2.db"
OLD_DB_PATH = DATA_DIR / "alpha_results.db"
LOG_FILE = DATA_DIR / "miner_production.log"
SETTINGS_FILE = DATA_DIR / "miner_settings.json"
CACHE_FILE = DATA_DIR / "data_fields_cache.json"

HF_TOKEN = os.getenv("HF_TOKEN")
SPACE_ID = os.getenv("SPACE_ID")
DEFAULT_DATASET_ID = f"{SPACE_ID}-backup" if SPACE_ID else None
HF_DATASET_ID = os.getenv("HF_DATASET_ID") or DEFAULT_DATASET_ID

def checkpoint_db():
    """
    Executes a checkpoint on SQLite WAL mode database to ensure all logs 
    are written into the main .db file before backing up.
    """
    if not DB_PATH.exists():
        return
    import sqlite3
    try:
        conn = sqlite3.connect(str(DB_PATH))
        cursor = conn.cursor()
        cursor.execute("PRAGMA wal_checkpoint(TRUNCATE)")
        conn.commit()
        conn.close()
        logger.info("SQLite WAL checkpoint completed successfully.")
    except Exception as e:
        logger.error(f"Failed to checkpoint SQLite DB: {e}")

def migrate_legacy_db():
    """Migrates data from old alpha_records table to the new alphas table if present."""
    if not DB_PATH.exists():
        return
        
    import sqlite3
    conn = None
    try:
        conn = sqlite3.connect(str(DB_PATH))
        cursor = conn.cursor()
        
        # Check if alphas table exists to ensure it has all columns (schema migration check)
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='alphas'")
        has_alphas = cursor.fetchone()
        if has_alphas:
            cursor.execute("PRAGMA table_info(alphas)")
            columns = [col[1] for col in cursor.fetchall()]
            expected_columns = {
                "trash_reason": "VARCHAR",
                "pasteurization": "VARCHAR NOT NULL DEFAULT 'On'",
                "nan_handling": "VARCHAR",
                "unit_handling": "VARCHAR",
                "lookback": "INTEGER",
                "test_period": "VARCHAR",
                "max_trade": "FLOAT",
                "max_position": "FLOAT",
                "language": "VARCHAR NOT NULL DEFAULT 'FASTEXPR'"
            }
            for col_name, col_type in expected_columns.items():
                if col_name not in columns:
                    logger.info(f"Schema migration: Adding missing column '{col_name}' to 'alphas' table...")
                    try:
                        cursor.execute(f"ALTER TABLE alphas ADD COLUMN {col_name} {col_type}")
                    except Exception as e:
                        logger.error(f"Failed to add column '{col_name}': {e}")
            conn.commit()
            
        # Check if alpha_records table exists
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='alpha_records'")
        has_legacy = cursor.fetchone()
        
        if not has_legacy:
            return
            
        logger.info("Legacy alpha_records table found. Migrating data to new alphas table...")
        
        # Create alphas table manually if not created yet by SQLAlchemy
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='alphas'")
        has_new = cursor.fetchone()
        
        if not has_new:
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS alphas (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    expression VARCHAR UNIQUE NOT NULL,
                    language VARCHAR NOT NULL DEFAULT 'FASTEXPR',
                    region VARCHAR NOT NULL,
                    universe VARCHAR NOT NULL,
                    delay INTEGER NOT NULL,
                    decay INTEGER NOT NULL,
                    truncation FLOAT NOT NULL,
                    neutralization VARCHAR NOT NULL,
                    pasteurization VARCHAR NOT NULL DEFAULT 'On',
                    nan_handling VARCHAR,
                    unit_handling VARCHAR,
                    lookback INTEGER,
                    test_period VARCHAR,
                    max_trade FLOAT,
                    max_position FLOAT,
                    sharpe FLOAT,
                    fitness FLOAT,
                    turnover FLOAT,
                    returns FLOAT,
                    status VARCHAR NOT NULL DEFAULT 'PENDING_GEN',
                    retry_count INTEGER NOT NULL DEFAULT 0,
                    error_log VARCHAR,
                    trash_reason VARCHAR,
                    alpha_id VARCHAR,
                    date_created TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
                    date_submitted TIMESTAMP,
                    created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
                )
            """)
            conn.commit()
            
        # Get data from legacy table
        cursor.execute("SELECT id, expression, region, universe, delay, decay, truncation, neutralization, sharpe, fitness, turnover, returns, status, error_log, alpha_id, date_created, date_submitted, retry_count FROM alpha_records")
        rows = cursor.fetchall()
        
        inserted_cnt = 0
        for row in rows:
            (r_id, expr, reg, univ, delay, decay, trunc, neut, sharpe, fit, turn, ret, status, err, a_id, dt_create, dt_submit, retries) = row
            
            # Map legacy status to new v2 status if needed
            new_status = status
            if status == 'FAILED':
                new_status = 'TRASHED'
            elif status == 'PASSED':
                new_status = 'PENDING_SUBMIT'
            elif status == 'GENERATED':
                new_status = 'PENDING_GEN'
                
            try:
                cursor.execute("""
                    INSERT OR IGNORE INTO alphas (
                        id, expression, region, universe, delay, decay, truncation, neutralization,
                        sharpe, fitness, turnover, returns, status, error_log, alpha_id, date_created, date_submitted, retry_count
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """, (r_id, expr, reg, univ, delay, decay, trunc, neut, sharpe, fit, turn, ret, new_status, err, a_id, dt_create, dt_submit, retries))
                inserted_cnt += 1
            except Exception:
                pass
                
        conn.commit()
        logger.info(f"Successfully migrated {inserted_cnt} legacy alphas from alpha_records to alphas.")
        
        # Drop legacy table
        cursor.execute("DROP TABLE alpha_records")
        conn.commit()
        logger.info("Dropped legacy alpha_records table to finalize migration.")
        
    except Exception as ex:
        logger.error(f"Failed to migrate legacy database: {ex}")
    finally:
        if conn:
            conn.close()

def restore_from_hf():
    """
    Downloads database files, settings, and cache from Hugging Face Dataset.
    If quantforge_v2.db doesn't exist but alpha_results.db does, copies it as fallback.
    """
    if not HF_TOKEN:
        logger.warning("HF_TOKEN environment variable is not set. Database restore from Hugging Face Dataset skipped.")
        return
    if not HF_DATASET_ID:
        logger.warning("HF_DATASET_ID and SPACE_ID environment variables are not set. Database restore skipped.")
        return

    logger.info(f"Checking for database backup in Hugging Face Dataset: {HF_DATASET_ID}...")
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    
    # 1. Try to restore quantforge_v2.db
    v2_restored = False
    try:
        from huggingface_hub import hf_hub_download
        cached_path = hf_hub_download(
            repo_id=HF_DATASET_ID,
            filename="quantforge_v2.db",
            repo_type="dataset",
            token=HF_TOKEN
        )
        shutil.copy(cached_path, DB_PATH)
        logger.info(f"Successfully restored quantforge_v2.db from Hugging Face Dataset to {DB_PATH}.")
        v2_restored = True
    except Exception as e:
        logger.info(f"Could not restore quantforge_v2.db (might not exist yet): {e}")

    # 2. Fallback to alpha_results.db if v2 wasn't restored and doesn't exist locally
    if not v2_restored and not DB_PATH.exists():
        try:
            from huggingface_hub import hf_hub_download
            cached_path = hf_hub_download(
                repo_id=HF_DATASET_ID,
                filename="alpha_results.db",
                repo_type="dataset",
                token=HF_TOKEN
            )
            shutil.copy(cached_path, DB_PATH)
            logger.info(f"Successfully restored alpha_results.db as quantforge_v2.db to {DB_PATH} (migration fallback).")
        except Exception as e2:
            logger.warning(f"Could not restore alpha_results.db as migration fallback: {e2}")

    # Migrate any old table formats in DB to new schema
    migrate_legacy_db()

    # 3. Restore miner_settings.json
    try:
        from huggingface_hub import hf_hub_download
        cached_path = hf_hub_download(
            repo_id=HF_DATASET_ID,
            filename="miner_settings.json",
            repo_type="dataset",
            token=HF_TOKEN
        )
        shutil.copy(cached_path, SETTINGS_FILE)
        logger.info("Successfully restored miner_settings.json from Hugging Face Dataset.")
    except Exception as e:
        logger.warning(f"Could not restore miner_settings.json: {e}")

    # 4. Restore data_fields_cache.json
    try:
        from huggingface_hub import hf_hub_download
        cached_path = hf_hub_download(
            repo_id=HF_DATASET_ID,
            filename="data_fields_cache.json",
            repo_type="dataset",
            token=HF_TOKEN
        )
        shutil.copy(cached_path, CACHE_FILE)
        logger.info("Successfully restored data_fields_cache.json from Hugging Face Dataset.")
    except Exception as e:
        logger.warning(f"Could not restore data_fields_cache.json: {e}")

def backup_worker():
    """Background worker that periodically uploads files to HF Dataset."""
    if not HF_TOKEN or not HF_DATASET_ID:
        return

    try:
        from huggingface_hub import HfApi
    except ImportError:
        logger.warning("huggingface_hub package is not installed. Background backup thread aborted.")
        return

    logger.info(f"Starting database backup loop to Hugging Face Dataset: {HF_DATASET_ID}...")
    api = HfApi(token=HF_TOKEN)

    # Ensure repository exists
    try:
        api.create_repo(
            repo_id=HF_DATASET_ID,
            repo_type="dataset",
            private=True,
            exist_ok=True
        )
    except Exception as e:
        logger.warning(f"Could not verify/create Hugging Face Dataset repository: {e}")

    is_first = True
    while True:
        try:
            # Sleep 1 minute on first cycle, 15 minutes subsequently
            if is_first:
                time.sleep(60)
                is_first = False
            else:
                time.sleep(15 * 60)

            # Upload quantforge_v2.db
            if DB_PATH.exists():
                logger.info("Uploading quantforge_v2.db to Hugging Face...")
                # Run checkpoint to merge WAL file content into the main database file
                checkpoint_db()
                temp_db = DB_PATH.parent / "quantforge_v2_backup_temp.db"
                try:
                    shutil.copy2(DB_PATH, temp_db)
                    api.upload_file(
                        path_or_fileobj=str(temp_db),
                        path_in_repo="quantforge_v2.db",
                        repo_id=HF_DATASET_ID,
                        repo_type="dataset"
                    )
                    logger.info("quantforge_v2.db backup uploaded successfully.")
                except Exception as ex:
                    logger.error(f"Failed to upload quantforge_v2.db: {ex}")
                finally:
                    if temp_db.exists():
                        try:
                            os.remove(temp_db)
                        except Exception:
                            pass

            # Upload miner_settings.json
            if SETTINGS_FILE.exists():
                logger.info("Uploading miner_settings.json to Hugging Face...")
                try:
                    api.upload_file(
                        path_or_fileobj=str(SETTINGS_FILE),
                        path_in_repo="miner_settings.json",
                        repo_id=HF_DATASET_ID,
                        repo_type="dataset"
                    )
                    logger.info("miner_settings.json backup uploaded successfully.")
                except Exception as ex:
                    logger.error(f"Failed to upload miner_settings.json: {ex}")

            # Upload data_fields_cache.json
            if CACHE_FILE.exists():
                logger.info("Uploading data_fields_cache.json to Hugging Face...")
                try:
                    api.upload_file(
                        path_or_fileobj=str(CACHE_FILE),
                        path_in_repo="data_fields_cache.json",
                        repo_id=HF_DATASET_ID,
                        repo_type="dataset"
                    )
                    logger.info("data_fields_cache.json backup uploaded successfully.")
                except Exception as ex:
                    logger.error(f"Failed to upload data_fields_cache.json: {ex}")

            # Upload miner_production.log
            if LOG_FILE.exists():
                logger.info("Uploading miner_production.log to Hugging Face...")
                temp_log = LOG_FILE.parent / "miner_production_backup_temp.log"
                try:
                    shutil.copy2(LOG_FILE, temp_log)
                    api.upload_file(
                        path_or_fileobj=str(temp_log),
                        path_in_repo="miner_production.log",
                        repo_id=HF_DATASET_ID,
                        repo_type="dataset"
                    )
                    logger.info("miner_production.log backup uploaded successfully.")
                except Exception as ex:
                    logger.error(f"Failed to upload miner_production.log: {ex}")
                finally:
                    if temp_log.exists():
                        try:
                            os.remove(temp_log)
                        except Exception:
                            pass

        except Exception as global_ex:
            logger.error(f"Error in backup loop: {global_ex}")

def backup_settings_now():
    """Immediately uploads miner_settings.json to Hugging Face Dataset."""
    if not HF_TOKEN or not HF_DATASET_ID or not SETTINGS_FILE.exists():
        return
    try:
        from huggingface_hub import HfApi
        api = HfApi(token=HF_TOKEN)
        logger.info("Triggering immediate backup of miner_settings.json to Hugging Face...")
        api.upload_file(
            path_or_fileobj=str(SETTINGS_FILE),
            path_in_repo="miner_settings.json",
            repo_id=HF_DATASET_ID,
            repo_type="dataset"
        )
        logger.info("miner_settings.json immediately backed up successfully.")
    except Exception as e:
        logger.error(f"Failed to perform immediate backup of settings: {e}")

def start_backup():
    """Launches the background backup thread."""
    if HF_TOKEN and HF_DATASET_ID:
        t = threading.Thread(target=backup_worker, daemon=True, name="HF_DB_Backup")
        t.start()
        logger.info("Hugging Face background backup thread started.")
    else:
        logger.warning("HF_TOKEN or HF_DATASET_ID not configured. Backup thread disabled.")

# Restore DB and configuration files on module import
restore_from_hf()