import asyncio
import argparse
import aiohttp
import base64
import os


async def api_request(url: str,
                      username: str, password: str, params: dict) -> dict:
    headers = {
        "Authorization": aiohttp.BasicAuth(username, password).encode(),
        "Accept-Encoding": "identity",
    }
    async with aiohttp.ClientSession() as session:
        async with session.post(f"{url}/sdapi/v1/txt2img",
                                headers=headers, json=params) as resp:
            return await resp.json()


async def save_images(username: str, password: str,
                      prompt: str, negative_prompt: str,
                      destination: str, params_dict: dict[str, str]) -> None:
    params = {
        "enable_hr": True,
        "sampler name": "DPM++ 2M Karras",
        "width": 896,
        "height": 896,
        "steps": 30,
        "cfg_scale": 12,
        "n_iter": 3,
        "batch_size": 1,
        "denoising_strength": 0.6,
        "prompt": prompt,
        "negative_prompt": negative_prompt,
    }
    response = await api_request(username, password, params)
    for i, image_data in enumerate(response["images"]):
        image_bytes = base64.decodebytes(image_data.encode())
        with open(os.path.join(destination, f"file{i + 1}.png"), "wb") as f:
            f.write(image_bytes)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--username", default="root",
        help="Username for basic authentication (default: root)")
    parser.add_argument(
        "--password", default="itsreallyme",
        help="Password for basic authentication (default: itsreallyme)")
    parser.add_argument("url", help="URL to use for API")
    parser.add_argument("prompt", help="Prompt for image generation")
    parser.add_argument("negative_prompt",
                        help="Negative prompt for image generation")
    parser.add_argument("destination", help="Directory to save images to")
    args = parser.parse_args()

    asyncio.run(
        save_images(
            args.url,
            args.username,
            args.password,
            args.prompt,
            args.negative_prompt,
            args.destination
            )
    )