database.py

from datetime import datetime, timezone
from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime
from sqlalchemy.orm import sessionmaker, declarative_base
from config import DATA_DIR, logger

import os
import threading
import time
import shutil

# Setup database path
DB_PATH = DATA_DIR / "alpha_results.db"

# Restore database from Hugging Face Dataset if applicable
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 restore_db_from_hf():
    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}...")
    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
        )
        DATA_DIR.mkdir(parents=True, exist_ok=True)
        shutil.copy(cached_path, DB_PATH)
        logger.info(f"Successfully restored database from Hugging Face Dataset to {DB_PATH}.")
    except ImportError:
        logger.warning("huggingface_hub package is not installed. Database restore skipped.")
    except Exception as e:
        logger.warning(f"Could not restore database from Hugging Face Dataset (this is normal if it is the first run): {e}")

    try:
        from huggingface_hub import hf_hub_download
        settings_path = DATA_DIR / "miner_settings.json"
        cached_settings_path = hf_hub_download(
            repo_id=HF_DATASET_ID,
            filename="miner_settings.json",
            repo_type="dataset",
            token=HF_TOKEN
        )
        shutil.copy(cached_settings_path, settings_path)
        logger.info("Successfully restored miner_settings.json from Hugging Face Dataset.")
    except Exception as e:
        logger.warning(f"Could not restore miner_settings.json from Hugging Face Dataset: {e}")
        
    try:
        from huggingface_hub import hf_hub_download
        cache_path = DATA_DIR / "data_fields_cache.json"
        cached_cache_path = hf_hub_download(
            repo_id=HF_DATASET_ID,
            filename="data_fields_cache.json",
            repo_type="dataset",
            token=HF_TOKEN
        )
        shutil.copy(cached_cache_path, cache_path)
        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 from Hugging Face Dataset: {e}")

# Try to restore backup before SQLAlchemy engine is initialized
restore_db_from_hf()

DATABASE_URL = f"sqlite:///{DB_PATH}"

def backup_db_to_hf():
    if not HF_TOKEN:
        logger.warning("HF_TOKEN environment variable is not set. Database backup to Hugging Face Dataset skipped.")
        return
    if not HF_DATASET_ID:
        logger.warning("HF_DATASET_ID and SPACE_ID environment variables are not set. Database backup skipped.")
        return

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

    logger.info(f"Starting background database backup thread (interval: 15 mins) to: {HF_DATASET_ID}...")
    api = HfApi(token=HF_TOKEN)
    
    # Create the dataset repo if it doesn't exist
    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_backup = True
    while True:
        try:
            # Sleep for 60 seconds on the first loop, and 15 minutes thereafter
            if is_first_backup:
                time.sleep(60)
                is_first_backup = False
            else:
                time.sleep(15 * 60)
            
            if not DB_PATH.exists():
                logger.debug("Database file does not exist yet. Skipping backup cycle.")
                continue
                
            logger.info(f"Backing up database to Hugging Face Dataset: {HF_DATASET_ID}...")
            temp_db_path = DB_PATH.parent / "alpha_results_backup_temp.db"
            try:
                shutil.copy2(DB_PATH, temp_db_path)
                api.upload_file(
                    path_or_fileobj=str(temp_db_path),
                    path_in_repo="alpha_results.db",
                    repo_id=HF_DATASET_ID,
                    repo_type="dataset",
                )
                logger.info("Database backup successfully uploaded to Hugging Face Dataset.")
            except Exception as upload_error:
                logger.error(f"Failed to upload database file to Hugging Face: {upload_error}")
            finally:
                if temp_db_path.exists():
                    os.remove(temp_db_path)
                    
            # Also backup miner_settings.json if it exists
            settings_path = DATA_DIR / "miner_settings.json"
            if settings_path.exists():
                logger.info("Backing up miner_settings.json to Hugging Face Dataset...")
                try:
                    api.upload_file(
                        path_or_fileobj=str(settings_path),
                        path_in_repo="miner_settings.json",
                        repo_id=HF_DATASET_ID,
                        repo_type="dataset",
                    )
                    logger.info("miner_settings.json backup successfully uploaded to Hugging Face Dataset.")
                except Exception as upload_error:
                    logger.error(f"Failed to upload miner_settings.json to Hugging Face: {upload_error}")
                    
            # Also backup data_fields_cache.json if it exists
            cache_path = DATA_DIR / "data_fields_cache.json"
            if cache_path.exists():
                logger.info("Backing up data_fields_cache.json to Hugging Face Dataset...")
                try:
                    api.upload_file(
                        path_or_fileobj=str(cache_path),
                        path_in_repo="data_fields_cache.json",
                        repo_id=HF_DATASET_ID,
                        repo_type="dataset",
                    )
                    logger.info("data_fields_cache.json backup successfully uploaded to Hugging Face Dataset.")
                except Exception as upload_error:
                    logger.error(f"Failed to upload data_fields_cache.json to Hugging Face: {upload_error}")
                    
            # Also backup miner_production.log if it exists
            prod_log_path = DATA_DIR / "miner_production.log"
            if prod_log_path.exists():
                logger.info("Backing up miner_production.log to Hugging Face Dataset...")
                try:
                    temp_log_path = DATA_DIR / "miner_production_backup_temp.log"
                    shutil.copy2(prod_log_path, temp_log_path)
                    api.upload_file(
                        path_or_fileobj=str(temp_log_path),
                        path_in_repo="miner_production.log",
                        repo_id=HF_DATASET_ID,
                        repo_type="dataset",
                    )
                    logger.info("miner_production.log backup successfully uploaded to Hugging Face Dataset.")
                except Exception as upload_error:
                    logger.error(f"Failed to upload miner_production.log to Hugging Face: {upload_error}")
                finally:
                    if temp_log_path.exists():
                        try:
                            os.remove(temp_log_path)
                        except Exception:
                            pass
                    
        except Exception as e:
            logger.error(f"Error during background database backup to Hugging Face: {e}")

def start_backup_thread():
    if HF_TOKEN and HF_DATASET_ID:
        t = threading.Thread(target=backup_db_to_hf, daemon=True, name="HF_DB_Backup")
        t.start()
        logger.info("Database backup thread started successfully.")
    else:
        logger.warning("HF_TOKEN or HF_DATASET_ID/SPACE_ID not configured. Backup thread not started.")

# Start the periodic background backup thread
start_backup_thread()

# Initialize SQLAlchemy Engine and Base
engine = create_engine(DATABASE_URL, echo=False, connect_args={"timeout": 30})
Base = declarative_base()
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

class AlphaRecord(Base):
    __tablename__ = "alpha_records"

    id = Column(Integer, primary_key=True, index=True)
    expression = Column(String, unique=True, nullable=False)
    region = Column(String, default='USA')
    universe = Column(String, default='TOP3000')
    delay = Column(Integer, default=1)
    decay = Column(Integer, default=10)
    truncation = Column(Float, default=0.01)
    neutralization = Column(String, default='SUBINDUSTRY')
    
    # Results
    sharpe = Column(Float, nullable=True)
    fitness = Column(Float, nullable=True)
    turnover = Column(Float, nullable=True)
    returns = Column(Float, nullable=True)
    
    # Metadata
    status = Column(String, default='GENERATED') # e.g., 'GENERATED', 'SIMULATING', 'FAILED', 'PASSED', 'SUBMITTED'
    error_log = Column(String, nullable=True)
    alpha_id = Column(String, nullable=True)
    date_created = Column(DateTime, default=lambda: datetime.now(timezone.utc))
    date_submitted = Column(DateTime, nullable=True)
    retry_count = Column(Integer, default=0)

    def __repr__(self):
        return f"<AlphaRecord(id={self.id}, expression='{self.expression}', status='{self.status}')>"

def init_db():
    """Initializes the SQLite database and creates all tables."""
    logger.info(f"Initializing database at {DB_PATH}...")
    Base.metadata.create_all(bind=engine)
    
    # Enable WAL mode for better concurrency
    try:
        import sqlite3
        conn = sqlite3.connect(str(DB_PATH))
        cursor = conn.cursor()
        cursor.execute("PRAGMA journal_mode=WAL")
        cursor.execute("PRAGMA synchronous=NORMAL")
        conn.commit()
        conn.close()
        logger.info("SQLite WAL mode and synchronous settings configured successfully.")
    except Exception as e:
        logger.error(f"Failed to configure SQLite WAL mode: {e}")
        
    # Check if alpha_id or date_submitted columns exist, if not, add them safely
    try:
        import sqlite3
        conn = sqlite3.connect(str(DB_PATH))
        cursor = conn.cursor()
        cursor.execute("PRAGMA table_info(alpha_records)")
        columns = [col[1] for col in cursor.fetchall()]
        
        if "alpha_id" not in columns:
            logger.info("Adding alpha_id column to alpha_records table...")
            cursor.execute("ALTER TABLE alpha_records ADD COLUMN alpha_id VARCHAR")
            conn.commit()
            logger.info("alpha_id column added successfully.")
            
        if "date_submitted" not in columns:
            logger.info("Adding date_submitted column to alpha_records table...")
            cursor.execute("ALTER TABLE alpha_records ADD COLUMN date_submitted TIMESTAMP")
            conn.commit()
            logger.info("date_submitted column added successfully.")
            
        if "retry_count" not in columns:
            logger.info("Adding retry_count column to alpha_records table...")
            cursor.execute("ALTER TABLE alpha_records ADD COLUMN retry_count INTEGER DEFAULT 0")
            conn.commit()
            logger.info("retry_count column added successfully.")
            
        # Migrate SP500 to TOPSP500 in existing records
        logger.info("Running database migration to replace SP500 with TOPSP500...")
        cursor.execute("UPDATE alpha_records SET universe = 'TOPSP500' WHERE universe = 'SP500'")
        conn.commit()
        logger.info(f"Universe migration query completed. Rows modified: {cursor.rowcount}")
        
        conn.close()
    except Exception as e:
        logger.error(f"Error checking/adding columns to database: {e}")
        
    logger.info("Database initialized successfully.")

# Automatically initialize the database and apply migrations on import
init_db()