import asyncio
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
from sqlalchemy import text, select, and_
import os
from dotenv import load_dotenv

load_dotenv()
DATABASE_URL = os.getenv("DATABASE_URL")

engine = create_async_engine(DATABASE_URL)
async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

async def migrate():
    from app.models.stock_ledger import StockLedger
    from app.models.transaction import Transaction, TransactionItem
    from app.models.patti_transaction import PattiTransaction, PattiTransactionItem

    async with engine.connect() as conn:
        print("Adding poti columns to stock_ledger...")
        await conn.execute(text("ALTER TABLE stock_ledger ADD COLUMN IF NOT EXISTS poti_in INTEGER DEFAULT 0"))
        await conn.execute(text("ALTER TABLE stock_ledger ADD COLUMN IF NOT EXISTS poti_out INTEGER DEFAULT 0"))
        await conn.execute(text("ALTER TABLE stock_ledger ADD COLUMN IF NOT EXISTS poti_balance_after INTEGER DEFAULT 0"))
        await conn.commit()

    async with async_session() as session:
        print("Populating poti data and calculating bag balances...")
        
        # We need to recalculate balances per (mirchi_name, lot_number)
        # First, get all entries ordered by creation
        res = await session.execute(select(StockLedger).order_by(StockLedger.created_at.asc()))
        entries = res.scalars().all()
        
        balances = {} # (mirchi_name, lot_number) -> current_poti_balance
        
        for entry in entries:
            poti_count = 0
            # Find the item to get the poti_count
            if entry.transaction_id:
                item_res = await session.execute(
                    select(TransactionItem).where(
                        and_(
                            TransactionItem.transaction_id == entry.transaction_id,
                            TransactionItem.mirchi_name == entry.mirchi_name
                        )
                    )
                )
                item = item_res.scalar_one_or_none()
                if item:
                    poti_count = item.poti_count
            elif entry.patti_transaction_id:
                item_res = await session.execute(
                    select(PattiTransactionItem).where(
                        and_(
                            PattiTransactionItem.transaction_id == entry.patti_transaction_id,
                            PattiTransactionItem.mirchi_name == entry.mirchi_name
                        )
                    )
                )
                item = item_res.scalar_one_or_none()
                if item:
                    poti_count = item.poti_count
            
            # Update entry with poti_in/poti_out
            if entry.movement_type in ['IN', 'RETURN_IN']:
                entry.poti_in = poti_count
                entry.poti_out = 0
            else:
                entry.poti_in = 0
                entry.poti_out = poti_count
            
            # Calculate running balance
            key = (entry.mirchi_name, entry.lot_number)
            current_bal = balances.get(key, 0)
            if entry.movement_type in ['IN', 'RETURN_IN']:
                current_bal += poti_count
            else:
                current_bal -= poti_count
            
            entry.poti_balance_after = current_bal
            balances[key] = current_bal
            
        await session.commit()
        print("Migration and poti balance recalculation completed!")

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