import argparse
import os
import requests
import fal_client
from dotenv import load_dotenv

load_dotenv()

# Map FAL_API_KEY to FAL_KEY if needed
if os.getenv("FAL_API_KEY") and not os.getenv("FAL_KEY"):
    os.environ["FAL_KEY"] = os.getenv("FAL_API_KEY")


def create_ai_image(prompt: str, aspect_ratio: str = "16:9", output_path: str = None) -> str:
    """Generate an image using fal.ai nano-banana model."""
    result = fal_client.subscribe(
        "fal-ai/nano-banana",
        arguments={
            "prompt": prompt,
            "aspect_ratio": aspect_ratio,
        },
    )
    image_url = result["images"][0]["url"]
    if output_path:
        response = requests.get(image_url)
        with open(output_path, "wb") as f:
            f.write(response.content)
        return output_path
    return image_url


def edit_ai_image(input_file: str, prompt: str, aspect_ratio: str = "16:9", output_path: str = None) -> str:
    """Edit an image using fal.ai nano-banana/edit model."""
    image_url = fal_client.upload_file(input_file)
    result = fal_client.subscribe(
        "fal-ai/nano-banana/edit",
        arguments={
            "image_urls": [image_url],
            "prompt": prompt,
            "aspect_ratio": aspect_ratio,
        },
    )
    result_url = result["images"][0]["url"]
    if output_path:
        response = requests.get(result_url)
        with open(output_path, "wb") as f:
            f.write(response.content)
        return output_path
    return result_url


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="AI image tools using fal.ai")
    subparsers = parser.add_subparsers(dest="command", required=True)

    create_parser = subparsers.add_parser("create", help="Create an AI-generated image")
    create_parser.add_argument("prompt", help="Text prompt for image generation")
    create_parser.add_argument("--aspect-ratio", default="16:9", help="Aspect ratio (default: 16:9)")
    create_parser.add_argument("--output", "-o", help="Output file path to save the image")

    edit_parser = subparsers.add_parser("edit", help="Edit an image with AI")
    edit_parser.add_argument("input_file", help="Input image file to edit")
    edit_parser.add_argument("prompt", help="Text prompt for image editing")
    edit_parser.add_argument("--aspect-ratio", default="16:9", help="Aspect ratio (default: 16:9)")
    edit_parser.add_argument("--output", "-o", help="Output file path to save the image")

    args = parser.parse_args()

    if args.command == "create":
        result = create_ai_image(args.prompt, args.aspect_ratio, args.output)
        if args.output:
            print(f"Image saved to: {result}")
        else:
            print(f"Generated image URL: {result}")
    elif args.command == "edit":
        result = edit_ai_image(args.input_file, args.prompt, args.aspect_ratio, args.output)
        if args.output:
            print(f"Image saved to: {result}")
        else:
            print(f"Edited image URL: {result}")
