#!/usr/bin/env python3
"""
Database Migration Script for Schema v4
Adds 3 new columns to fundamental_data table
"""

import sqlite3
from pathlib import Path

DB_PATH = Path("data/trading_system.db")

def migrate():
    """Add Schema v4 columns to fundamental_data table"""
    print("=" * 80)
    print("SCHEMA v4 DATABASE MIGRATION")
    print("=" * 80)

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

    # Backup first
    backup_path = DB_PATH.parent / f"{DB_PATH.name}.backup.schema_v4"
    import shutil
    shutil.copy2(DB_PATH, backup_path)
    print(f"\n✅ Backup created: {backup_path}")

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

    print("\n📊 Adding Schema v4 columns to fundamental_data table...")

    try:
        # Check if columns already exist
        cursor.execute("PRAGMA table_info(fundamental_data)")
        existing_columns = [row[1] for row in cursor.fetchall()]

        columns_to_add = [
            ("roic", "REAL"),
            ("interest_coverage", "REAL"),
            ("piotroski_fscore", "INTEGER")
        ]

        for col_name, col_type in columns_to_add:
            if col_name in existing_columns:
                print(f"   ⚠️  Column '{col_name}' already exists, skipping")
            else:
                cursor.execute(f"ALTER TABLE fundamental_data ADD COLUMN {col_name} {col_type}")
                print(f"   ✅ Added column: {col_name} ({col_type})")

        conn.commit()

        # Verify
        cursor.execute("PRAGMA table_info(fundamental_data)")
        all_columns = [row[1] for row in cursor.fetchall()]

        print("\n📋 Current fundamental_data schema:")
        for i, col in enumerate(all_columns, 1):
            marker = "🆕" if col in ['roic', 'interest_coverage', 'piotroski_fscore'] else "  "
            print(f"   {marker} {i}. {col}")

        print("\n✅ Migration completed successfully!")
        print("\n💡 Next steps:")
        print("   1. Run: python test_schema_v4.py")
        print("   2. Run full market scan: python warren_scan.py --market IT")

        return True

    except Exception as e:
        print(f"\n❌ Migration failed: {e}")
        conn.rollback()
        return False

    finally:
        conn.close()

if __name__ == "__main__":
    migrate()