#!/usr/bin/env python3
"""
Real functionality test for upload WITH authentication

Companion test to test_simple_upload_no_auth.py that validates the complete
upload pipeline when proper authentication is provided.
"""

import requests
import json
import tempfile
import zipfile
import os
from pathlib import Path
from datetime import datetime
import sys

# Add project root to path for real imports
sys.path.insert(0, str(Path(__file__).parent))

def load_environment_config():
    """Load real environment configuration"""
    env_file = Path(".env.local")
    if not env_file.exists():
        env_file = Path(".env")
    
    config = {}
    if env_file.exists():
        with open(env_file) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    config[key.strip()] = value.strip().strip('"\'')
    
    # Also check environment variables
    for key in ['RPA_TOKEN_ADMIN', 'RPA_TOKEN_READONLY', 'ENCRYPTION_PASSWORD']:
        if key in os.environ:
            config[key] = os.environ[key]
    
    return config

def create_realistic_rawsession_data():
    """Create realistic RawSession data matching actual system format"""
    return {
        "schema_version": "rawsession_v1",
        "session_id": f"test_auth_session_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
        "agent_version": "0.1.0",
        "environment": {
            "platform": "linux",
            "hostname": "test-machine",
            "screen": {
                "primary_resolution": [1920, 1080],
                "display_scale": 1.0
            }
        },
        "user": {
            "id": "test_user_auth",
            "label": "Test User with Auth"
        },
        "context": {
            "customer": "Test Customer Auth",
            "training_label": "Upload_Auth_Test",
            "notes": "Real functionality test for upload WITH authentication"
        },
        "started_at": datetime.now().isoformat() + "Z",
        "ended_at": datetime.now().isoformat() + "Z",
        "events": [
            {
                "t": 0.5,
                "type": "mouse_click",
                "button": "left",
                "pos": [450, 320],
                "window": {
                    "title": "Test Application",
                    "app_name": "test_app"
                },
                "screenshot_id": "shot_0001"
            },
            {
                "t": 1.2,
                "type": "key_combo",
                "keys": ["CTRL", "C"],
                "window": {
                    "title": "Test Application",
                    "app_name": "test_app"
                },
                "screenshot_id": "shot_0002"
            }
        ],
        "screenshots": [
            {
                "screenshot_id": "shot_0001",
                "relative_path": "shots/shot_0001.png",
                "captured_at": datetime.now().isoformat() + "Z"
            },
            {
                "screenshot_id": "shot_0002",
                "relative_path": "shots/shot_0002.png",
                "captured_at": datetime.now().isoformat() + "Z"
            }
        ]
    }

def create_realistic_test_zip():
    """Create a realistic ZIP file matching agent_v0 format"""
    temp_dir = Path(tempfile.mkdtemp())
    session_data = create_realistic_rawsession_data()
    session_id = session_data["session_id"]
    
    # Create session directory structure
    session_dir = temp_dir / session_id
    session_dir.mkdir(parents=True)
    shots_dir = session_dir / "shots"
    shots_dir.mkdir()
    
    # Create session JSON file
    session_file = session_dir / f"{session_id}.json"
    with open(session_file, 'w') as f:
        json.dump(session_data, f, indent=2)
    
    # Create dummy screenshot files
    png_data = b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\tpHYs\x00\x00\x0b\x13\x00\x00\x0b\x13\x01\x00\x9a\x9c\x18\x00\x00\x00\x12IDATx\x9cc```bPPP\x00\x02\xac\x01\x00\x00\x05\x00\x01\r\n\x87\xdc\x00\x00\x00\x00IEND\xaeB`\x82'
    
    for shot_id in ["shot_0001", "shot_0002"]:
        screenshot_file = shots_dir / f"{shot_id}.png"
        with open(screenshot_file, 'wb') as f:
            f.write(png_data)
    
    # Create ZIP file
    zip_path = temp_dir / f"{session_id}.zip"
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        # Add session JSON
        zipf.write(session_file, f"{session_id}.json")
        # Add screenshots
        for shot_id in ["shot_0001", "shot_0002"]:
            screenshot_file = shots_dir / f"{shot_id}.png"
            zipf.write(screenshot_file, f"shots/{shot_id}.png")
    
    return zip_path, session_id

def test_upload_with_auth():
    """Test upload with proper authentication using realistic data"""
    
    # Load real environment configuration
    config = load_environment_config()
    
    # Check if we have authentication tokens
    admin_token = config.get('RPA_TOKEN_ADMIN')
    readonly_token = config.get('RPA_TOKEN_READONLY')
    
    if not admin_token and not readonly_token:
        print("⚠️  No authentication tokens found in environment")
        print("This test requires RPA_TOKEN_ADMIN or RPA_TOKEN_READONLY")
        print("Check .env.local or .env files, or set environment variables")
        return False
    
    # Use admin token if available, otherwise readonly
    auth_token = admin_token or readonly_token
    token_type = "admin" if admin_token else "readonly"
    
    zip_path = None
    try:
        # Create realistic test data
        zip_path, session_id = create_realistic_test_zip()
        
        url = "http://127.0.0.1:8000/api/traces/upload"
        
        print(f"Testing authenticated upload to: {url}")
        print(f"Session ID: {session_id}")
        print(f"Using {token_type} token: {auth_token[:10]}...")
        print(f"ZIP file size: {zip_path.stat().st_size} bytes")
        
        # Test the actual upload endpoint with authentication
        headers = {
            'Authorization': f'Bearer {auth_token}',
            'X-API-Key': auth_token  # Some APIs use this header instead
        }
        
        with open(zip_path, 'rb') as f:
            files = {'file': (zip_path.name, f, 'application/zip')}
            data = {'session_id': session_id}
            
            response = requests.post(
                url,
                files=files,
                data=data,
                headers=headers,
                timeout=30
            )
        
        print(f"Status Code: {response.status_code}")
        print(f"Response Headers: {dict(response.headers)}")
        print(f"Response Body: {response.text}")
        
        # Analyze the response for real functionality validation
        if response.status_code == 200:
            print("✅ SUCCESS: Authenticated upload worked!")
            
            # Try to parse response as JSON
            try:
                response_data = response.json()
                print(f"Response data: {response_data}")
                
                # Validate response structure
                if 'status' in response_data and response_data['status'] == 'success':
                    print("✅ Server confirmed successful processing")
                
                if 'session_id' in response_data:
                    print(f"✅ Server confirmed session_id: {response_data['session_id']}")
                
                # Check if processing was queued
                if 'queued' in response_data:
                    print(f"✅ Processing queued: {response_data['queued']}")
                
            except json.JSONDecodeError:
                print("Response is not JSON, but upload was accepted")
            
            return True
            
        elif response.status_code == 401:
            print("❌ AUTHENTICATION FAILED: Token was rejected")
            print("Check if the token is valid and not expired")
            return False
            
        elif response.status_code == 403:
            print("❌ AUTHORIZATION FAILED: Token lacks required permissions")
            print(f"The {token_type} token may not have upload permissions")
            return False
            
        else:
            print(f"❌ UNEXPECTED STATUS: {response.status_code}")
            print("This might indicate a server issue or configuration problem")
            return False
            
    except requests.exceptions.ConnectionError as e:
        print(f"❌ CONNECTION ERROR: Server not running at {url}")
        print(f"Error: {e}")
        print("Start the server with: ./run.sh --server")
        return False
        
    except Exception as e:
        print(f"❌ UNEXPECTED ERROR: {e}")
        import traceback
        traceback.print_exc()
        return False
        
    finally:
        # Clean up test files
        if zip_path and zip_path.exists():
            import shutil
            temp_dir = zip_path.parent
            try:
                shutil.rmtree(temp_dir)
                print(f"Cleaned up test directory: {temp_dir}")
            except Exception as e:
                print(f"Warning: Could not clean up {temp_dir}: {e}")

def validate_server_running():
    """Check if the RPA Vision V3 server is running"""
    try:
        response = requests.get("http://127.0.0.1:8000/api/traces/status", timeout=5)
        return response.status_code in [200, 401, 403]
    except:
        return False

def check_server_processing_pipeline():
    """Test if the server can process uploaded data"""
    try:
        # Try to get queue status
        response = requests.get("http://127.0.0.1:8000/api/traces/queue", timeout=5)
        if response.status_code in [200, 401, 403]:
            print("✅ Server processing pipeline endpoints are available")
            return True
    except:
        pass
    
    print("⚠️  Could not verify processing pipeline status")
    return False

if __name__ == "__main__":
    print("🚀 RPA Vision V3 - Real Functionality Test: Upload WITH Authentication")
    print("=" * 75)
    
    # Check if server is running first
    if not validate_server_running():
        print("❌ Server is not running!")
        print("Please start the server first:")
        print("  ./run.sh --server")
        print("  # or")
        print("  python server/api_upload.py")
        exit(1)
    
    print("✅ Server is running")
    
    # Check processing pipeline
    check_server_processing_pipeline()
    print()
    
    success = test_upload_with_auth()
    
    print()
    print("=" * 75)
    if success:
        print("🎉 Authenticated upload test completed successfully!")
        print("The server properly processed the upload with authentication.")
    else:
        print("❌ Authenticated upload test failed!")
        print("Check server logs, authentication configuration, and tokens.")
    
    exit(0 if success else 1)