"""
Test script for PNL correction functionality
Creates a test EXPIRED trade with transactions and verifies the correction
"""

from optrabot.database import SessionLocal
from optrabot import models, crud
from datetime import datetime, date
import pytz

def create_test_expired_trade():
    """Create a test EXPIRED trade with transactions but wrong PNL"""
    with SessionLocal() as session:
        # Create test trade
        trade = models.Trade(
            account='TEST_ACCOUNT',
            symbol='SPX',
            strategy='TEST_STRATEGY',
            status='EXPIRED',
            realizedPNL=0.0  # Wrong PNL!
        )
        session.add(trade)
        session.commit()
        session.refresh(trade)
        
        # Add transactions that should result in positive PNL
        # Sell to open: Credit $100 (negative price for credit)
        tx1 = models.Transaction(
            tradeid=trade.id,
            id=1,
            type='SELL',
            sectype='C',
            timestamp=datetime.now(pytz.UTC),
            expiration=date(2025, 10, 23),
            strike=5800.0,
            contracts=1,
            price=-1.00,  # Negative for credit received
            fee=0.65,
            commission=0.0,
            notes='Test opening transaction'
        )
        
        # Expired worthless: No cost
        tx2 = models.Transaction(
            tradeid=trade.id,
            id=2,
            type='EXP',
            sectype='C',
            timestamp=datetime.now(pytz.UTC),
            expiration=date(2025, 10, 23),
            strike=5800.0,
            contracts=1,
            price=0.00,  # Expired worthless
            fee=0.0,
            commission=0.0,
            notes='Test expiration transaction'
        )
        
        session.add(tx1)
        session.add(tx2)
        session.commit()
        
        print(f'Created test trade {trade.id} with status=EXPIRED and realizedPNL=0.0')
        print(f'Expected PNL after correction: $99.35 (100 - 0.65 fee)')
        return trade.id

def verify_correction(trade_id):
    """Verify that the PNL was corrected"""
    with SessionLocal() as session:
        from sqlalchemy import select
        statement = select(models.Trade).filter_by(id=trade_id)
        trade = session.scalars(statement).first()
        
        if trade:
            print(f'\nTrade {trade_id} after correction:')
            print(f'  Status: {trade.status}')
            print(f'  RealizedPNL: ${trade.realizedPNL}')
            print(f'  Transactions: {len(trade.transactions)}')
            
            if trade.realizedPNL == 99.35:
                print('✅ PNL correction successful!')
            else:
                print(f'❌ PNL incorrect! Expected $99.35, got ${trade.realizedPNL}')
        else:
            print(f'❌ Trade {trade_id} not found')

def cleanup_test_trade(trade_id):
    """Remove test trade"""
    with SessionLocal() as session:
        from sqlalchemy import text
        session.execute(text(f'DELETE FROM transactions WHERE tradeid = {trade_id}'))
        session.execute(text(f'DELETE FROM trades WHERE id = {trade_id}'))
        session.commit()
        print(f'\nTest trade {trade_id} cleaned up')

if __name__ == '__main__':
    print('=== PNL Correction Test ===\n')
    
    # Create test trade
    trade_id = create_test_expired_trade()
    
    # Run the correction
    print('\nRunning updateDatabase() to trigger correction...')
    from optrabot.database import updateDatabase
    updateDatabase()
    
    # Verify
    verify_correction(trade_id)
    
    # Cleanup
    cleanup_test_trade(trade_id)
    
    # Clear the flag for testing again
    with SessionLocal() as session:
        crud.delete_setting(session, 'pnl_correction_expired_trades_done')
        print('Correction flag cleared for next test')
