#!/usr/bin/env python3
"""
Test script to verify branch filtering fixes for Reports & Statements, 
Payment Receipts, and Notifications pages, plus admin update functionality.
"""

import os
import sys
import django
from datetime import datetime, timedelta
from decimal import Decimal

# Setup Django environment
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'branch_system.settings')
django.setup()

from django.contrib.auth import get_user_model
from django.test import Client, RequestFactory
from django.contrib.sessions.middleware import SessionMiddleware
from django.contrib.auth.middleware import AuthenticationMiddleware
from django.contrib.messages.middleware import MessageMiddleware
from users.models import Branch, CustomUser
from loans.models import Loan, Repayment
from utils.models import Notification
from utils.views import payment_receipts, notifications
from reports.views import reports_dashboard
from reports.comprehensive_reports import reports_service

User = get_user_model()

def log_message(message, level='INFO'):
    """Log messages with timestamp"""
    timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    print(f"[{timestamp}] {level}: {message}")

def setup_test_data():
    """Setup test data for branch filtering tests"""
    log_message("Setting up test data...")
    
    # Create test branches
    branch1, created = Branch.objects.get_or_create(
        name="Main Branch",
        defaults={
            'is_main_branch': True,
            'is_active': True,
            'address': 'Main Street 123'
        }
    )
    
    branch2, created = Branch.objects.get_or_create(
        name="Secondary Branch", 
        defaults={
            'is_main_branch': False,
            'is_active': True,
            'address': 'Second Street 456'
        }
    )
    
    # Get or create test users for each branch
    try:
        admin_user = CustomUser.objects.get(email='admin@test.com')
    except CustomUser.DoesNotExist:
        admin_user = CustomUser.objects.create(
            email='admin@test.com',
            first_name='Admin',
            last_name='User',
            role='admin',
            is_staff=True,
            is_superuser=True,
            branch=branch1,
            status='active',
            phone_number='+254700000001'
        )
        admin_user.set_password('testpass123')
        admin_user.save()
    
    # Create borrowers for each branch
    try:
        borrower1 = CustomUser.objects.get(email='borrower1@test.com')
    except CustomUser.DoesNotExist:
        borrower1 = CustomUser.objects.create(
            email='borrower1@test.com',
            first_name='John',
            last_name='Doe',
            role='borrower',
            branch=branch1,
            status='active',
            phone_number='+254700000002'
        )
    
    try:
        borrower2 = CustomUser.objects.get(email='borrower2@test.com')
    except CustomUser.DoesNotExist:
        borrower2 = CustomUser.objects.create(
            email='borrower2@test.com',
            first_name='Jane',
            last_name='Smith',
            role='borrower',
            branch=branch2,
            status='active',
            phone_number='+254700000003'
        )
    
    log_message(f"✅ Created/verified branches: {branch1.name}, {branch2.name}")
    log_message(f"✅ Created/verified users: {admin_user.email}, {borrower1.email}, {borrower2.email}")
    
    return {
        'branches': [branch1, branch2],
        'admin_user': admin_user,
        'borrowers': [borrower1, borrower2]
    }

def test_payment_receipts_filtering(test_data):
    """Test payment receipts branch filtering"""
    log_message("Testing Payment Receipts branch filtering...")
    
    try:
        factory = RequestFactory()
        admin_user = test_data['admin_user']
        branch1, branch2 = test_data['branches']
        
        # Create test request with branch selection
        request = factory.get('/utils/payment-receipts/')
        request.user = admin_user
        
        # Add session middleware
        middleware = SessionMiddleware(lambda req: None)
        middleware.process_request(request)
        request.session.save()
        
        # Test with branch1 selected
        request.session['selected_branch_id'] = str(branch1.id)
        
        # Call the view
        response = payment_receipts(request)
        
        if response.status_code == 200:
            log_message("✅ Payment Receipts view executed successfully with branch filtering")
        else:
            log_message(f"❌ Payment Receipts view failed with status {response.status_code}", 'ERROR')
            
        # Test with branch2 selected
        request.session['selected_branch_id'] = str(branch2.id)
        response = payment_receipts(request)
        
        if response.status_code == 200:
            log_message("✅ Payment Receipts view works with different branch selection")
        else:
            log_message(f"❌ Payment Receipts view failed with branch2", 'ERROR')
            
        return True
        
    except Exception as e:
        log_message(f"❌ Error testing Payment Receipts: {e}", 'ERROR')
        return False

def test_notifications_filtering(test_data):
    """Test notifications branch filtering"""
    log_message("Testing Notifications branch filtering...")
    
    try:
        factory = RequestFactory()
        admin_user = test_data['admin_user']
        branch1, branch2 = test_data['branches']
        
        # Create test notifications for different branches
        borrower1, borrower2 = test_data['borrowers']
        
        # Create notifications for each borrower
        notif1 = Notification.objects.create(
            user=borrower1,
            title="Test Notification Branch 1",
            message="This is for branch 1",
            notification_type="info"
        )
        
        notif2 = Notification.objects.create(
            user=borrower2,
            title="Test Notification Branch 2", 
            message="This is for branch 2",
            notification_type="info"
        )
        
        # Create test request
        request = factory.get('/utils/notifications/')
        request.user = admin_user
        
        # Add session middleware
        middleware = SessionMiddleware(lambda req: None)
        middleware.process_request(request)
        request.session.save()
        
        # Test with branch1 selected
        request.session['selected_branch_id'] = str(branch1.id)
        
        response = notifications(request)
        
        if response.status_code == 200:
            log_message("✅ Notifications view executed successfully with branch filtering")
        else:
            log_message(f"❌ Notifications view failed with status {response.status_code}", 'ERROR')
            
        return True
        
    except Exception as e:
        log_message(f"❌ Error testing Notifications: {e}", 'ERROR')
        return False

def test_reports_dashboard_filtering(test_data):
    """Test reports dashboard branch filtering"""
    log_message("Testing Reports & Statements Dashboard branch filtering...")
    
    try:
        factory = RequestFactory()
        admin_user = test_data['admin_user']
        branch1, branch2 = test_data['branches']
        
        # Create test request
        request = factory.get('/reports/dashboard/')
        request.user = admin_user
        
        # Add session middleware
        middleware = SessionMiddleware(lambda req: None)
        middleware.process_request(request)
        request.session.save()
        
        # Test with branch1 selected
        request.session['selected_branch_id'] = str(branch1.id)
        
        response = reports_dashboard(request)
        
        if response.status_code == 200:
            log_message("✅ Reports Dashboard view executed successfully with branch filtering")
        else:
            log_message(f"❌ Reports Dashboard view failed with status {response.status_code}", 'ERROR')
            
        # Test the reports service directly with branch filtering
        dashboard_data = reports_service.generate_comprehensive_dashboard_data(branch_id=branch1.id)
        
        if dashboard_data and 'generated_at' in dashboard_data:
            log_message("✅ Reports service supports branch filtering")
        else:
            log_message("❌ Reports service branch filtering failed", 'ERROR')
            
        return True
        
    except Exception as e:
        log_message(f"❌ Error testing Reports Dashboard: {e}", 'ERROR')
        return False

def test_admin_update_functionality(test_data):
    """Test admin update functionality including password changes"""
    log_message("Testing Admin Update functionality...")
    
    try:
        from users.views import admin_update
        
        factory = RequestFactory()
        admin_user = test_data['admin_user']
        
        # Create a test staff user to update
        try:
            staff_user = CustomUser.objects.get(email='staff@test.com')
        except CustomUser.DoesNotExist:
            staff_user = CustomUser.objects.create(
                email='staff@test.com',
                first_name='Staff',
                last_name='User',
                role='staff',
                is_staff=True,
                branch=test_data['branches'][0],
                status='active',
                phone_number='+254700000004'
            )
            staff_user.set_password('oldpass123')
            staff_user.save()
        
        # Test GET request (view form)
        request = factory.get(f'/users/admins/{staff_user.id}/update/')
        request.user = admin_user
        
        # Add session and messages middleware
        middleware = SessionMiddleware(lambda req: None)
        middleware.process_request(request)
        request.session.save()
        
        response = admin_update(request, staff_user.id)
        
        if response.status_code == 200:
            log_message("✅ Admin Update GET request successful")
        else:
            log_message(f"❌ Admin Update GET failed with status {response.status_code}", 'ERROR')
        
        # Test POST request (update user)
        post_data = {
            'first_name': 'Updated Staff',
            'last_name': 'User',
            'email': 'staff@test.com',
            'phone_number': '+254700000004',
            'status': 'active',
            'branch': str(test_data['branches'][1].id),  # Change branch
            'is_staff': 'on',
            'new_password': 'newpass123',
            'confirm_password': 'newpass123'
        }
        
        request = factory.post(f'/users/admins/{staff_user.id}/update/', post_data)
        request.user = admin_user
        
        # Add middleware
        middleware = SessionMiddleware(lambda req: None)
        middleware.process_request(request)
        request.session.save()
        
        # Add messages framework
        from django.contrib.messages.storage.fallback import FallbackStorage
        setattr(request, '_messages', FallbackStorage(request))
        
        response = admin_update(request, staff_user.id)
        
        # Refresh user from database
        staff_user.refresh_from_db()
        
        # Check if updates were applied
        if staff_user.first_name == 'Updated Staff':
            log_message("✅ Admin basic info update successful")
        else:
            log_message("❌ Admin basic info update failed", 'ERROR')
            
        if staff_user.branch_id == test_data['branches'][1].id:
            log_message("✅ Admin branch update successful")
        else:
            log_message("❌ Admin branch update failed", 'ERROR')
            
        # Test password change
        if staff_user.check_password('newpass123'):
            log_message("✅ Admin password update successful")
        else:
            log_message("❌ Admin password update failed", 'ERROR')
            
        return True
        
    except Exception as e:
        log_message(f"❌ Error testing Admin Update: {e}", 'ERROR')
        return False

def run_comprehensive_test():
    """Run all tests"""
    log_message("Starting comprehensive branch filtering and admin update tests...")
    
    # Setup test data
    test_data = setup_test_data()
    
    # Run tests
    tests = [
        ("Payment Receipts Filtering", test_payment_receipts_filtering),
        ("Notifications Filtering", test_notifications_filtering),
        ("Reports Dashboard Filtering", test_reports_dashboard_filtering),
        ("Admin Update Functionality", test_admin_update_functionality)
    ]
    
    results = {}
    
    for test_name, test_func in tests:
        log_message(f"\n--- Running {test_name} Test ---")
        try:
            results[test_name] = test_func(test_data)
        except Exception as e:
            log_message(f"❌ {test_name} test failed with exception: {e}", 'ERROR')
            results[test_name] = False
    
    # Summary
    log_message("\n" + "="*50)
    log_message("TEST RESULTS SUMMARY")
    log_message("="*50)
    
    passed = 0
    total = len(results)
    
    for test_name, result in results.items():
        status = "✅ PASSED" if result else "❌ FAILED"
        log_message(f"{test_name}: {status}")
        if result:
            passed += 1
    
    log_message(f"\nOverall: {passed}/{total} tests passed")
    
    if passed == total:
        log_message("🎉 All tests passed! Branch filtering and admin update fixes are working correctly.")
    else:
        log_message("⚠️  Some tests failed. Please review the errors above.")
    
    return passed == total

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