#!/usr/bin/env python3
"""
Test script for AI Mockup Service Authentication and Wireframe Saving
This script tests the complete flow from authentication to wireframe saving
"""

import requests
import json
import time
import os
from datetime import datetime

# Configuration
AI_MOCKUP_SERVICE_URL = "http://localhost:8021"
USER_AUTH_SERVICE_URL = "http://localhost:8011"
POSTGRES_HOST = "localhost"
POSTGRES_PORT = "5433"  # Docker mapped port
POSTGRES_DB = "dev_pipeline"
POSTGRES_USER = "pipeline_admin"
POSTGRES_PASSWORD = "secure_pipeline_2024"

def test_health_checks():
    """Test health endpoints"""
    print("🔍 Testing health checks...")
    
    # Test AI Mockup Service health
    try:
        response = requests.get(f"{AI_MOCKUP_SERVICE_URL}/health", timeout=10)
        if response.status_code == 200:
            print("✅ AI Mockup Service is healthy")
            print(f"   Status: {response.json().get('status')}")
            print(f"   Database: {response.json().get('services', {}).get('database')}")
            print(f"   User Auth: {response.json().get('services', {}).get('user_auth')}")
        else:
            print(f"❌ AI Mockup Service health check failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ AI Mockup Service health check error: {e}")
        return False
    
    # Test User Auth Service health
    try:
        response = requests.get(f"{USER_AUTH_SERVICE_URL}/health", timeout=10)
        if response.status_code == 200:
            print("✅ User Auth Service is healthy")
        else:
            print(f"❌ User Auth Service health check failed: {response.status_code}")
            return False
    except Exception as e:
        print(f"❌ User Auth Service health check error: {e}")
        return False
    
    return True

def test_authentication():
    """Test authentication flow"""
    print("\n🔐 Testing authentication...")
    
    # Test registration
    test_user = {
        "username": f"testuser_{int(time.time())}",
        "email": f"testuser_{int(time.time())}@example.com",
        "password": "TestPassword123!",
        "first_name": "Test",
        "last_name": "User"
    }
    
    try:
        response = requests.post(f"{USER_AUTH_SERVICE_URL}/api/auth/register", 
                               json=test_user, timeout=10)
        if response.status_code == 201:
            print("✅ User registration successful")
            user_data = response.json().get('data', {})
            user_id = user_data.get('user', {}).get('id')
            if user_id:
                print("✅ User ID received")
                
                # For testing purposes, manually verify the email by updating the database
                # This bypasses the email verification requirement
                try:
                    import psycopg2
                    conn = psycopg2.connect(
                        host=POSTGRES_HOST,
                        database=POSTGRES_DB,
                        user=POSTGRES_USER,
                        password=POSTGRES_PASSWORD,
                        port=POSTGRES_PORT
                    )
                    with conn.cursor() as cur:
                        cur.execute("UPDATE users SET email_verified = true WHERE id = %s", (user_id,))
                        conn.commit()
                    conn.close()
                    print("✅ Email verification bypassed for testing")
                except Exception as e:
                    print(f"⚠️ Could not bypass email verification: {e}")
                    return None, None
                
                # Now try to login
                login_response = requests.post(f"{USER_AUTH_SERVICE_URL}/api/auth/login",
                                            json={"email": test_user["email"], "password": test_user["password"]},
                                            timeout=10)
                
                if login_response.status_code == 200:
                    login_data = login_response.json().get('data', {})
                    access_token = login_data.get('tokens', {}).get('accessToken')
                    if access_token:
                        print("✅ Access token received")
                        return access_token, test_user
                    else:
                        print("❌ No access token in login response")
                        return None, None
                else:
                    print(f"❌ Login failed: {login_response.status_code}")
                    print(f"   Response: {login_response.text}")
                    return None, None
            else:
                print("❌ No user ID in response")
                return None, None
        else:
            print(f"❌ User registration failed: {response.status_code}")
            print(f"   Response: {response.text}")
            return None, None
    except Exception as e:
        print(f"❌ User registration error: {e}")
        return None, None

def test_wireframe_generation(access_token):
    """Test wireframe generation"""
    print("\n🎨 Testing wireframe generation...")
    
    headers = {"Authorization": f"Bearer {access_token}"}
    prompt = "Create a simple landing page with header, hero section, and footer"
    
    try:
        response = requests.post(f"{AI_MOCKUP_SERVICE_URL}/generate-wireframe/desktop",
                              json={"prompt": prompt},
                              headers=headers,
                              timeout=30)
        
        if response.status_code == 200:
            print("✅ Wireframe generation successful")
            result = response.json()
            if result.get('svg'):
                print("✅ SVG wireframe received")
                return result.get('svg')
            else:
                print("❌ No SVG in response")
                return None
        else:
            print(f"❌ Wireframe generation failed: {response.status_code}")
            print(f"   Response: {response.text}")
            return None
    except Exception as e:
        print(f"❌ Wireframe generation error: {e}")
        return None

def test_wireframe_saving(access_token, svg_data):
    """Test wireframe saving"""
    print("\n💾 Testing wireframe saving...")
    
    headers = {"Authorization": f"Bearer {access_token}"}
    
    wireframe_data = {
        "wireframe": {
            "name": f"Test Wireframe {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
            "description": "Test wireframe created by automated test",
            "device_type": "desktop",
            "dimensions": {"width": 1440, "height": 1024},
            "metadata": {"prompt": "Test prompt", "generator": "test"}
        },
        "elements": [
            {
                "id": "test-element-1",
                "type": "shape",
                "data": {"type": "rectangle", "props": {"w": 200, "h": 100}},
                "position": {"x": 100, "y": 100},
                "size": {"width": 200, "height": 100},
                "style": {"color": "#3B82F6", "fill": "#EFF6FF"},
                "parent_id": None,
                "z_index": 0
            }
        ]
    }
    
    try:
        response = requests.post(f"{AI_MOCKUP_SERVICE_URL}/api/wireframes",
                              json=wireframe_data,
                              headers=headers,
                              timeout=10)
        
        if response.status_code == 201:
            print("✅ Wireframe saved successfully")
            result = response.json()
            wireframe_id = result.get('wireframe_id')
            if wireframe_id:
                print(f"✅ Wireframe ID: {wireframe_id}")
                return wireframe_id
            else:
                print("❌ No wireframe ID in response")
                return None
        else:
            print(f"❌ Wireframe saving failed: {response.status_code}")
            print(f"   Response: {response.text}")
            return None
    except Exception as e:
        print(f"❌ Wireframe saving error: {e}")
        return None

def test_wireframe_retrieval(access_token, wireframe_id):
    """Test wireframe retrieval"""
    print("\n📖 Testing wireframe retrieval...")
    
    headers = {"Authorization": f"Bearer {access_token}"}
    
    try:
        response = requests.get(f"{AI_MOCKUP_SERVICE_URL}/api/wireframes/{wireframe_id}",
                              headers=headers,
                              timeout=10)
        
        if response.status_code == 200:
            print("✅ Wireframe retrieved successfully")
            result = response.json()
            if result.get('wireframe') and result.get('elements'):
                print("✅ Wireframe data and elements received")
                return True
            else:
                print("❌ Incomplete wireframe data")
                return False
        else:
            print(f"❌ Wireframe retrieval failed: {response.status_code}")
            print(f"   Response: {response.text}")
            return False
    except Exception as e:
        print(f"❌ Wireframe retrieval error: {e}")
        return False

def main():
    """Run all tests"""
    print("🚀 Starting AI Mockup Service Integration Tests")
    print("=" * 50)
    
    # Test 1: Health checks
    if not test_health_checks():
        print("\n❌ Health checks failed. Please ensure all services are running.")
        return
    
    # Test 2: Authentication
    access_token, user_data = test_authentication()
    if not access_token:
        print("\n❌ Authentication failed. Please check user-auth service.")
        return
    
    # Test 3: Wireframe generation
    svg_data = test_wireframe_generation(access_token)
    if not svg_data:
        print("\n❌ Wireframe generation failed.")
        return
    
    # Test 4: Wireframe saving
    wireframe_id = test_wireframe_saving(access_token, svg_data)
    if not wireframe_id:
        print("\n❌ Wireframe saving failed.")
        return
    
    # Test 5: Wireframe retrieval
    if not test_wireframe_retrieval(access_token, wireframe_id):
        print("\n❌ Wireframe retrieval failed.")
        return
    
    print("\n" + "=" * 50)
    print("🎉 All tests passed! Wireframe saving is working correctly.")
    print(f"✅ User: {user_data.get('username')}")
    print(f"✅ Wireframe ID: {wireframe_id}")
    print("✅ Authentication, generation, saving, and retrieval all working")

if __name__ == "__main__":
    main()
