import argparse
import os
import sys
import requests
from dotenv import load_dotenv
from stream_utils import StreamResponseHandler

# Load environment variables
load_dotenv()


def get_default_image():
    """Get the default image path and convert it to a data URI."""
    image_path = "./assets/lakeview.jpg"
    if os.path.exists(image_path):
        try:
            from src.utils import image_path_to_uri

            image_uri = image_path_to_uri(image_path)
            print(f"Using default image: {image_path}")
            return image_uri
        except Exception as e:
            print(f"Error converting image to URI: {e}")
            return None
    else:
        print(f"Warning: Default image not found at {image_path}")
        return None


def upload_image(handler, image_path):
    """
    Upload an image to the server.

    Args:
        handler (StreamResponseHandler): The stream response handler.
        image_path (str): Path to the image file to upload.

    Returns:
        str: The URL of the uploaded image, or None if upload failed.
    """
    if not os.path.exists(image_path):
        handler.console.print(
            f"[bold red]Error:[/] Image file not found at {image_path}"
        )
        return None

    try:
        handler.console.print(f"Uploading image: [bold]{image_path}[/]")
        with open(image_path, "rb") as f:
            files = {"file": (os.path.basename(image_path), f)}
            response = requests.post("http://localhost:8000/upload", files=files)
            if response.status_code == 200:
                image_url = response.json().get("image_url")
                handler.console.print(
                    f"Image uploaded successfully. URL: [bold green]{image_url}[/]"
                )
                return image_url
            else:
                handler.console.print(
                    f"[bold red]Failed to upload image.[/] Status code: {response.status_code}"
                )
                handler.console.print(f"Response: {response.text}")
                return None
    except Exception as e:
        handler.console.print(f"[bold red]Error uploading image:[/] {e}")
        return None


def main():
    # Create a stream response handler
    handler = StreamResponseHandler()

    # Parse command line arguments
    parser = argparse.ArgumentParser(description="Test the image edit streaming API.")
    parser.add_argument(
        "--instruction", "-i", required=True, help="The edit instruction."
    )
    parser.add_argument("--image", "-img", help="The URL of the image to edit.")
    parser.add_argument("--upload", "-u", help="Path to an image file to upload first.")

    args = parser.parse_args()

    # Check if the server is running
    if not handler.check_server_health():
        sys.exit(1)

    image_url = args.image

    # If upload is specified, upload the image first
    if args.upload:
        image_url = upload_image(handler, args.upload)
        if not image_url:
            handler.console.print(
                "[yellow]Warning:[/] Failed to upload image. Continuing without image URL."
            )

    # Use the default image if no image URL is provided
    if not image_url:
        image_url = get_default_image()
        if not image_url:
            handler.console.print(
                "[yellow]No image URL provided and default image not available.[/]"
            )
            handler.console.print("The agent may ask for an image if needed.")

    # Prepare the payload for the edit request
    payload = {"edit_instruction": args.instruction}

    if image_url:
        payload["image_url"] = image_url

    # Stream the edit request
    endpoint_url = "http://localhost:8000/edit/stream"
    handler.stream_response(endpoint_url, payload=payload, title="Image Edit Response")


if __name__ == "__main__":
    main()