kitchen/backend/migrations/add_invoice_features.py

"""
Migration script to add new invoice features:
- document_type, order_number columns on invoices
- duplicate_status, duplicate_of_id, related_document_id columns on invoices
- line_items table

Run this script once after deploying the new code.
"""
import asyncio
import logging
from sqlalchemy import text
from database import engine

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


async def run_migration():
    """Add new columns and tables for invoice features."""
    # Always run all migrations - each has its own try/except to handle existing columns
    logger.info("Running database migrations...")

    # Run each migration in its own transaction
    migrations = [
        "ALTER TABLE invoices ADD COLUMN IF NOT EXISTS document_type VARCHAR(50) DEFAULT 'invoice'",
        "ALTER TABLE invoices ADD COLUMN IF NOT EXISTS order_number VARCHAR(100)",
        "ALTER TABLE invoices ADD COLUMN IF NOT EXISTS duplicate_status VARCHAR(50)",
        "ALTER TABLE invoices ADD COLUMN IF NOT EXISTS duplicate_of_id INTEGER REFERENCES invoices(id)",
        "ALTER TABLE invoices ADD COLUMN IF NOT EXISTS related_document_id INTEGER REFERENCES invoices(id)",
    ]

    for sql in migrations:
        try:
            async with engine.begin() as conn:
                await conn.execute(text(sql))
                logger.info(f"Executed: {sql[:60]}...")
        except Exception as e:
            if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
                logger.info(f"Column already exists, skipping")
            else:
                logger.warning(f"Migration warning: {e}")

    # Create line_items table
    create_line_items = """
    CREATE TABLE IF NOT EXISTS line_items (
        id SERIAL PRIMARY KEY,
        invoice_id INTEGER NOT NULL REFERENCES invoices(id) ON DELETE CASCADE,
        description TEXT,
        quantity NUMERIC(10, 3),
        unit_price NUMERIC(10, 2),
        amount NUMERIC(10, 2),
        product_code VARCHAR(100),
        line_number INTEGER DEFAULT 0,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    """
    try:
        async with engine.begin() as conn:
            await conn.execute(text(create_line_items))
            logger.info("Created line_items table")
    except Exception as e:
        logger.warning(f"line_items table: {e}")

    # Create index
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "CREATE INDEX IF NOT EXISTS ix_line_items_invoice_id ON line_items(invoice_id)"
            ))
            logger.info("Created index on line_items.invoice_id")
    except Exception as e:
        logger.warning(f"Index: {e}")

    # Add aliases column to suppliers table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE suppliers ADD COLUMN IF NOT EXISTS aliases JSON DEFAULT '[]'"
            ))
            logger.info("Added aliases column to suppliers")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("aliases column already exists")
        else:
            logger.warning(f"aliases column: {e}")

    # Add net_total column to invoices table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE invoices ADD COLUMN IF NOT EXISTS net_total NUMERIC(10, 2)"
            ))
            logger.info("Added net_total column to invoices")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("net_total column already exists")
        else:
            logger.warning(f"net_total column: {e}")

    # Add is_non_stock column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS is_non_stock BOOLEAN DEFAULT FALSE"
            ))
            logger.info("Added is_non_stock column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("is_non_stock column already exists")
        else:
            logger.warning(f"is_non_stock column: {e}")

    # Add vendor_name column to invoices table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE invoices ADD COLUMN IF NOT EXISTS vendor_name VARCHAR(255)"
            ))
            logger.info("Added vendor_name column to invoices")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("vendor_name column already exists")
        else:
            logger.warning(f"vendor_name column: {e}")

    # Add ocr_raw_json column to invoices table (for storing full Azure response)
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE invoices ADD COLUMN IF NOT EXISTS ocr_raw_json TEXT"
            ))
            logger.info("Added ocr_raw_json column to invoices")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("ocr_raw_json column already exists")
        else:
            logger.warning(f"ocr_raw_json column: {e}")

    # Add supplier_match_type column to invoices table (for fuzzy matching)
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE invoices ADD COLUMN IF NOT EXISTS supplier_match_type VARCHAR(20)"
            ))
            logger.info("Added supplier_match_type column to invoices")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("supplier_match_type column already exists")
        else:
            logger.warning(f"supplier_match_type column: {e}")

    # Create field_mappings table
    create_field_mappings = """
    CREATE TABLE IF NOT EXISTS field_mappings (
        id SERIAL PRIMARY KEY,
        kitchen_id INTEGER NOT NULL REFERENCES kitchens(id),
        supplier_id INTEGER REFERENCES suppliers(id),
        source_field VARCHAR(100) NOT NULL,
        target_field VARCHAR(100) NOT NULL,
        field_type VARCHAR(20) DEFAULT 'invoice',
        transform VARCHAR(50) DEFAULT 'direct',
        priority INTEGER DEFAULT 0,
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
    )
    """
    try:
        async with engine.begin() as conn:
            await conn.execute(text(create_field_mappings))
            logger.info("Created field_mappings table")
    except Exception as e:
        logger.warning(f"field_mappings table: {e}")

    # Create index on field_mappings
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "CREATE INDEX IF NOT EXISTS ix_field_mappings_kitchen_id ON field_mappings(kitchen_id)"
            ))
            await conn.execute(text(
                "CREATE INDEX IF NOT EXISTS ix_field_mappings_supplier_id ON field_mappings(supplier_id)"
            ))
            logger.info("Created indexes on field_mappings")
    except Exception as e:
        logger.warning(f"field_mappings indexes: {e}")

    # Add unit column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS unit VARCHAR(50)"
            ))
            logger.info("Added unit column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("unit column already exists")
        else:
            logger.warning(f"unit column: {e}")

    # Add order_quantity column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS order_quantity NUMERIC(10, 3)"
            ))
            logger.info("Added order_quantity column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("order_quantity column already exists")
        else:
            logger.warning(f"order_quantity column: {e}")

    # Add tax_rate column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS tax_rate VARCHAR(50)"
            ))
            logger.info("Added tax_rate column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("tax_rate column already exists")
        else:
            logger.warning(f"tax_rate column: {e}")

    # Add tax_amount column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS tax_amount NUMERIC(10, 2)"
            ))
            logger.info("Added tax_amount column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("tax_amount column already exists")
        else:
            logger.warning(f"tax_amount column: {e}")

    # Add raw_content column to line_items table (for pack size parsing)
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS raw_content TEXT"
            ))
            logger.info("Added raw_content column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("raw_content column already exists")
        else:
            logger.warning(f"raw_content column: {e}")

    # Add pack_quantity column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS pack_quantity INTEGER"
            ))
            logger.info("Added pack_quantity column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("pack_quantity column already exists")
        else:
            logger.warning(f"pack_quantity column: {e}")

    # Add unit_size column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS unit_size NUMERIC(10, 3)"
            ))
            logger.info("Added unit_size column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("unit_size column already exists")
        else:
            logger.warning(f"unit_size column: {e}")

    # Add unit_size_type column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS unit_size_type VARCHAR(10)"
            ))
            logger.info("Added unit_size_type column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("unit_size_type column already exists")
        else:
            logger.warning(f"unit_size_type column: {e}")

    # Add portions_per_unit column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS portions_per_unit INTEGER DEFAULT 1"
            ))
            logger.info("Added portions_per_unit column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("portions_per_unit column already exists")
        else:
            logger.warning(f"portions_per_unit column: {e}")

    # Add cost_per_item column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS cost_per_item NUMERIC(10, 4)"
            ))
            logger.info("Added cost_per_item column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("cost_per_item column already exists")
        else:
            logger.warning(f"cost_per_item column: {e}")

    # Add cost_per_portion column to line_items table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ADD COLUMN IF NOT EXISTS cost_per_portion NUMERIC(10, 4)"
            ))
            logger.info("Added cost_per_portion column to line_items")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("cost_per_portion column already exists")
        else:
            logger.warning(f"cost_per_portion column: {e}")

    # Create product_definitions table for persistent portion/pack data
    create_product_definitions = """
    CREATE TABLE IF NOT EXISTS product_definitions (
        id SERIAL PRIMARY KEY,
        kitchen_id INTEGER NOT NULL REFERENCES kitchens(id),
        supplier_id INTEGER REFERENCES suppliers(id),
        product_code VARCHAR(100),
        description_pattern VARCHAR(255),
        pack_quantity INTEGER,
        unit_size NUMERIC(10, 3),
        unit_size_type VARCHAR(10),
        portions_per_unit INTEGER,
        portion_description VARCHAR(100),
        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
        UNIQUE(kitchen_id, supplier_id, product_code)
    )
    """
    try:
        async with engine.begin() as conn:
            await conn.execute(text(create_product_definitions))
            logger.info("Created product_definitions table")
    except Exception as e:
        logger.warning(f"product_definitions table: {e}")

    # Create indexes on product_definitions
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "CREATE INDEX IF NOT EXISTS ix_product_definitions_kitchen_id ON product_definitions(kitchen_id)"
            ))
            await conn.execute(text(
                "CREATE INDEX IF NOT EXISTS ix_product_definitions_supplier_id ON product_definitions(supplier_id)"
            ))
            await conn.execute(text(
                "CREATE INDEX IF NOT EXISTS ix_product_definitions_product_code ON product_definitions(product_code)"
            ))
            logger.info("Created indexes on product_definitions")
    except Exception as e:
        logger.warning(f"product_definitions indexes: {e}")

    # Remove default from portions_per_unit column (make it truly nullable)
    # PostgreSQL doesn't have a simple ALTER COLUMN DROP DEFAULT, but we can change the default to NULL
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE line_items ALTER COLUMN portions_per_unit DROP DEFAULT"
            ))
            logger.info("Removed default from portions_per_unit column")
    except Exception as e:
        logger.warning(f"portions_per_unit default removal: {e}")

    # Add saved_by_user_id column to product_definitions table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE product_definitions ADD COLUMN IF NOT EXISTS saved_by_user_id INTEGER REFERENCES users(id)"
            ))
            logger.info("Added saved_by_user_id column to product_definitions")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("saved_by_user_id column already exists")
        else:
            logger.warning(f"saved_by_user_id column: {e}")

    # Add source_invoice_id column to product_definitions table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE product_definitions ADD COLUMN IF NOT EXISTS source_invoice_id INTEGER REFERENCES invoices(id) ON DELETE SET NULL"
            ))
            logger.info("Added source_invoice_id column to product_definitions")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("source_invoice_id column already exists")
        else:
            logger.warning(f"source_invoice_id column: {e}")

    # Add source_invoice_number column to product_definitions table
    try:
        async with engine.begin() as conn:
            await conn.execute(text(
                "ALTER TABLE product_definitions ADD COLUMN IF NOT EXISTS source_invoice_number VARCHAR(100)"
            ))
            logger.info("Added source_invoice_number column to product_definitions")
    except Exception as e:
        if "already exists" in str(e).lower() or "duplicate column" in str(e).lower():
            logger.info("source_invoice_number column already exists")
        else:
            logger.warning(f"source_invoice_number column: {e}")

    logger.info("Migration completed successfully!")


if __name__ == "__main__":
    asyncio.run(run_migration())