import colorsys
import os
import random
from pathlib import Path

import requests
from openai import OpenAI
from tqdm import tqdm

client = OpenAI(
    base_url='https://external.api.recraft.ai/v1',
    api_key=os.getenv("RECRAFT_API_KEY"),
)


def generate_random_rgb_color() -> list[int]:
    r = random.randint(0, 255)
    g = random.randint(0, 255)
    b = random.randint(0, 255)
    return [r, g, b]

def generate_color_palette(n=5):
    return [generate_random_rgb_color() for _ in range(n)]


def create_icons(prompt: str) -> str:
    colors = [{"rgb": colors} for colors in generate_color_palette(5)]
    print(colors)
    response = client.images.generate(
        prompt=prompt,
        model='recraft20b',
        extra_body={
            'style': 'icon',
           # 'substyle': 'broken_line',
            'substyle': 'offset_fill',
            'controls': {
                'colors': colors
            },
        }
    )

    return response.data[0].url


def download_svg(url, output_file):
    try:
        response = requests.get(url)
        response.raise_for_status()  # Raise an error for bad HTTP responses
        if 'image/svg+xml' in response.headers.get('Content-Type', ''):
            with open(output_file, 'wb') as file:
                file.write(response.content)
            print(f"SVG file saved as {output_file}")
        else:
            print("The URL does not point to an SVG file.")
    except requests.exceptions.RequestException as e:
        print(f"An error occurred: {e}")

def create_and_save_icon(prompt: str, output_path: Path):
    output_path.mkdir(exist_ok=True)
    url = create_icons(prompt)
    print(url)
    prompt = prompt.replace(" ", "_")
    download_svg(url, output_path / f"{prompt}.svg")

if __name__ == "__main__":
    output_path = Path(f"../output/recraft_v0/")
    with open("../data/adjective_noun_pairs.txt", 'r') as fp:
        pairs = fp.read().split("\n")

    svg_files = []
    if output_path.exists():
        svg_files = os.listdir(output_path)

    svg_files = {file[:-4].replace("_", " ") for file in svg_files if file.endswith(".svg")}

    for pair in tqdm(pairs[:440]):
        if pair in svg_files:
            print(f"skipping {pair}")
            continue
        create_and_save_icon(pair, output_path)
