import onnxruntime as ort
from PIL import Image
import numpy as np
import os
from tqdm import tqdm
import requests
import hashlib
from typing import List, Union
import shutil
from pathlib import Path

def process_image(image:Image.Image) -> np.ndarray:
    '''
    Convert an image to a numpy array.
    :param image: the image to convert
    :return: the numpy array
    '''

    image = image.convert("RGB").resize((512,512))
    image = np.array(image).astype(np.float32) / 255
    image = image.transpose((2,0,1)).reshape(1,3,512,512).transpose((0,2,3,1))
    return image

def download(url:str, save_path:str, md5:str, length:str) -> bool:
    '''
    Download a file from url to save_path.
    If the file already exists, check its md5.
    If the md5 matches, return True,if the md5 doesn't match, return False.
    :param url: the url of the file to download
    :param save_path: the path to save the file
    :param md5: the md5 of the file
    :param length: the length of the file
    :return: True if the file is downloaded successfully, False otherwise
    '''

    try:
        response = requests.get(url=url, stream=True)
        with open(save_path, "wb") as f:
            with tqdm.wrapattr(response.raw, "read", total=length, desc="Downloading") as r_raw:
                shutil.copyfileobj(r_raw, f) 
        return True if hashlib.md5(open(save_path, "rb").read()).hexdigest() == md5 else False
    except Exception as e:
        print(e)
        return False

def download_model():
    '''
    Download the model and tags file from the server.
    :return: the path to the model and tags file
    '''

    model_url = "https://huggingface.co/chinoll/deepdanbooru/resolve/main/deepdanbooru.onnx"
    tags_url = "https://huggingface.co/chinoll/deepdanbooru/resolve/main/tags.txt"
    model_md5 = "16be4e40ebcc0b1d1915bbf31f00969f"
    tags_md5 = "a3f764de985cdeba89f1d232a4204402"
    model_length = 643993025
    tags_length = 133810

    home = str(Path.home()) + "/.deepdanbooru_onnx/"
    if not os.path.exists(home):
        os.mkdir(home)

    model_name = "deepdanbooru.onnx"
    tags_name = "tags.txt"

    model_path = home + model_name
    tags_path = home + tags_name
    if os.path.exists(model_path):
        if hashlib.md5(open(model_path, "rb").read()).hexdigest() != model_md5:
            os.remove(model_path)
            if not download(model_url, model_path, model_md5, model_length):
                raise ValueError("Model download failed")

    else:
        if not download(model_url, model_path, model_md5, model_length):
            raise ValueError("Model download failed")

    if os.path.exists(tags_path):
        if hashlib.md5(open(tags_path, "rb").read()).hexdigest() != tags_md5:
            os.remove(tags_path)
            if not download(tags_url, tags_path, tags_md5, tags_length):
                raise ValueError("Tags download failed")
    else:
        if not download(tags_url, tags_path, tags_md5, tags_length):
            raise ValueError("Tags download failed")
    return model_path, tags_path

class DeepDanbooru:
    def __init__(self, mode: str = "auto", model_path: Union[str, None] =None, tags_path: Union[str, None] = None, threshold: Union[float, int] = 0.6, pin_memory: bool = False, batch_size: int = 1):
        '''
        Initialize the DeepDanbooru class.
        :param mode: the mode of the model, "cpu" or "gpu" or "auto"
        :param model_path: the path to the model file
        :param tags_path: the path to the tags file
        :param threshold: the threshold of the model
        :param pin_memory: whether to use pin memory
        :param batch_size: the batch size of the model
        '''

        providers = {
            "cpu":"CPUExecutionProvider",
            "gpu":"CUDAExecutionProvider",
            "tensorrt": "TensorrtExecutionProvider",
            "auto":"CUDAExecutionProvider" if "CUDAExecutionProvider" in ort.get_available_providers() else "CPUExecutionProvider",
        }

        if not (isinstance(threshold,float) or isinstance(threshold,int)):
            raise TypeError("threshold must be float or int")
        if threshold < 0 or threshold > 1:
            raise ValueError("threshold must be between 0 and 1")
        if mode not in providers:
            raise ValueError("Mode not supported. Please choose from: cpu, gpu, tensorrt")
        if providers[mode] not in ort.get_available_providers():
            raise ValueError(f"Your device is not supported {mode}. Please choose from: cpu")
        if model_path is not None and not os.path.exists(model_path):
            raise FileNotFoundError("Model file not found")
        if tags_path is not None and not os.path.exists(tags_path):
            raise FileNotFoundError("Tags file not found")

        if model_path is None or tags_path is None:
            model_path, tags_path = download_model()

        self.session = ort.InferenceSession(model_path, providers=[providers[mode]])
        self.tags = [i.replace("\n","") for i in open(tags_path, "r").readlines()]

        self.input_name = self.session.get_inputs()[0].name
        self.output_name = [output.name for output in self.session.get_outputs()]
        self.threshold = threshold
        self.pin_memory = pin_memory
        self.batch_size = batch_size
        self.mode = mode
    
    def __str__(self) -> str:
        return f"DeepDanbooru(mode={self.mode}, threshold={self.threshold}, pin_memory={self.pin_memory}, batch_size={self.batch_size})"

    def __repr__(self) -> str:
        return self.__str__()

    def from_image_inference(self,image:Image.Image) -> dict:
        image = process_image(image)
        return self.predict(image)

    def from_ndarray_inferece(self, image:np.ndarray) -> dict:
        if image.shape != (1,512,512,3):
            raise ValueError(f"Image must be {(1,512,512,3)}")
        return self.predict(image)

    def from_path_inference(self,image:str) -> dict:
        return self.from_image_inference(Image.open(image))

    def from_list_inference(self,image:Union[list,tuple]) -> List[dict]:
        if self.pin_memory:
            image = [process_image(Image.open(i)) for i in image]
        for i in [image[i:i + self.batch_size] for i in range(0, len(image), self.batch_size)]:
            imagelist = i
            if not self.pin_memory:
                imagelist = []
                for j in i:
                    imagelist.append(process_image(Image.open(j)))
            _image = np.vstack(imagelist)
            results = self.inference(_image)
            for i in range(results.shape[0]):
                image_tag = {}
                for tag, score in zip(self.tags, results[i]):
                    if score >= self.threshold:
                        image_tag[tag] = score
                yield image_tag

    def inference(self,image):
        return self.session.run(self.output_name, {self.input_name:image})[0]

    def predict(self,image):
        result = self.inference(image)
        image_tag = {}
        for tag, score in zip(self.tags, result[0]):
            if score >= self.threshold:
                image_tag[tag] = score
        return image_tag

    def __call__(self, image) -> Union[dict, List[dict]]:
        if isinstance(image,str):
            return self.from_path_inference(image)
        elif isinstance(image,np.ndarray):
            return self.from_ndarray_inferece(image)
        elif isinstance(image,list) or isinstance(image,tuple):
            return self.from_list_inference(image)
        elif isinstance(image,Image.Image):
            return self.from_image_inference(image)
        else:
            raise ValueError("Image must be a path to an image or a numpy array")