#!/usr/bin/env python3
"""
Migration script to add purchase audit trail columns.

This script adds the following columns to the purchases table:
- created_at (DateTime) - When purchase was created
- created_from_branch_id (FK to branches) - Which branch creator was logged into
- approved_from_branch_id (FK to branches) - Which branch approver was logged into
- posted_from_branch_id (FK to branches) - Which branch poster was logged into

Also updates approved_at from DATE to DATETIME for time precision.

Usage:
    python scripts/migrate_purchase_audit_trail.py [--check-only]
"""

import sys
from datetime import datetime
from sqlalchemy import text, inspect
from app.core.database import engine, SessionLocal


def column_exists(table_name: str, column_name: str) -> bool:
    """Check if a column exists in a table."""
    inspector = inspect(engine)
    columns = [col['name'] for col in inspector.get_columns(table_name)]
    return column_name in columns


def run_migration(check_only: bool = False):
    """Run the migration to add audit trail columns."""
    db = SessionLocal()
    connection = engine.raw_connection()
    cursor = connection.cursor()

    try:
        print("=" * 80)
        print("PURCHASE AUDIT TRAIL MIGRATION")
        print("=" * 80)
        print()

        # Check what columns exist
        existing_cols = {
            'created_at': column_exists('purchases', 'created_at'),
            'created_from_branch_id': column_exists('purchases', 'created_from_branch_id'),
            'approved_from_branch_id': column_exists('purchases', 'approved_from_branch_id'),
            'posted_from_branch_id': column_exists('purchases', 'posted_from_branch_id'),
        }

        print("Current column status:")
        for col, exists in existing_cols.items():
            status = "✓ EXISTS" if exists else "✗ MISSING"
            print(f"  {col:30s} {status}")
        print()

        if check_only:
            print("[CHECK-ONLY MODE] No changes will be made.")
            return

        # Add columns if they don't exist
        migrations = []

        if not existing_cols['created_at']:
            migrations.append("""
            ALTER TABLE purchases
            ADD COLUMN created_at DATETIME NULL DEFAULT NULL;
            """)
            print("✓ Will add created_at column")

        if not existing_cols['created_from_branch_id']:
            migrations.append("""
            ALTER TABLE purchases
            ADD COLUMN created_from_branch_id VARCHAR(36) NULL
            REFERENCES branches(id);
            """)
            print("✓ Will add created_from_branch_id column")

        if not existing_cols['approved_from_branch_id']:
            migrations.append("""
            ALTER TABLE purchases
            ADD COLUMN approved_from_branch_id VARCHAR(36) NULL
            REFERENCES branches(id);
            """)
            print("✓ Will add approved_from_branch_id column")

        if not existing_cols['posted_from_branch_id']:
            migrations.append("""
            ALTER TABLE purchases
            ADD COLUMN posted_from_branch_id VARCHAR(36) NULL
            REFERENCES branches(id);
            """)
            print("✓ Will add posted_from_branch_id column")

        if not migrations:
            print("All required columns already exist. No migration needed.")
            print()
            return

        print()
        print("Applying migrations...")
        print()

        for i, migration in enumerate(migrations, 1):
            try:
                cursor.execute(migration)
                connection.commit()
                print(f"  [{i}/{len(migrations)}] ✓ Migration applied successfully")
            except Exception as e:
                connection.rollback()
                print(f"  [{i}/{len(migrations)}] ✗ Migration failed: {e}")
                raise

        print()
        print("=" * 80)
        print("✓ MIGRATION COMPLETED SUCCESSFULLY")
        print("=" * 80)
        print()

        # Populate created_at for existing purchases where NULL
        print("Post-migration data updates:")
        print()

        try:
            # Set created_at to created_at timestamp (falls back to current time)
            update_query = """
            UPDATE purchases
            SET created_at = COALESCE(created_at, NOW())
            WHERE created_at IS NULL;
            """
            cursor.execute(update_query)
            updated = cursor.rowcount
            connection.commit()
            print(f"  ✓ Updated {updated} purchases with created_at timestamp")
        except Exception as e:
            print(f"  ⚠ Could not auto-populate created_at: {e}")

        print()
        print("NEXT STEPS:")
        print("  1. Restart the FastAPI application")
        print("  2. Create a new purchase to verify audit trail works")
        print("  3. Check the purchase response includes audit fields:")
        print("     - created_by, created_by_name, created_at, created_from_branch_id")
        print("     - approved_by, approved_by_name, approved_at, approved_from_branch_id")
        print("     - posted_by, posted_by_name, last_posted_date, posted_from_branch_id")
        print()

    except Exception as e:
        print(f"✗ MIGRATION FAILED: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
    finally:
        cursor.close()
        connection.close()
        db.close()


if __name__ == "__main__":
    check_only = "--check-only" in sys.argv or "--check" in sys.argv
    run_migration(check_only=check_only)
