#!/usr/bin/env python3
"""
Migration script to add market_snapshots table to existing database.

Usage:
    python scripts/migrate_add_market_snapshots.py
"""

import sqlite3
from pathlib import Path

# Database path
DB_PATH = Path("data/trading_system.db")

def migrate():
    """Add market_snapshots table if it doesn't exist"""

    if not DB_PATH.exists():
        print(f"❌ Database not found: {DB_PATH}")
        return False

    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()

    try:
        # Check if table already exists
        cursor.execute("""
            SELECT name FROM sqlite_master
            WHERE type='table' AND name='market_snapshots'
        """)

        if cursor.fetchone():
            print("ℹ️  Table 'market_snapshots' already exists. Skipping migration.")
            return True

        # Create market_snapshots table
        print("Creating market_snapshots table...")
        cursor.execute("""
            CREATE TABLE market_snapshots (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                date DATETIME NOT NULL,
                market VARCHAR(20) NOT NULL,
                ticker_count INTEGER,
                avg_discount_pct REAL,
                undervalued_count INTEGER,
                overvalued_count INTEGER,
                sector_breakdown TEXT,
                top_opportunities TEXT,
                created_at DATETIME
            )
        """)

        # Create index
        print("Creating index idx_market_snapshot_date...")
        cursor.execute("""
            CREATE INDEX idx_market_snapshot_date
            ON market_snapshots (market, date)
        """)

        conn.commit()
        print("✅ Migration completed successfully!")
        print(f"✅ Table 'market_snapshots' created")
        print(f"✅ Index 'idx_market_snapshot_date' created")

        # Verify
        cursor.execute("SELECT COUNT(*) FROM market_snapshots")
        count = cursor.fetchone()[0]
        print(f"✅ Verified: table is empty (count={count})")

        return True

    except sqlite3.Error as e:
        print(f"❌ Migration failed: {e}")
        conn.rollback()
        return False

    finally:
        conn.close()

if __name__ == "__main__":
    print("=" * 60)
    print("Warren AI - Database Migration")
    print("Adding market_snapshots table for Market Radar feature")
    print("=" * 60)
    print()

    success = migrate()

    print()
    if success:
        print("✅ Database migration completed successfully!")
    else:
        print("❌ Database migration failed!")

    print("=" * 60)