"""
Database Migration: Billing Enhancements
=========================================

This migration adds support for:
1. Custom billing cycles (recurring date ranges like Jan 1 - Dec 31)
2. Dynamic pricing with minimum/suggested amounts
3. Donation tracking separate from base subscription fees

Changes:
- SubscriptionPlan: Add 8 new fields for custom cycles and dynamic pricing
- Subscription: Add 2 new fields for donation tracking

Backward compatibility:
- Keeps existing fields (price_cents, stripe_price_id) for safe rollback
- Backfills new fields from existing data

Usage:
  python migrate_billing_enhancements.py
"""

import os
import sys
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker
from datetime import datetime

# Import database connection
from database import engine, SessionLocal

def run_migration():
    """Execute the migration to add billing enhancement fields."""

    print("=" * 60)
    print("LOAF Membership - Billing Enhancements Migration")
    print("=" * 60)
    print()

    # Create session
    db = SessionLocal()

    try:
        # =====================================================================
        # STEP 1: Add new columns to subscription_plans table
        # =====================================================================
        print("STEP 1: Adding new columns to subscription_plans table...")

        subscription_plan_migrations = [
            # Custom billing cycle fields
            "ALTER TABLE subscription_plans ADD COLUMN IF NOT EXISTS custom_cycle_enabled BOOLEAN DEFAULT FALSE",
            "ALTER TABLE subscription_plans ADD COLUMN IF NOT EXISTS custom_cycle_start_month INTEGER",
            "ALTER TABLE subscription_plans ADD COLUMN IF NOT EXISTS custom_cycle_start_day INTEGER",
            "ALTER TABLE subscription_plans ADD COLUMN IF NOT EXISTS custom_cycle_end_month INTEGER",
            "ALTER TABLE subscription_plans ADD COLUMN IF NOT EXISTS custom_cycle_end_day INTEGER",

            # Dynamic pricing fields
            "ALTER TABLE subscription_plans ADD COLUMN IF NOT EXISTS minimum_price_cents INTEGER DEFAULT 3000 NOT NULL",
            "ALTER TABLE subscription_plans ADD COLUMN IF NOT EXISTS suggested_price_cents INTEGER",
            "ALTER TABLE subscription_plans ADD COLUMN IF NOT EXISTS allow_donation BOOLEAN DEFAULT TRUE",
        ]

        for sql in subscription_plan_migrations:
            try:
                db.execute(text(sql))
                print(f"  ✓ {sql.split('ADD COLUMN IF NOT EXISTS')[1].split()[0]}")
            except Exception as e:
                print(f"  ⚠ {sql.split('ADD COLUMN IF NOT EXISTS')[1].split()[0]} - {str(e)}")

        db.commit()
        print("  ✓ Subscription plans table columns added successfully")
        print()

        # =====================================================================
        # STEP 2: Backfill subscription_plans data
        # =====================================================================
        print("STEP 2: Backfilling subscription_plans data...")

        backfill_plans_sql = """
        UPDATE subscription_plans
        SET
            minimum_price_cents = COALESCE(minimum_price_cents, price_cents),
            suggested_price_cents = COALESCE(suggested_price_cents, price_cents),
            custom_cycle_enabled = COALESCE(custom_cycle_enabled, FALSE),
            allow_donation = COALESCE(allow_donation, TRUE)
        WHERE minimum_price_cents IS NULL OR suggested_price_cents IS NULL
        """

        result = db.execute(text(backfill_plans_sql))
        db.commit()
        print(f"  ✓ Backfilled {result.rowcount} subscription plan records")
        print()

        # =====================================================================
        # STEP 3: Add new columns to subscriptions table
        # =====================================================================
        print("STEP 3: Adding new columns to subscriptions table...")

        subscription_migrations = [
            "ALTER TABLE subscriptions ADD COLUMN IF NOT EXISTS base_subscription_cents INTEGER",
            "ALTER TABLE subscriptions ADD COLUMN IF NOT EXISTS donation_cents INTEGER DEFAULT 0 NOT NULL",
        ]

        for sql in subscription_migrations:
            try:
                db.execute(text(sql))
                print(f"  ✓ {sql.split('ADD COLUMN IF NOT EXISTS')[1].split()[0]}")
            except Exception as e:
                print(f"  ⚠ {sql.split('ADD COLUMN IF NOT EXISTS')[1].split()[0]} - {str(e)}")

        db.commit()
        print("  ✓ Subscriptions table columns added successfully")
        print()

        # =====================================================================
        # STEP 4: Backfill subscriptions data
        # =====================================================================
        print("STEP 4: Backfilling subscriptions data...")

        backfill_subscriptions_sql = """
        UPDATE subscriptions
        SET
            base_subscription_cents = COALESCE(base_subscription_cents, amount_paid_cents, 0),
            donation_cents = COALESCE(donation_cents, 0)
        WHERE base_subscription_cents IS NULL
        """

        result = db.execute(text(backfill_subscriptions_sql))
        db.commit()
        print(f"  ✓ Backfilled {result.rowcount} subscription records")
        print()

        # =====================================================================
        # STEP 5: Verify migration
        # =====================================================================
        print("STEP 5: Verifying migration...")

        # Check subscription_plans columns
        check_plans_sql = """
        SELECT
            COUNT(*) as total,
            COUNT(minimum_price_cents) as has_minimum,
            COUNT(suggested_price_cents) as has_suggested,
            COUNT(custom_cycle_enabled) as has_custom_enabled
        FROM subscription_plans
        """

        result = db.execute(text(check_plans_sql)).fetchone()
        print(f"  Subscription Plans: {result[0]} total")
        print(f"    - {result[1]} have minimum_price_cents")
        print(f"    - {result[2]} have suggested_price_cents")
        print(f"    - {result[3]} have custom_cycle_enabled")

        # Check subscriptions columns
        check_subscriptions_sql = """
        SELECT
            COUNT(*) as total,
            COUNT(base_subscription_cents) as has_base,
            COUNT(donation_cents) as has_donation,
            SUM(base_subscription_cents) as total_base,
            SUM(donation_cents) as total_donations
        FROM subscriptions
        """

        result = db.execute(text(check_subscriptions_sql)).fetchone()
        print(f"  Subscriptions: {result[0]} total")
        print(f"    - {result[1]} have base_subscription_cents")
        print(f"    - {result[2]} have donation_cents")
        print(f"    - Total base: ${(result[3] or 0) / 100:.2f}")
        print(f"    - Total donations: ${(result[4] or 0) / 100:.2f}")
        print()

        # =====================================================================
        # STEP 6: Add constraints (optional validation)
        # =====================================================================
        print("STEP 6: Adding data validation constraints...")

        # Add check constraint for custom cycle months (1-12)
        try:
            db.execute(text("""
                ALTER TABLE subscription_plans
                ADD CONSTRAINT IF NOT EXISTS check_custom_cycle_months
                CHECK (
                    (custom_cycle_start_month IS NULL OR (custom_cycle_start_month >= 1 AND custom_cycle_start_month <= 12)) AND
                    (custom_cycle_end_month IS NULL OR (custom_cycle_end_month >= 1 AND custom_cycle_end_month <= 12))
                )
            """))
            print("  ✓ Added month validation constraint (1-12)")
        except Exception as e:
            print(f"  ⚠ Month constraint: {str(e)}")

        # Add check constraint for custom cycle days (1-31)
        try:
            db.execute(text("""
                ALTER TABLE subscription_plans
                ADD CONSTRAINT IF NOT EXISTS check_custom_cycle_days
                CHECK (
                    (custom_cycle_start_day IS NULL OR (custom_cycle_start_day >= 1 AND custom_cycle_start_day <= 31)) AND
                    (custom_cycle_end_day IS NULL OR (custom_cycle_end_day >= 1 AND custom_cycle_end_day <= 31))
                )
            """))
            print("  ✓ Added day validation constraint (1-31)")
        except Exception as e:
            print(f"  ⚠ Day constraint: {str(e)}")

        # Add check constraint for minimum price (>= $30)
        try:
            db.execute(text("""
                ALTER TABLE subscription_plans
                ADD CONSTRAINT IF NOT EXISTS check_minimum_price
                CHECK (minimum_price_cents >= 3000)
            """))
            print("  ✓ Added minimum price constraint ($30+)")
        except Exception as e:
            print(f"  ⚠ Minimum price constraint: {str(e)}")

        db.commit()
        print()

        # =====================================================================
        # MIGRATION COMPLETE
        # =====================================================================
        print("=" * 60)
        print("✓ MIGRATION COMPLETED SUCCESSFULLY")
        print("=" * 60)
        print()
        print("Summary:")
        print("  - Added 8 new columns to subscription_plans table")
        print("  - Added 2 new columns to subscriptions table")
        print("  - Backfilled all existing data")
        print("  - Added validation constraints")
        print()
        print("Next Steps:")
        print("  1. Update models.py with new column definitions")
        print("  2. Update server.py endpoints for custom billing")
        print("  3. Update frontend components for dynamic pricing")
        print()
        print("Rollback Notes:")
        print("  - Old fields (price_cents, stripe_price_id) are preserved")
        print("  - To rollback, revert code and ignore new columns")
        print("  - Database can coexist with old and new fields")
        print()

    except Exception as e:
        db.rollback()
        print()
        print("=" * 60)
        print("✗ MIGRATION FAILED")
        print("=" * 60)
        print(f"Error: {str(e)}")
        print()
        print("No changes have been committed to the database.")
        sys.exit(1)

    finally:
        db.close()

def rollback_migration():
    """Rollback the migration by removing new columns (use with caution)."""

    print("=" * 60)
    print("ROLLBACK: Billing Enhancements Migration")
    print("=" * 60)
    print()
    print("WARNING: This will remove all custom billing cycle and donation data!")
    confirm = input("Type 'ROLLBACK' to confirm: ")

    if confirm != "ROLLBACK":
        print("Rollback cancelled.")
        return

    db = SessionLocal()

    try:
        print("Removing columns from subscription_plans...")
        db.execute(text("ALTER TABLE subscription_plans DROP COLUMN IF EXISTS custom_cycle_enabled"))
        db.execute(text("ALTER TABLE subscription_plans DROP COLUMN IF EXISTS custom_cycle_start_month"))
        db.execute(text("ALTER TABLE subscription_plans DROP COLUMN IF EXISTS custom_cycle_start_day"))
        db.execute(text("ALTER TABLE subscription_plans DROP COLUMN IF EXISTS custom_cycle_end_month"))
        db.execute(text("ALTER TABLE subscription_plans DROP COLUMN IF EXISTS custom_cycle_end_day"))
        db.execute(text("ALTER TABLE subscription_plans DROP COLUMN IF EXISTS minimum_price_cents"))
        db.execute(text("ALTER TABLE subscription_plans DROP COLUMN IF EXISTS suggested_price_cents"))
        db.execute(text("ALTER TABLE subscription_plans DROP COLUMN IF EXISTS allow_donation"))

        print("Removing columns from subscriptions...")
        db.execute(text("ALTER TABLE subscriptions DROP COLUMN IF EXISTS base_subscription_cents"))
        db.execute(text("ALTER TABLE subscriptions DROP COLUMN IF EXISTS donation_cents"))

        db.commit()
        print("✓ Rollback completed successfully")

    except Exception as e:
        db.rollback()
        print(f"✗ Rollback failed: {str(e)}")
        sys.exit(1)

    finally:
        db.close()

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Billing Enhancements Migration")
    parser.add_argument(
        "--rollback",
        action="store_true",
        help="Rollback the migration (removes new columns)"
    )

    args = parser.parse_args()

    if args.rollback:
        rollback_migration()
    else:
        run_migration()