#!/usr/bin/env python3
"""
Test script for generating MultiTalk videos via API
"""

import requests
import json
import time
import base64
import os
from typing import Dict, Any, Optional


def encode_file_to_base64(file_path: str) -> str:
    """Encode a file to base64 string"""
    with open(file_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def save_base64_to_file(base64_string: str, output_path: str):
    """Save base64 string to file"""
    with open(output_path, "wb") as f:
        f.write(base64.b64decode(base64_string))


def test_generate_video(
    api_url: str,
    image_path: str,
    audio_path: str,
    prompt: str,
    second_audio_path: Optional[str] = None,
    num_frames: int = 49,
    sampling_steps: int = 20,
    turbo: bool = True,
    seed: Optional[int] = None,
) -> Dict[str, Any]:
    """
    Test the video generation functionality

    Args:
        api_url: Base URL of the MultiTalk API
        image_path: Path to input image file
        audio_path: Path to first audio file
        prompt: Text prompt for generation
        second_audio_path: Optional path to second audio file for multi-person
        num_frames: Number of frames to generate
        sampling_steps: Number of sampling steps
        turbo: Whether to use turbo mode
        seed: Random seed for reproducibility

    Returns:
        Dictionary with test results
    """
    print("🎬 Testing MultiTalk Video Generation")
    print(f"API URL: {api_url}")
    print(f"Image: {image_path}")
    print(f"Audio: {audio_path}")
    print(f"Second Audio: {second_audio_path or 'None'}")
    print(f"Prompt: {prompt}")
    print(f"Frames: {num_frames}, Steps: {sampling_steps}, Turbo: {turbo}")
    print("-" * 50)

    results = {
        "health_check": False,
        "models_ready": False,
        "generation_success": False,
        "video_saved": False,
        "metadata": None,
        "total_time": 0,
        "generation_time": 0,
        "errors": [],
    }

    start_time = time.time()

    try:
        # 1. Health check
        print("1. Checking API health...")
        health_response = requests.get(f"{api_url}/health")

        if health_response.status_code == 200:
            results["health_check"] = True
            print("✅ Health check passed")
        else:
            results["errors"].append(
                f"Health check failed: {health_response.status_code}"
            )
            print(f"❌ Health check failed: {health_response.status_code}")
            return results

        # 2. Check models status
        print("\n2. Checking models status...")
        status_response = requests.get(f"{api_url}/models/status")

        if status_response.status_code == 200:
            status_data = status_response.json()
            results["models_ready"] = status_data["status"] == "ready"
            print(f"📊 Models status: {status_data['status']}")

            if not results["models_ready"]:
                print("❌ Models not ready. Please run download test first.")
                results["errors"].append("Models not ready for generation")
                # return results
            else:
                print("✅ Models ready for generation")
        else:
            results["errors"].append(
                f"Models status check failed: {status_response.status_code}"
            )
            print(f"❌ Models status check failed: {status_response.status_code}")
            # return results

        if not results["models_ready"]:
            # 2. Check models status
            print("\n2. Downloading models...")
            status_response = requests.get(f"{api_url}/models/download")
            if status_response.status_code == 200:
                status_data = status_response.json()

                print(f"📊 Models downloaded: {status_data['status']}")

        # 3. Prepare input files
        print("\n3. Preparing input files...")

        # Check if files exist
        if not os.path.exists(image_path):
            results["errors"].append(f"Image file not found: {image_path}")
            print(f"❌ Image file not found: {image_path}")
            return results

        if not os.path.exists(audio_path):
            results["errors"].append(f"Audio file not found: {audio_path}")
            print(f"❌ Audio file not found: {audio_path}")
            return results

        if second_audio_path and not os.path.exists(second_audio_path):
            results["errors"].append(
                f"Second audio file not found: {second_audio_path}"
            )
            print(f"❌ Second audio file not found: {second_audio_path}")
            return results

        # Encode files to base64
        print("📦 Encoding files to base64...")
        image_b64 = encode_file_to_base64(image_path)
        audio_b64 = encode_file_to_base64(audio_path)
        second_audio_b64 = None

        if second_audio_path:
            second_audio_b64 = encode_file_to_base64(second_audio_path)

        print(f"✅ Image encoded: {len(image_b64)} chars")
        print(f"✅ Audio encoded: {len(audio_b64)} chars")
        if second_audio_b64:
            print(f"✅ Second audio encoded: {len(second_audio_b64)} chars")

        # 4. Prepare request payload
        print("\n4. Preparing request payload...")

        payload = {
            "image": image_b64,
            "first_audio": audio_b64,
            "prompt": prompt,
            "num_frames": num_frames,
            "sampling_steps": sampling_steps,
            "turbo": turbo,
        }

        if second_audio_b64:
            payload["second_audio"] = second_audio_b64

        if seed is not None:
            payload["seed"] = seed

        print(f"✅ Payload prepared with {len(payload)} fields")

        # 5. Generate video
        print("\n5. Starting video generation...")
        print("⏳ This may take several minutes...")

        generation_start = time.time()

        generate_response = requests.post(
            f"{api_url}/generate",
            json=payload,
        )

        results["generation_time"] = time.time() - generation_start

        if generate_response.status_code == 200:
            response_data = generate_response.json()
            results["generation_success"] = True
            results["metadata"] = response_data.get("metadata", {})

            print("✅ Video generation successful!")
            print(f"⏱️ Generation time: {results['generation_time']:.2f} seconds")

            # Print metadata
            if results["metadata"]:
                print("\n📋 Generation metadata:")
                for key, value in results["metadata"].items():
                    print(f"   {key}: {value}")

            # 6. Save generated video
            print("\n6. Saving generated video...")

            if "video" in response_data:
                # Generate output filename
                timestamp = int(time.time())
                output_filename = f"generated_video_{timestamp}.mp4"

                # Save video
                save_base64_to_file(response_data["video"], output_filename)
                results["video_saved"] = True

                # Get file size
                file_size = os.path.getsize(output_filename)
                print(f"✅ Video saved: {output_filename}")
                print(f"📦 File size: {file_size / (1024*1024):.2f} MB")
            else:
                results["errors"].append("No video data in response")
                print("❌ No video data in response")

        else:
            results["errors"].append(
                f"Generation failed: {generate_response.status_code}"
            )
            print(f"❌ Generation failed: {generate_response.status_code}")
            try:
                error_detail = generate_response.json()
                print(f"Error details: {error_detail}")
                results["errors"].append(str(error_detail))
            except:
                print(f"Response text: {generate_response.text}")
                results["errors"].append(generate_response.text)

    except requests.exceptions.Timeout:
        results["errors"].append("Request timeout")
        print("❌ Request timeout")
    except requests.exceptions.ConnectionError:
        results["errors"].append("Connection error")
        print("❌ Connection error")
    except Exception as e:
        results["errors"].append(f"Unexpected error: {str(e)}")
        print(f"❌ Unexpected error: {str(e)}")

    finally:
        results["total_time"] = time.time() - start_time
        print(f"\n⏱️ Total time: {results['total_time']:.2f} seconds")

    return results


def print_summary(results: Dict[str, Any]):
    """Print a summary of the test results"""
    print("\n" + "=" * 50)
    print("📋 TEST SUMMARY")
    print("=" * 50)

    print(f"Health Check: {'✅' if results['health_check'] else '❌'}")
    print(f"Models Ready: {'✅' if results['models_ready'] else '❌'}")
    print(f"Generation Success: {'✅' if results['generation_success'] else '❌'}")
    print(f"Video Saved: {'✅' if results['video_saved'] else '❌'}")
    print(f"Generation Time: {results['generation_time']:.2f} seconds")
    print(f"Total Time: {results['total_time']:.2f} seconds")

    if results["metadata"]:
        print("\n📊 Generation Details:")
        for key, value in results["metadata"].items():
            print(f"   {key}: {value}")

    if results["errors"]:
        print("\n❌ Errors encountered:")
        for error in results["errors"]:
            print(f"   • {error}")


if __name__ == "__main__":
    # Configuration
    API_URL = "https://abdul-khaliqqq18--multitalk-inference-fastapi-app.modal.run"  # Replace with your actual Modal URL

    # Test files (you need to provide these)
    IMAGE_PATH = r"C:\Users\vinoth.v.kolandairaj\OneDrive - Accenture\Downloads\IMG_0276.JPG"  # Path to your test image
    AUDIO_PATH = r"C:\Users\vinoth.v.kolandairaj\OneDrive - Accenture\Downloads\harvard.wav"  # Path to your test audio
    SECOND_AUDIO_PATH = None  # Path to second audio for multi-person (optional)

    # Generation parameters
    PROMPT = "A person talking naturally with subtle head movements and expressions"
    NUM_FRAMES = 25  # Must be in format 4n+1 (25, 29, 33, 37, 41, 45, 49, etc.)
    SAMPLING_STEPS = 10
    TURBO = True
    SEED = 42  # For reproducible results

    # You can also get these from environment variables
    import os

    API_URL = API_URL
    IMAGE_PATH = IMAGE_PATH
    AUDIO_PATH = AUDIO_PATH
    SECOND_AUDIO_PATH = os.getenv("TEST_SECOND_AUDIO_PATH", SECOND_AUDIO_PATH)

    print("🚀 Starting MultiTalk Video Generation Test")
    print(f"Using API URL: {API_URL}")
    print("Make sure to:")
    print("1. Replace API_URL with your actual Modal deployment URL")
    print("2. Provide valid test_image.jpg and test_audio.wav files")
    print("3. Ensure models are downloaded (run download test first)")
    print()

    # Check if test files exist
    if not os.path.exists(IMAGE_PATH):
        print(f"❌ Test image not found: {IMAGE_PATH}")
        print("Please provide a test image file")
        exit(1)

    if not os.path.exists(AUDIO_PATH):
        print(f"❌ Test audio not found: {AUDIO_PATH}")
        print("Please provide a test audio file")
        exit(1)

    # Run the test
    results = test_generate_video(
        api_url=API_URL,
        image_path=IMAGE_PATH,
        audio_path=AUDIO_PATH,
        prompt=PROMPT,
        second_audio_path=SECOND_AUDIO_PATH,
        num_frames=NUM_FRAMES,
        sampling_steps=SAMPLING_STEPS,
        turbo=TURBO,
        seed=SEED,
    )

    # Print summary
    print_summary(results)

    # Save results to file
    with open("generation_test_results.json", "w") as f:
        # Remove the actual video data from results before saving
        results_copy = results.copy()
        if "video_data" in results_copy:
            del results_copy["video_data"]
        json.dump(results_copy, f, indent=2)

    print(f"\n💾 Results saved to generation_test_results.json")

    # Exit with appropriate code
    if (
        results["generation_success"]
        and results["video_saved"]
        and not results["errors"]
    ):
        print("\n🎉 All tests passed!")
        exit(0)
    else:
        print("\n⚠️ Some tests failed!")
        exit(1)