#!/usr/bin/env python3
"""
Migration Script: FundamentalData v1 → v2
=========================================

Adds 19 new fields for advanced fundamental analysis:
- Capital structure (debt, cash, enterprise value)
- Cash flows (operating, capex, FCF)
- Margins (gross, operating, EBITDA)
- Revenue/earnings metrics
- Share/dividend details

Schema version field:
- Existing records: schema_version = 1
- New records: schema_version = 2 (default in model)

Usage:
    python scripts/migrate_fundamentals_v2.py
"""

from datetime import datetime
from loguru import logger
from sqlalchemy import text

from src.database.db_manager import DatabaseManager


def add_column_if_not_exists(session, column_name, column_type="REAL"):
    """Add a column to fundamental_data table if it doesn't exist"""
    try:
        # Try to select from the column to see if it exists
        session.execute(text(f"SELECT {column_name} FROM fundamental_data LIMIT 1"))
        logger.info(f"  ✓ Column '{column_name}' already exists")
        return True
    except Exception:
        # Column doesn't exist, add it
        try:
            session.execute(text(f"ALTER TABLE fundamental_data ADD COLUMN {column_name} {column_type}"))
            session.commit()
            logger.success(f"  + Added column '{column_name}' ({column_type})")
            return True
        except Exception as e:
            logger.error(f"  ✗ Failed to add column '{column_name}': {e}")
            return False


def migrate_to_v2():
    """Execute migration from v1 to v2 schema"""

    logger.info("=" * 70)
    logger.info("FundamentalData Schema Migration: v1 → v2")
    logger.info("=" * 70)

    db = DatabaseManager()

    try:
        with db.get_session() as session:
            # Step 1: Check current state
            result = session.execute(
                text("SELECT COUNT(*) FROM fundamental_data")
            ).scalar()
            logger.info(f"Total FundamentalData records: {result}")

            # Step 2: Add new columns (19 fields + schema_version)
            logger.info("\nAdding new columns...")

            columns_to_add = [
                # Capital Structure (4)
                ("total_debt", "REAL"),
                ("total_cash", "REAL"),
                ("net_debt", "REAL"),
                # enterprise_value already exists, but we'll check

                # Cash Flow Metrics (3)
                ("operating_cashflow", "REAL"),
                ("capital_expenditures", "REAL"),
                ("free_cashflow", "REAL"),

                # Margin Metrics (3)
                ("gross_margin", "REAL"),
                ("ebitda_margin", "REAL"),
                ("ebitda", "REAL"),

                # Revenue/Earnings (4)
                ("total_revenue", "REAL"),
                ("ebit", "REAL"),
                ("net_income", "REAL"),
                ("trailing_eps", "REAL"),

                # Shares/Dividend Details (4)
                ("shares_outstanding", "REAL"),
                ("dividend_rate", "REAL"),
                ("ex_dividend_date", "DATETIME"),
                ("dividend_payment_date", "DATETIME"),

                # Schema Version (1)
                ("schema_version", "INTEGER DEFAULT 2"),
            ]

            all_success = True
            for col_name, col_type in columns_to_add:
                if not add_column_if_not_exists(session, col_name, col_type):
                    all_success = False

            if not all_success:
                logger.error("Some columns failed to add")
                return False

            # Step 3: Update existing records to schema_version = 1
            logger.info("\nMarking existing records as schema_version = 1...")

            update_result = session.execute(
                text("""
                    UPDATE fundamental_data
                    SET schema_version = 1
                    WHERE schema_version IS NULL
                """)
            )
            session.commit()

            updated_count = update_result.rowcount
            logger.success(f"✓ Updated {updated_count} records to schema_version = 1")

            # Step 4: Verify migration
            v1_count = session.execute(
                text("SELECT COUNT(*) FROM fundamental_data WHERE schema_version = 1")
            ).scalar()

            v2_count = session.execute(
                text("SELECT COUNT(*) FROM fundamental_data WHERE schema_version = 2")
            ).scalar()

            logger.info("\n" + "=" * 70)
            logger.info("Migration Complete - Summary:")
            logger.info(f"  - v1 records (basic fundamentals): {v1_count}")
            logger.info(f"  - v2 records (extended fundamentals): {v2_count}")
            logger.info(f"  - Total records: {v1_count + v2_count}")
            logger.info("=" * 70)
            logger.success("✓ Migration successful!")

            return True

    except Exception as e:
        logger.error(f"Migration failed: {e}")
        return False


def verify_schema():
    """Verify the new schema is correctly applied"""

    logger.info("\nVerifying new schema...")

    db = DatabaseManager()

    try:
        with db.get_session() as session:
            # Check if new columns exist by trying to select them
            result = session.execute(
                text("""
                    SELECT
                        total_debt,
                        total_cash,
                        net_debt,
                        operating_cashflow,
                        free_cashflow,
                        gross_margin,
                        ebitda,
                        total_revenue,
                        shares_outstanding,
                        schema_version
                    FROM fundamental_data
                    LIMIT 1
                """)
            ).fetchone()

            if result is not None:
                logger.success("✓ All new columns are accessible")
                return True
            else:
                logger.warning("No records found to verify")
                return True

    except Exception as e:
        logger.error(f"Schema verification failed: {e}")
        logger.info("This might mean the database needs to be recreated with the new schema")
        return False


if __name__ == "__main__":
    logger.info(f"Starting migration at {datetime.now()}")

    # Run migration
    success = migrate_to_v2()

    if success:
        # Verify schema
        verify_schema()

        logger.info("\n" + "=" * 70)
        logger.success("Migration completed successfully!")
        logger.info("=" * 70)
        logger.info("\nNext steps:")
        logger.info("  1. Run warren_scan.py to collect extended fundamentals")
        logger.info("  2. New scans will populate schema_version = 2 records")
        logger.info("  3. Old v1 records remain valid (backward compatible)")
    else:
        logger.error("\n" + "=" * 70)
        logger.error("Migration failed!")
        logger.error("=" * 70)
        logger.info("\nTroubleshooting:")
        logger.info("  1. Check database file exists and is writable")
        logger.info("  2. Verify no other processes are using the database")
        logger.info("  3. Consider backing up data and recreating database")