"""Comprehensive test suite for all FORGE AI tools"""
import asyncio
import httpx
import time
from datetime import datetime

BASE_URL = "http://localhost:8020/api/v1"

test_results = {
    "image_generation": {},
    "image_upscaling": None,
    "background_removal": None,
    "video_generation": {},
    "video_upscaling": None,
}

async def wait_for_job(job_id: str, timeout: int = 180):
    """Wait for a job to complete"""
    async with httpx.AsyncClient() as client:
        start = time.time()
        while time.time() - start < timeout:
            resp = await client.get(f"{BASE_URL}/jobs/{job_id}")
            job = resp.json()
            status = job.get("status")
            progress = job.get("progress", 0)

            print(f"  Job {job_id[:8]}: {status} ({progress}%)")

            if status == "completed":
                return job
            elif status == "failed":
                print(f"  ERROR: {job.get('error_message')}")
                return job

            await asyncio.sleep(3)

        print(f"  TIMEOUT after {timeout}s")
        return None

async def test_image_provider(provider: str, model: str, options: dict):
    """Test an image generation provider"""
    print(f"\n{'='*60}")
    print(f"Testing {provider.upper()} Image Generation")
    print(f"{'='*60}")

    async with httpx.AsyncClient() as client:
        payload = {
            "prompt": f"A beautiful landscape painting, testing {provider}",
            "provider": provider,
            "model": model,
            "provider_options": options
        }

        print(f"Creating job...")
        resp = await client.post(f"{BASE_URL}/modules/image/generate", json=payload)

        if resp.status_code != 200:
            print(f"✗ FAILED: {resp.status_code} - {resp.text}")
            return {"success": False, "error": resp.text}

        job = resp.json()
        job_id = job["id"]
        print(f"✓ Job created: {job_id[:8]}...")

        result = await wait_for_job(job_id)

        if result and result.get("status") == "completed":
            output_assets = result.get("output_asset_ids") or []
            print(f"✓ SUCCESS: Generated {len(output_assets)} image(s)")
            return {"success": True, "job_id": job_id, "assets": output_assets}
        else:
            print(f"✗ FAILED")
            return {"success": False, "error": result.get("error_message") if result else "Timeout"}

async def main():
    """Run all tests"""
    print(f"\n{'#'*60}")
    print(f"# FORGE AI COMPREHENSIVE TOOL TEST")
    print(f"# Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{'#'*60}\n")

    # Test image generation providers
    print("\n" + "="*60)
    print("PHASE 1: IMAGE GENERATION PROVIDERS")
    print("="*60)

    providers_to_test = [
        ("openai", "gpt-image-1", {"quality": "low", "n": 1}),
        ("ideogram", "V_3", {"aspect_ratio": "ASPECT_1_1", "num_images": 1}),
        ("flux", "flux-2-pro", {"width": 512, "height": 512, "steps": 20}),
        ("stable-diffusion", "sd3.5-large", {"aspect_ratio": "1:1"}),
        ("nano-banana", "gemini-2.5-flash-image", {"aspect_ratio": "1:1", "image_size": "1K"}),
    ]

    for provider, model, options in providers_to_test:
        result = await test_image_provider(provider, model, options)
        test_results["image_generation"][provider] = result
        await asyncio.sleep(2)

    # Print summary
    print(f"\n\n{'#'*60}")
    print(f"# TEST RESULTS SUMMARY")
    print(f"{'#'*60}\n")

    print("IMAGE GENERATION:")
    for provider, result in test_results["image_generation"].items():
        status = "✓ PASS" if result.get("success") else "✗ FAIL"
        assets = f"({len(result.get('assets', []))} images)" if result.get("success") else ""
        error = f" - {result.get('error', '')[:50]}" if not result.get("success") else ""
        print(f"  {status} {provider:20s} {assets}{error}")

    print(f"\nCompleted: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")

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