#!/usr/bin/env python3
"""
Database migration script to add document_path columns to FuelRecord and RecurringExpense tables.
This migration adds support for file attachments to fuel and tax (recurring expense) entries.
"""
import sqlite3
import os
import sys

def migrate_database():
    """Add document_path columns to FuelRecord and RecurringExpense tables."""
    # Support both Docker and local paths
    db_paths = [
        '/app/data/masina_dock.db',
        './data/masina_dock.db',
        '../data/masina_dock.db'
    ]
    
    db_path = None
    for path in db_paths:
        if os.path.exists(path):
            db_path = path
            break
    
    if not db_path:
        print("Database does not exist yet. Migration will be applied on first run.")
        return True
    
    print(f"Migrating database at: {db_path}")
    
    try:
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()
        
        # Check FuelRecord table
        cursor.execute("PRAGMA table_info(fuel_record)")
        fuel_columns = [column[1] for column in cursor.fetchall()]
        
        if 'document_path' not in fuel_columns:
            print("Adding document_path column to fuel_record table...")
            cursor.execute("ALTER TABLE fuel_record ADD COLUMN document_path VARCHAR(255)")
            print("✓ Added document_path to fuel_record")
        else:
            print("✓ fuel_record.document_path already exists")
        
        # Check RecurringExpense table
        cursor.execute("PRAGMA table_info(recurring_expense)")
        expense_columns = [column[1] for column in cursor.fetchall()]
        
        if 'document_path' not in expense_columns:
            print("Adding document_path column to recurring_expense table...")
            cursor.execute("ALTER TABLE recurring_expense ADD COLUMN document_path VARCHAR(255)")
            print("✓ Added document_path to recurring_expense")
        else:
            print("✓ recurring_expense.document_path already exists")
        
        conn.commit()
        print("\n✓ Database migration completed successfully!")
        return True
        
    except sqlite3.Error as e:
        print(f"✗ Migration error: {e}", file=sys.stderr)
        if conn:
            conn.rollback()
        return False
    except Exception as e:
        print(f"✗ Unexpected error: {e}", file=sys.stderr)
        return False
    finally:
        if conn:
            conn.close()

if __name__ == '__main__':
    success = migrate_database()
    sys.exit(0 if success else 1)