#!/usr/bin/env python3
"""
Test OpenAI Responses API image generation functionality
"""
import requests
import time
import os


def test_responses_image_generation():
    """Test image generation using OpenAI Responses API

    This test will stop execution immediately upon encountering any error
    to avoid expensive API calls.
    """
    print("🧪 Testing OpenAI Responses API Image Generation...")
    print("⚠️  Note: Test will stop on first error to avoid expensive API calls\n")

    # Get a test token (you'll need to use a valid token for your site)
    # For this test, we'll use a placeholder - replace with actual token
    test_token = os.environ.get('TEST_API_TOKEN', '8ff428ef0ed57da0fad18cedef459f4ada77f04b635bb60580ce23ca9145127c')

    base_url = "http://127.0.0.1:8000/api/v1/image-responses"

    # Test 1: Non-streaming image generation
    print("\n1. Testing non-streaming image generation...")
    try:
        image_request = {
            "prompt": "A serene landscape with mountains and a lake at sunset",
            "user": "test_user",
            "model": "gpt-4.1-mini",
            "stream": False
        }

        headers = {
            "Authorization": f"Bearer {test_token}",
            "Content-Type": "application/json"
        }

        response = requests.post(
            f"{base_url}/generate-image/",
            json=image_request,
            headers=headers
        )

        if response.status_code == 200:
            urls = response.json()
            print(f"✅ Non-streaming image generation successful!")
            print(f"   📸 Generated image URL: {urls[0] if urls else 'No URL'}")
        else:
            print(f"❌ Non-streaming generation failed: {response.status_code}")
            print(f"   Error: {response.text}")
            print("\n🛑 Stopping test execution due to error (to avoid expensive API calls)")
            return False

    except Exception as e:
        print(f"❌ Non-streaming test error: {e}")
        print("\n🛑 Stopping test execution due to error (to avoid expensive API calls)")
        return False


    # Test 2: Store and stream image generation
    print("\n2. Testing store and stream image generation...")
    try:
        image_request = {
            "prompt": "A futuristic city with flying cars and neon lights",
            "user": "test_user",
            "model": "gpt-4.1-mini",
            "stream": True
        }

        headers = {
            "Authorization": f"Bearer {test_token}",
            "Content-Type": "application/json"
        }

        # Store the request
        store_response = requests.post(
            f"{base_url}/store-image/",
            json=image_request,
            headers=headers
        )

        if store_response.status_code == 200:
            key = store_response.json()["message"]
            print(f"✅ Image request stored with key: {key}")

            # Stream the generation
            stream_response = requests.get(
                f"{base_url}/stream-image-generation/{key}",
                stream=True
            )

            if stream_response.status_code == 200:
                print("✅ Streaming image generation started!")
                final_url = None
                progress_count = 0
                stream_error = False

                for line in stream_response.iter_lines():
                    if line:
                        line_text = line.decode('utf-8')
                        if line_text.startswith('data: '):
                            content = line_text[6:]
                            if content == '__chatbot_end__':
                                break
                            elif content == '{progress}':
                                progress_count += 1
                                print(f"   📊 Progress update {progress_count}")
                            elif content.startswith('http'):
                                final_url = content
                                print(f"   📸 Final image URL: {final_url}")
                            elif content.startswith('Error:'):
                                print(f"   ❌ Error during generation: {content}")
                                stream_error = True
                                break

                if stream_error or not final_url:
                    print("\n🛑 Stopping test execution due to streaming error (to avoid expensive API calls)")
                    return False

                print("✅ Streaming image generation completed successfully!")

            else:
                print(f"❌ Streaming failed: {stream_response.status_code}")
                print(f"   Error: {stream_response.text}")
                print("\n🛑 Stopping test execution due to error (to avoid expensive API calls)")
                return False

        else:
            print(f"❌ Store request failed: {store_response.status_code}")
            print(f"   Error: {store_response.text}")
            print("\n🛑 Stopping test execution due to error (to avoid expensive API calls)")
            return False

    except Exception as e:
        print(f"❌ Streaming test error: {e}")
        print("\n🛑 Stopping test execution due to error (to avoid expensive API calls)")
        return False

    # exit(0)  # Temporary exit to avoid running streaming test in this environment

    # Test 3: Different prompt styles
    print("\n3. Testing various prompt styles...")
    test_prompts = [
        "A minimalist logo design for a tech startup",
        "An oil painting of a cat wearing a space suit",
        "A photorealistic image of a coffee cup on a wooden table"
    ]

    for i, prompt in enumerate(test_prompts, 1):
        try:
            print(f"\n   Test {i}: {prompt[:50]}...")

            image_request = {
                "prompt": prompt,
                "user": "test_user",
                "model": "gpt-4.1-mini",
                "stream": False
            }

            headers = {
                "Authorization": f"Bearer {test_token}",
                "Content-Type": "application/json"
            }

            response = requests.post(
                f"{base_url}/generate-image/",
                json=image_request,
                headers=headers
            )

            if response.status_code == 200:
                urls = response.json()
                print(f"   ✅ Generated: {urls[0] if urls else 'No URL'}")
            else:
                print(f"   ❌ Failed: {response.status_code}")
                print(f"   Error: {response.text}")
                print("\n🛑 Stopping test execution due to error (to avoid expensive API calls)")
                return False

            # Small delay between requests
            time.sleep(1)

        except Exception as e:
            print(f"   ❌ Error: {e}")
            print("\n🛑 Stopping test execution due to error (to avoid expensive API calls)")
            return False

    # All tests passed!
    print("\n🎉 OpenAI Responses API image generation testing complete!")
    print("   ✅ All tests passed successfully!")
    print("\n📋 Summary:")
    print("   ✅ Non-streaming generation endpoint working")
    print("   ✅ Store and stream endpoints working")
    print("   ✅ Image saving and URL generation working")
    print("   ✅ Multiple prompt styles supported")
    print("\n📝 Notes:")
    print("   - Replace TEST_API_TOKEN environment variable with valid token")
    print("   - Ensure FastAPI server is running on http://127.0.0.1:8000")
    print("   - Check DRUPAL_HOST environment variable is set correctly")
    print("   - Verify site directories exist for image storage")

    return True


if __name__ == "__main__":
    test_responses_image_generation()