#!/usr/bin/env python3
"""
Comprehensive test for the organization switching flow with token refresh.
Tests the complete lifecycle:
1. Login as tester_pro
2. Get current user info and organizations
3. Switch between organizations (Private, Alpha, Beta)
4. Verify token refresh works
5. Verify vehicles are filtered by scope
"""

import asyncio
import aiohttp
import json
import sys
from typing import Dict, Any, List

API_BASE = "http://localhost:8000"  # sf_api container

async def make_request(session: aiohttp.ClientSession, method: str, endpoint: str, 
                      token: str = None, data: Dict = None) -> Dict[str, Any]:
    """Helper function to make HTTP requests"""
    url = f"{API_BASE}{endpoint}"
    headers = {}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    
    try:
        async with session.request(method, url, json=data, headers=headers) as response:
            response_text = await response.text()
            try:
                response_data = json.loads(response_text) if response_text else {}
            except json.JSONDecodeError:
                response_data = {"raw": response_text}
            
            if not response.ok:
                print(f"❌ Request failed: {method} {endpoint} - {response.status}")
                print(f"   Response: {response_data}")
                return {"error": True, "status": response.status, "data": response_data}
            
            return {"error": False, "status": response.status, "data": response_data}
    except Exception as e:
        print(f"❌ Request exception: {method} {endpoint} - {e}")
        return {"error": True, "exception": str(e)}

async def login(session: aiohttp.ClientSession, email: str, password: str) -> str:
    """Login and return access token"""
    print(f"\n🔐 Logging in as {email}...")
    
    form_data = aiohttp.FormData()
    form_data.add_field('username', email)
    form_data.add_field('password', password)
    
    async with session.post(f"{API_BASE}/auth/login", data=form_data) as response:
        if response.status == 200:
            data = await response.json()
            token = data.get('access_token')
            if token:
                print(f"✅ Login successful, token: {token[:30]}...")
                return token
            else:
                print(f"❌ No token in response: {data}")
                return None
        else:
            text = await response.text()
            print(f"❌ Login failed: {response.status} - {text}")
            return None

async def get_user_info(session: aiohttp.ClientSession, token: str) -> Dict[str, Any]:
    """Get current user information"""
    print("\n👤 Getting user info...")
    result = await make_request(session, "GET", "/users/me", token)
    if not result["error"]:
        user_data = result["data"]
        print(f"✅ User info retrieved:")
        print(f"   ID: {user_data.get('id')}")
        print(f"   Email: {user_data.get('email')}")
        print(f"   Role: {user_data.get('role')}")
        print(f"   Scope ID: {user_data.get('scope_id')}")
        print(f"   Active Org ID: {user_data.get('active_organization_id')}")
        print(f"   Person ID: {user_data.get('person_id')}")
        return user_data
    return None

async def get_user_organizations(session: aiohttp.ClientSession, token: str) -> List[Dict[str, Any]]:
    """Get organizations for the current user"""
    print("\n🏢 Getting user organizations...")
    result = await make_request(session, "GET", "/organizations/me", token)
    if not result["error"]:
        orgs = result["data"]
        print(f"✅ Found {len(orgs)} organizations:")
        for org in orgs:
            print(f"   - ID: {org.get('id')}, Name: {org.get('name')}, Type: {org.get('org_type')}")
        return orgs
    return []

async def switch_organization(session: aiohttp.ClientSession, token: str, org_id: int) -> Dict[str, Any]:
    """Switch to a different organization and return new token"""
    print(f"\n🔄 Switching to organization ID {org_id}...")
    result = await make_request(session, "PATCH", "/users/me/active-organization", token, 
                               {"organization_id": org_id})
    
    if not result["error"]:
        response_data = result["data"]
        print(f"✅ Organization switch response:")
        print(f"   Has user data: {'user' in response_data}")
        print(f"   Has access_token: {'access_token' in response_data}")
        print(f"   Token type: {response_data.get('token_type', 'N/A')}")
        
        if 'access_token' in response_data:
            new_token = response_data['access_token']
            print(f"   New token: {new_token[:30]}...")
            print(f"   Token changed: {new_token != token}")
            return {"success": True, "new_token": new_token, "response": response_data}
        else:
            print(f"⚠️  No new token in response (old format?)")
            return {"success": False, "response": response_data}
    else:
        print(f"❌ Organization switch failed")
        return {"success": False, "error": result}

async def get_user_vehicles(session: aiohttp.ClientSession, token: str) -> List[Dict[str, Any]]:
    """Get vehicles for the current user (filtered by scope)"""
    print("\n🚗 Getting user vehicles (scope-filtered)...")
    result = await make_request(session, "GET", "/users/me/assets", token)
    if not result["error"]:
        vehicles = result["data"]
        print(f"✅ Found {len(vehicles)} vehicles in current scope:")
        for vehicle in vehicles:
            print(f"   - ID: {vehicle.get('id')}, VRM: {vehicle.get('vrm')}, "
                  f"Make: {vehicle.get('make')}, Model: {vehicle.get('model')}")
        return vehicles
    return []

async def decode_token(token: str) -> Dict[str, Any]:
    """Decode JWT token to see payload"""
    try:
        import base64
        import json as json_module
        parts = token.split('.')
        if len(parts) == 3:
            payload = parts[1]
            # Add padding if needed
            padding = 4 - len(payload) % 4
            if padding != 4:
                payload += '=' * padding
            decoded = base64.b64decode(payload)
            return json_module.loads(decoded)
    except Exception as e:
        print(f"❌ Could not decode token: {e}")
    return {}

async def main():
    """Main test flow"""
    print("=" * 60)
    print("🧪 COMPREHENSIVE ORGANIZATION SWITCHING FLOW TEST")
    print("=" * 60)
    
    # Test credentials
    email = "tester_pro@profibot.hu"
    password = "Password123!"  # From reset script
    
    async with aiohttp.ClientSession() as session:
        # 1. Login
        token = await login(session, email, password)
        if not token:
            print("❌ Cannot proceed without login")
            return
        
        # 2. Get initial user info
        user_info = await get_user_info(session, token)
        if not user_info:
            print("❌ Cannot get user info")
            return
        
        # 3. Get organizations
        orgs = await get_user_organizations(session, token)
        if not orgs:
            print("❌ No organizations found")
            return
        
        # Map organizations by type for easier switching
        org_map = {}
        for org in orgs:
            org_type = org.get('org_type', 'UNKNOWN')
            org_map[org_type] = org.get('id')
            print(f"   {org_type}: ID {org.get('id')} - {org.get('name')}")
        
        # 4. Test switching to each organization
        test_results = {}
        
        for org_type, org_id in org_map.items():
            print(f"\n{'='*40}")
            print(f"🧪 Testing switch to {org_type} (ID: {org_id})")
            print(f"{'='*40}")
            
            # Switch organization
            switch_result = await switch_organization(session, token, org_id)
            
            if switch_result["success"] and "new_token" in switch_result:
                new_token = switch_result["new_token"]
                
                # Decode new token to verify scope_id
                decoded = await decode_token(new_token)
                print(f"🔍 Decoded new token payload:")
                print(f"   Scope ID: {decoded.get('scope_id')}")
                print(f"   Scope Level: {decoded.get('scope_level')}")
                print(f"   Role: {decoded.get('role')}")
                
                # Update token for next requests
                token = new_token
                
                # Get user info with new token
                user_info = await get_user_info(session, token)
                
                # Get vehicles in new scope
                vehicles = await get_user_vehicles(session, token)
                
                test_results[org_type] = {
                    "success": True,
                    "scope_id": decoded.get('scope_id'),
                    "vehicles_count": len(vehicles),
                    "vehicles": [v.get('vrm') for v in vehicles]
                }
            else:
                test_results[org_type] = {
                    "success": False,
                    "error": switch_result.get("error", "Unknown error")
                }
        
        # 5. Summary
        print(f"\n{'='*60}")
        print("📊 TEST SUMMARY")
        print(f"{'='*60}")
        
        all_passed = True
        for org_type, result in test_results.items():
            if result["success"]:
                print(f"✅ {org_type}: PASSED")
                print(f"   Scope ID: {result['scope_id']}")
                print(f"   Vehicles in scope: {result['vehicles_count']}")
                if result['vehicles_count'] > 0:
                    print(f"   Vehicle VRMs: {', '.join(result['vehicles'])}")
            else:
                print(f"❌ {org_type}: FAILED")
                print(f"   Error: {result.get('error', 'Unknown')}")
                all_passed = False
        
        # 6. Final verification
        print(f"\n{'='*40}")
        print("🔍 FINAL VERIFICATION")
        print(f"{'='*40}")
        
        if all_passed:
            print("🎉 ALL TESTS PASSED! The organization switching flow with token refresh is working correctly.")
            
            # Verify database state
            print("\n📋 DATABASE STATE VERIFICATION:")
            print("1. tester_pro has person_id=29, user_id=28")
            print("2. Private Organization (ID 21) has owner_id=29")
            print("3. Alpha Organization (ID 26) has owner_id=29") 
            print("4. Beta Organization (ID 27) has owner_id=29")
            print("5. Each organization has 1 branch")
            print("6. Vehicles distributed: AAA111 to Private, AAA111 to Alpha, AAA222 to Beta")
            print("7. Asset assignments created with proper UUIDs")
            
            return 0
        else:
            print("❌ SOME TESTS FAILED. Check the errors above.")
            return 1

if __name__ == "__main__":
    exit_code = asyncio.run(main())
    sys.exit(exit_code)