import gradio as gr
import asyncio
from threading import RLock
from pathlib import Path
from huggingface_hub import InferenceClient
import os


HF_TOKEN = os.environ.get("HF_TOKEN") if os.environ.get("HF_TOKEN") else None # If private or gated models aren't used, ENV setting is unnecessary.
server_timeout = 600
inference_timeout = 300


lock = RLock()
loaded_models = {}
model_info_dict = {}


def to_list(s):
    return [x.strip() for x in s.split(",")]


def list_sub(a, b):
    return [e for e in a if e not in b]


def list_uniq(l):
        return sorted(set(l), key=l.index)


def is_repo_name(s):
    import re
    return re.fullmatch(r'^[^/]+?/[^/]+?$', s)


def get_status(model_name: str):
    from huggingface_hub import InferenceClient
    client = InferenceClient(token=HF_TOKEN, timeout=10)
    return client.get_model_status(model_name)


def is_loadable(model_name: str, force_gpu: bool = False):
    try:
        status = get_status(model_name)
    except Exception as e:
        print(e)
        print(f"Couldn't load {model_name}.")
        return False
    gpu_state = isinstance(status.compute_type, dict) and "gpu" in status.compute_type.keys()
    if status is None or status.state not in ["Loadable", "Loaded"] or (force_gpu and not gpu_state):
        print(f"Couldn't load {model_name}. Model state:'{status.state}', GPU:{gpu_state}")
    return status is not None and status.state in ["Loadable", "Loaded"] and (not force_gpu or gpu_state)


def find_model_list(author: str="", tags: list[str]=[], not_tag="", sort: str="last_modified", limit: int=30, force_gpu=False, check_status=False):
    from huggingface_hub import HfApi
    api = HfApi(token=HF_TOKEN)
    default_tags = ["diffusers"]
    if not sort: sort = "last_modified"
    limit = limit * 20 if check_status and force_gpu else limit * 5
    models = []
    try:
        model_infos = api.list_models(author=author, task="text-to-image",
                                       tags=list_uniq(default_tags + tags), cardData=True, sort=sort, limit=limit)
    except Exception as e:
        print(f"Error: Failed to list models.")
        print(e)
        return models
    for model in model_infos:
        if not model.private and not model.gated or HF_TOKEN is not None:
           loadable = is_loadable(model.id, force_gpu) if check_status else True
           if not_tag and not_tag in model.tags or not loadable: continue
           models.append(model.id)
           if len(models) == limit: break
    return models


def get_t2i_model_info_dict(repo_id: str):
    from huggingface_hub import HfApi
    api = HfApi(token=HF_TOKEN)
    info = {"md": "None"}
    try:
        if not is_repo_name(repo_id) or not api.repo_exists(repo_id=repo_id): return info
        model = api.model_info(repo_id=repo_id, token=HF_TOKEN)
    except Exception as e:
        print(f"Error: Failed to get {repo_id}'s info.")
        print(e)
        return info
    if model.private or model.gated and HF_TOKEN is None: return info
    try:
        tags = model.tags
    except Exception as e:
        print(e)
        return info
    if not 'diffusers' in model.tags: return info
    if 'diffusers:FluxPipeline' in tags: info["ver"] = "FLUX.1"
    elif 'diffusers:StableDiffusionXLPipeline' in tags: info["ver"] = "SDXL"
    elif 'diffusers:StableDiffusionPipeline' in tags: info["ver"] = "SD1.5"
    elif 'diffusers:StableDiffusion3Pipeline' in tags: info["ver"] = "SD3"
    else: info["ver"] = "Other"
    info["url"] = f"https://huggingface.co/{repo_id}/"
    info["tags"] = model.card_data.tags if model.card_data and model.card_data.tags else []
    info["downloads"] = model.downloads
    info["likes"] = model.likes
    info["last_modified"] = model.last_modified.strftime("lastmod: %Y-%m-%d")
    un_tags = ['text-to-image', 'stable-diffusion', 'stable-diffusion-api', 'safetensors', 'stable-diffusion-xl']
    descs = [info["ver"]] + list_sub(info["tags"], un_tags) + [f'DLs: {info["downloads"]}'] + [f'❤: {info["likes"]}'] + [info["last_modified"]]
    info["md"] = f'Model Info: {", ".join(descs)} [Model Repo]({info["url"]})'
    return info


def rename_image(image_path: str | None, model_name: str, save_path: str | None = None):
    from PIL import Image, ImageFile
    ImageFile.LOAD_TRUNCATED_IMAGES = True
    from datetime import datetime, timezone, timedelta
    if image_path is None: return None
    dt_now = datetime.now(timezone(timedelta(hours=9)))
    filename = f"{model_name.split('/')[-1]}_{dt_now.strftime('%Y%m%d_%H%M%S')}.png"
    try:
        if Path(image_path).exists():
            png_path = "image.png"
            Image.open(image_path).convert('RGBA').save(png_path, "PNG")
            if save_path is not None:
                new_path = str(Path(png_path).resolve().rename(Path(save_path).resolve()))
            else:
                new_path = str(Path(png_path).resolve().rename(Path(filename).resolve()))
            return new_path
        else:
            return None
    except Exception as e:
        print(e)
        return None


def save_gallery(image_path: str | None, images: list[tuple] | None):
    if images is None: images = []
    files = [i[0] for i in images]
    if image_path is None: return images, files
    files.insert(0, str(image_path))
    images.insert(0, (str(image_path), Path(image_path).stem))
    return images, files


# https://github.com/gradio-app/gradio/blob/main/gradio/external.py
# https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
from typing import Literal
def load_from_model(model_name: str, hf_token: str | Literal[False] | None = None):
    import httpx
    import huggingface_hub
    from gradio.exceptions import ModelNotFoundError, TooManyRequestsError
    model_url = f"https://huggingface.co/{model_name}"
    api_url = f"https://api-inference.huggingface.co/models/{model_name}"
    print(f"Fetching model from: {model_url}")

    headers = ({} if hf_token in [False, None] else {"Authorization": f"Bearer {hf_token}"})
    response = httpx.request("GET", api_url, headers=headers)
    if response.status_code != 200:
        raise ModelNotFoundError(
            f"Could not find model: {model_name}. If it is a private or gated model, please provide your Hugging Face access token (https://huggingface.co/settings/tokens) as the argument for the `hf_token` parameter."
        )
    p = response.json().get("pipeline_tag")
    if p != "text-to-image": raise ModelNotFoundError(f"This model isn't for text-to-image or unsupported: {model_name}.")
    headers["X-Wait-For-Model"] = "true"
    client = huggingface_hub.InferenceClient(model=model_name, headers=headers,
                                              token=hf_token, timeout=server_timeout)
    inputs = gr.components.Textbox(label="Input")
    outputs = gr.components.Image(label="Output")
    fn = client.text_to_image

    def query_huggingface_inference_endpoints(*data, **kwargs):
        try:
            data = fn(*data, **kwargs)  # type: ignore
        except huggingface_hub.utils.HfHubHTTPError as e:
            if "429" in str(e):
                raise TooManyRequestsError() from e
        except Exception as e:
            raise Exception(e)
        return data

    interface_info = {
        "fn": query_huggingface_inference_endpoints,
        "inputs": inputs,
        "outputs": outputs,
        "title": model_name,
    }
    return gr.Interface(**interface_info)


def load_model(model_name: str):
    global loaded_models
    global model_info_dict
    if model_name in loaded_models.keys(): return loaded_models[model_name]
    try:
        loaded_models[model_name] = load_from_model(model_name, hf_token=HF_TOKEN)
        print(f"Loaded: {model_name}")
    except Exception as e:
        if model_name in loaded_models.keys(): del loaded_models[model_name]
        print(f"Failed to load: {model_name}")
        print(e)
        return None
    try:
        model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
        print(f"Assigned: {model_name}")
    except Exception as e:
        if model_name in model_info_dict.keys(): del model_info_dict[model_name]
        print(f"Failed to assigned: {model_name}")
        print(e)
    return loaded_models[model_name]


def load_model_api(model_name: str):
    global loaded_models
    global model_info_dict
    if model_name in loaded_models.keys(): return loaded_models[model_name]
    try:
        client = InferenceClient(timeout=5)
        status = client.get_model_status(model_name, token=HF_TOKEN)
        if status is None or status.framework != "diffusers" or status.state not in ["Loadable", "Loaded"]:
            print(f"Failed to load by API: {model_name}")
            return None
        else:
            loaded_models[model_name] = InferenceClient(model_name, token=HF_TOKEN, timeout=server_timeout)
            print(f"Loaded by API: {model_name}")
    except Exception as e:
        if model_name in loaded_models.keys(): del loaded_models[model_name]
        print(f"Failed to load by API: {model_name}")
        print(e)
        return None
    try:
        model_info_dict[model_name] = get_t2i_model_info_dict(model_name)
        print(f"Assigned by API: {model_name}")
    except Exception as e:
        if model_name in model_info_dict.keys(): del model_info_dict[model_name]
        print(f"Failed to assigned by API: {model_name}")
        print(e)
    return loaded_models[model_name]


def load_models(models: list):
    for model in models:
        load_model(model)


positive_prefix = {
    "Pony": to_list("score_9, score_8_up, score_7_up"),
    "Pony Anime": to_list("source_anime, anime, score_9, score_8_up, score_7_up"),
}
positive_suffix = {
    "Common": to_list("highly detailed, masterpiece, best quality, very aesthetic, absurdres"),
    "Anime": to_list("anime artwork, anime style, studio anime, highly detailed"),
}
negative_prefix = {
    "Pony": to_list("score_6, score_5, score_4"),
    "Pony Anime": to_list("score_6, score_5, score_4, source_pony, source_furry, source_cartoon"),
    "Pony Real": to_list("score_6, score_5, score_4, source_anime, source_pony, source_furry, source_cartoon"),
}
negative_suffix = {
    "Common": to_list("lowres, (bad), bad hands, bad feet, text, error, fewer, extra, missing, worst quality, jpeg artifacts, low quality, watermark, unfinished, displeasing, oldest, early, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]"),
    "Pony Anime": to_list("busty, ugly face, mutated hands, low res, blurry face, black and white, the simpsons, overwatch, apex legends"),
    "Pony Real": to_list("ugly, airbrushed, simple background, cgi, cartoon, anime"),
}
positive_all = negative_all = []
for k, v in (positive_prefix | positive_suffix).items():
    positive_all = positive_all + v + [s.replace("_", " ") for s in v]
positive_all = list_uniq(positive_all)
for k, v in (negative_prefix | negative_suffix).items():
    negative_all = negative_all + v + [s.replace("_", " ") for s in v]
positive_all = list_uniq(positive_all)


def recom_prompt(prompt: str = "", neg_prompt: str = "", pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
    def flatten(src):
        return [item for row in src for item in row]
    prompts = to_list(prompt)
    neg_prompts = to_list(neg_prompt)
    prompts = list_sub(prompts, positive_all)
    neg_prompts = list_sub(neg_prompts, negative_all)
    last_empty_p = [""] if not prompts and type != "None" else []
    last_empty_np = [""] if not neg_prompts and type != "None" else []
    prefix_ps = flatten([positive_prefix.get(s, []) for s in pos_pre])
    suffix_ps = flatten([positive_suffix.get(s, []) for s in pos_suf])
    prefix_nps = flatten([negative_prefix.get(s, []) for s in neg_pre])
    suffix_nps = flatten([negative_suffix.get(s, []) for s in neg_suf])
    prompt = ", ".join(list_uniq(prefix_ps + prompts + suffix_ps) + last_empty_p)
    neg_prompt = ", ".join(list_uniq(prefix_nps + neg_prompts + suffix_nps) + last_empty_np)
    return prompt, neg_prompt


recom_prompt_type = {
    "None": ([], [], [], []),
    "Auto": ([], [], [], []),
    "Common": ([], ["Common"], [], ["Common"]),
    "Animagine": ([], ["Common", "Anime"], [], ["Common"]),
    "Pony": (["Pony"], ["Common"], ["Pony"], ["Common"]),
    "Pony Anime": (["Pony", "Pony Anime"], ["Common", "Anime"], ["Pony", "Pony Anime"], ["Common", "Pony Anime"]),
    "Pony Real": (["Pony"], ["Common"], ["Pony", "Pony Real"], ["Common", "Pony Real"]),
}


enable_auto_recom_prompt = False
def insert_recom_prompt(prompt: str = "", neg_prompt: str = "", type: str = "None"):
    global enable_auto_recom_prompt
    if type == "Auto":  enable_auto_recom_prompt = True
    else: enable_auto_recom_prompt = False
    pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
    return recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)


def set_recom_prompt_preset(type: str = "None"):
    pos_pre, pos_suf, neg_pre, neg_suf = recom_prompt_type.get(type, ([], [], [], []))
    return pos_pre, pos_suf, neg_pre, neg_suf


def get_recom_prompt_type():
    type = list(recom_prompt_type.keys())
    type.remove("Auto")
    return type


def get_positive_prefix():
    return list(positive_prefix.keys())


def get_positive_suffix():
    return list(positive_suffix.keys())


def get_negative_prefix():
    return list(negative_prefix.keys())


def get_negative_suffix():
    return list(negative_suffix.keys())


def get_tag_type(pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = []):
    tag_type = "danbooru"
    words = pos_pre + pos_suf + neg_pre + neg_suf
    for word in words:
        if "Pony" in word:
            tag_type = "e621"
            break
    return tag_type


def get_model_info_md(model_name: str):
    if model_name in model_info_dict.keys(): return model_info_dict[model_name].get("md", "")


def change_model(model_name: str):
    load_model_api(model_name)
    return get_model_info_md(model_name)


def warm_model(model_name: str):
    model = load_model_api(model_name)
    if model:
        try:
            print(f"Warming model: {model_name}")
            infer_body(model, " ")
        except Exception as e:
            print(e)


# https://huggingface.co/docs/api-inference/detailed_parameters
# https://huggingface.co/docs/huggingface_hub/package_reference/inference_client
def infer_body(client: InferenceClient | gr.Interface | object, prompt: str, neg_prompt: str | None = None,
               height: int | None = None, width: int | None = None,
               steps: int | None = None, cfg: int | None = None, seed: int = -1):
    png_path = "image.png"
    kwargs = {}
    if height is not None and height >= 256: kwargs["height"] = height
    if width is not None and width >= 256: kwargs["width"] = width
    if steps is not None and steps >= 1: kwargs["num_inference_steps"] = steps
    if cfg is not None and cfg > 0: cfg = kwargs["guidance_scale"] = cfg
    if seed >= 0: kwargs["seed"] = seed
    try:
        if isinstance(client, InferenceClient):
            image = client.text_to_image(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
        elif isinstance(client, gr.Interface):
            image = client.fn(prompt=prompt, negative_prompt=neg_prompt, **kwargs, token=HF_TOKEN)
        else: return None
        if isinstance(image, tuple): return None
        image.save(png_path)
        return str(Path(png_path).resolve())
    except Exception as e:
        print(e)
        raise Exception(e)


async def infer(model_name: str, prompt: str, neg_prompt: str | None = None,
               height: int | None = None, width: int | None = None,
               steps: int | None = None, cfg: int | None = None, seed: int = -1,
               save_path: str | None = None, timeout: float = inference_timeout):
    import random
    noise = ""
    if seed < 0:
        rand = random.randint(1, 500)
        for i in range(rand):
            noise += " "
    model = load_model(model_name)
    if not model: return None
    task = asyncio.create_task(asyncio.to_thread(infer_body, model, f"{prompt} {noise}", neg_prompt,
                                                 height, width, steps, cfg, seed))
    await asyncio.sleep(0)
    try:
        result = await asyncio.wait_for(task, timeout=timeout)
    except asyncio.TimeoutError as e:
        print(e)
        print(f"Task timed out: {model_name}")
        if not task.done(): task.cancel()
        result = None
        raise Exception(f"Task timed out: {model_name}")
    except Exception as e:
        print(e)
        if not task.done(): task.cancel()
        result = None
        raise Exception(e)
    if task.done() and result is not None:
        with lock:
            image = rename_image(result, model_name, save_path)
        return image
    return None


# https://github.com/aio-libs/pytest-aiohttp/issues/8 # also AsyncInferenceClient is buggy.
def infer_fn(model_name: str, prompt: str, neg_prompt: str | None = None, height: int | None = None,
             width: int | None = None, steps: int | None = None, cfg: int | None = None, seed: int = -1,
             pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], save_path: str | None = None):
    if model_name == 'NA':
        return None
    try:
        loop = asyncio.get_running_loop()
    except Exception:
        loop = asyncio.new_event_loop()
    try:
        prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
        result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
                                               steps, cfg, seed, save_path, inference_timeout))
    except (Exception, asyncio.CancelledError) as e:
        print(e)
        print(f"Task aborted: {model_name}, Error: {e}")
        result = None
        raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
    finally:
        loop.close()
    return result


def infer_rand_fn(model_name_dummy: str, prompt: str, neg_prompt: str | None = None, height: int | None = None,
             width: int | None = None, steps: int | None = None, cfg: int | None = None, seed: int = -1,
             pos_pre: list = [], pos_suf: list = [], neg_pre: list = [], neg_suf: list = [], save_path: str | None = None):
    import random
    if model_name_dummy == 'NA':
        return None
    random.seed()
    model_name = random.choice(list(loaded_models.keys()))
    try:
        loop = asyncio.get_running_loop()
    except Exception:
        loop = asyncio.new_event_loop()
    try:
        prompt, neg_prompt = recom_prompt(prompt, neg_prompt, pos_pre, pos_suf, neg_pre, neg_suf)
        result = loop.run_until_complete(infer(model_name, prompt, neg_prompt, height, width,
                                               steps, cfg, seed, save_path, inference_timeout))
    except (Exception, asyncio.CancelledError) as e:
        print(e)
        print(f"Task aborted: {model_name}, Error: {e}")
        result = None
        raise gr.Error(f"Task aborted: {model_name}, Error: {e}")
    finally:
        loop.close()
    return result