#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from sqlalchemy import create_engine, func
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
import os
from datetime import datetime
import time

# Get database connection details from environment variables
DB_HOST = os.environ.get('DB_HOST', 'postgres')
DB_PORT = os.environ.get('DB_PORT', '5432')
DB_NAME = os.environ.get('DB_NAME', 'evyos')
DB_USER = os.environ.get('DB_USER', 'evyos')
DB_PASS = os.environ.get('DB_PASS', 'evyos')

# Create SQLAlchemy engine and session
engine = create_engine(f'postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}')
Session = sessionmaker(bind=engine)
session = Session()

print("\n" + "=" * 50)
print("UPDATING ACCOUNTS WITH ZERO REMAINDER_BALANCE")
print("=" * 50)

try:
    # First, get the count of accounts that need updating
    count_query = """SELECT COUNT(*) FROM account_records 
                    WHERE build_parts_id IS NOT NULL 
                    AND currency_value > 0 
                    AND (remainder_balance = 0 OR remainder_balance IS NULL) 
                    AND bank_date >= '2022-01-01'"""
    count_result = session.execute(count_query).scalar()
    print(f"Found {count_result} accounts with zero remainder_balance")
    
    # Then update all those accounts
    update_query = """UPDATE account_records 
                    SET remainder_balance = currency_value 
                    WHERE build_parts_id IS NOT NULL 
                    AND currency_value > 0 
                    AND (remainder_balance = 0 OR remainder_balance IS NULL) 
                    AND bank_date >= '2022-01-01'"""
    result = session.execute(update_query)
    session.commit()
    print(f"Updated {result.rowcount} accounts with zero remainder_balance")
    
    # Verify the update
    verify_query = """SELECT COUNT(*) FROM account_records 
                     WHERE build_parts_id IS NOT NULL 
                     AND currency_value > 0 
                     AND (remainder_balance = 0 OR remainder_balance IS NULL) 
                     AND bank_date >= '2022-01-01'"""
    verify_result = session.execute(verify_query).scalar()
    print(f"Remaining accounts with zero remainder_balance: {verify_result}")
    
except Exception as e:
    print(f"Error updating accounts: {str(e)}")
    session.rollback()
finally:
    session.close()

print("=" * 50)