#!/usr/bin/env python3
"""
Apply v9c Subscription System Migration

This script checks if v9c migration is needed and applies it.
Run automatically on container startup via main.py startup event.
"""
import os
import psycopg2
from psycopg2.extras import RealDictCursor


def get_db_connection():
    """Get PostgreSQL connection."""
    return psycopg2.connect(
        host=os.getenv("DB_HOST", "postgres"),
        port=int(os.getenv("DB_PORT", 5432)),
        database=os.getenv("DB_NAME", "mitai_prod"),
        user=os.getenv("DB_USER", "mitai_prod"),
        password=os.getenv("DB_PASSWORD", ""),
        cursor_factory=RealDictCursor
    )


def migration_needed(conn):
    """Check if v9c migration is needed."""
    cur = conn.cursor()

    # Check if tiers table exists
    cur.execute("""
        SELECT EXISTS (
            SELECT FROM information_schema.tables
            WHERE table_name = 'tiers'
        )
    """)
    tiers_exists = cur.fetchone()['exists']

    # Check if features table exists
    cur.execute("""
        SELECT EXISTS (
            SELECT FROM information_schema.tables
            WHERE table_name = 'features'
        )
    """)
    features_exists = cur.fetchone()['exists']

    cur.close()

    # Migration needed if either table is missing
    return not (tiers_exists and features_exists)


def apply_migration():
    """Apply v9c migration if needed."""
    print("[v9c Migration] Checking if migration is needed...")

    try:
        conn = get_db_connection()

        if not migration_needed(conn):
            print("[v9c Migration] Already applied, skipping.")
            conn.close()
            return

        print("[v9c Migration] Applying subscription system migration...")

        # Read migration SQL
        migration_path = os.path.join(
            os.path.dirname(__file__),
            "migrations",
            "v9c_subscription_system.sql"
        )

        with open(migration_path, 'r', encoding='utf-8') as f:
            migration_sql = f.read()

        # Execute migration
        cur = conn.cursor()
        cur.execute(migration_sql)
        conn.commit()
        cur.close()
        conn.close()

        print("[v9c Migration] ✅ Migration completed successfully!")

        # Verify tables created
        conn = get_db_connection()
        cur = conn.cursor()
        cur.execute("""
            SELECT table_name FROM information_schema.tables
            WHERE table_schema = 'public'
            AND table_name IN ('tiers', 'features', 'tier_limits', 'access_grants', 'coupons')
            ORDER BY table_name
        """)
        tables = [r['table_name'] for r in cur.fetchall()]
        print(f"[v9c Migration] Created tables: {', '.join(tables)}")

        # Verify initial data
        cur.execute("SELECT COUNT(*) as count FROM tiers")
        tier_count = cur.fetchone()['count']
        cur.execute("SELECT COUNT(*) as count FROM features")
        feature_count = cur.fetchone()['count']
        cur.execute("SELECT COUNT(*) as count FROM tier_limits")
        limit_count = cur.fetchone()['count']

        print(f"[v9c Migration] Initial data: {tier_count} tiers, {feature_count} features, {limit_count} tier limits")

        cur.close()
        conn.close()

    except Exception as e:
        print(f"[v9c Migration] ❌ Error: {e}")
        raise


if __name__ == "__main__":
    apply_migration()