import os
import gradio as gr
from run_cmd import run_cmd
from PIL import Image
import uuid
import numpy as np

temp_path = "./upscaled"


def inference(img, size, type):
    if not os.path.exists(temp_path):
        os.mkdir(temp_path)

    image = Image.open(img)

    OUTPUT_PATH = os.path.join(
        temp_path, f"{str(uuid.uuid4())[0:12]}_{size}.png")

    image.save(OUTPUT_PATH)

    if type == "Manga":
        run_cmd(f"python inference_manga_v2.py {OUTPUT_PATH}")
    else:
        run_cmd(f"python inference.py {OUTPUT_PATH} {type}")

    img_out = Image.open(OUTPUT_PATH)

    if size == "x2":
        img_out = img_out.resize(
            (img_out.width // 2, img_out.height // 2), resample=Image.BICUBIC)

    img_out = np.array(img_out)

    return img_out, gr.DownloadButton(
        value=OUTPUT_PATH,
        visible=True,
    )