import os
from typing import List

import numpy as np
import onnxruntime as ort
from PIL import Image
from PIL.Image import Image as PILImage
from rembg.sessions import BaseSession


class CustomBaseSession(BaseSession):
    def __init__(self, model_name: str):
        sess_opts = ort.SessionOptions()

        if "OMP_NUM_THREADS" in os.environ:
            sess_opts.inter_op_num_threads = int(os.environ["OMP_NUM_THREADS"])

        super().__init__(model_name, sess_opts)


class CustomSessionContainer:
    def __init__(self, mean_x, mean_y, mean_z, std_x, std_y, std_z, width, height) -> None:
        self.mean_x = mean_x
        self.mean_y = mean_y
        self.mean_z = mean_z
        self.std_x = std_x
        self.std_y = std_y
        self.std_z = std_z
        self.width = width
        self.height = height


class CustomAbstractSession(CustomBaseSession, CustomSessionContainer):
    def predict(self, img: PILImage, *args, **kwargs) -> List[PILImage]:
        ort_outs = self.inner_session.run(
            None,
            self.normalize(
                img,
                (self.mean_x, self.mean_y, self.mean_z),
                (self.std_x, self.std_y, self.std_z),
                (self.width, self.height)
            ),
        )

        pred = ort_outs[0][:, 0, :, :]

        ma = np.max(pred)
        mi = np.min(pred)

        pred = (pred - mi) / (ma - mi)
        pred = np.squeeze(pred)

        mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
        mask = mask.resize(img.size, Image.LANCZOS)

        return [mask]

    @classmethod
    def download_models(cls, *args, **kwargs):
        return os.path.join(cls.u2net_home(), f"{cls.name()}")

    def from_container(self, container: CustomSessionContainer):
        self.mean_x = container.mean_x
        self.mean_y = container.mean_y
        self.mean_z = container.mean_z
        self.std_x = container.std_x
        self.std_y = container.std_y
        self.std_z = container.std_z
        self.width = container.width
        self.height = container.height

        return self