import requests
import io
import base64

from typing import Any

from PIL import Image


class ChatAPI:
    def __init__(self, api_key: str, model: str = ""):
        self.api_key = api_key
        self.model = model

    def call_api(self, prompt: Any):
        _ = prompt

        return "Not implemented"

    def get_tokens(self):
        return {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0}


class VisionAPI:
    def __init__(self, api_key: str, model: str = ""):
        self.api_key = api_key
        self.model = model

    def call_api(self, prompt: Any):
        _ = prompt

        return "Not implemented"

    def get_tokens(self):
        return {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0}


class ImageAPI:
    def __init__(
        self,
        api_key: str,
        model: str = "",
        aspect_ratio: str = "1:1",
        strength: float = 0.5,
        seed: int = 0,
        cfg_scale: int = 10,
        style_preset: str = "",
        **kwargs,
    ):

        self.api_key = api_key
        self.model = model
        self.aspect_ratio = aspect_ratio
        self.seed = seed
        self.cfg_scale = cfg_scale
        self.strength = strength
        self.style_preset = style_preset

        self.root_url = "https://api.stability.ai/v2beta/stable-image"

        self.response = None
        self.tokens = None

        self.model_type = None
        self.model_name = None
        self.action = None

        self.allowed_ar = [
            "16:9",
            "1:1",
            "21:9",
            "2:3",
            "3:2",
            "4:5",
            "5:4",
            "9:16",
            "9:21",
        ]

        self.allowed_sp = [
            "3d-model",
            "analog-film",
            "anime",
            "cinematic",
            "comic-book",
            "digital-art",
            "enhance",
            "fantasy-art",
            "isometric",
            "line-art",
            "low-poly",
            "modeling-compound",
            "neon-punk",
            "origami",
            "photographic",
            "pixel-art",
            "tile-texture",
        ]

        self.__set_model_info()
        self.__build_url()
        self.__check_aspect_ratio()
        self.__check_style_preset()

    def __set_model_info(self):
        splitted = self.model.split("/")

        self.model_type = splitted[-1]
        self.action = splitted[1]
        self.model_name = splitted[-2] if splitted[-2] != "default" else None

    def __build_url(self):
        self.url = f"{self.root_url}/{self.action}/{self.model_type}"

    def __check_aspect_ratio(self):
        if self.aspect_ratio not in self.allowed_ar:
            self.aspect_ratio = "1:1"

    def __check_style_preset(self):
        if self.style_preset not in self.allowed_sp:
            self.style_preset = "photographic"

    def __build_upscale_data(self, prompt: Any, image: Any):

        _ = image

        payload = {
            "output_format": "png",
        }

        if self.model_type != "fast":
            payload["prompt"] = prompt

        return payload

    def __build_generate_data(self, prompt: Any, image: Any):

        payload = {
            "prompt": prompt,
            "output_format": "png",
            "seed": self.seed,
            "cfg_scale": self.cfg_scale,
            "style_preset": self.style_preset,
        }

        if self.model_name:
            payload["model"] = self.model_name

        if image:
            payload["strength"] = self.strength
            payload["mode"] = "image-to-image"
        else:
            payload["aspect_ratio"] = self.aspect_ratio
            payload["mode"] = "text-to-image"

        return payload

    def __build_data(self, prompt: Any, image: Any):

        build_data_dict = {
            "upscale": self.__build_upscale_data,
            "generate": self.__build_generate_data,
        }

        build_function = build_data_dict[self.action]
        return build_function(prompt, image)

    def __build_files(self, image: Any):
        if not image:
            return {"none": ""}
        else:
            if isinstance(image, str):
                try:
                    img = Image.open(image)
                except Exception:
                    img = Image.open(io.BytesIO(base64.b64decode(image)))

                img_byte_arr = io.BytesIO()

                img = self._resize_image(img)
                img.save(img_byte_arr, format="PNG")

                img_byte_arr = img_byte_arr.getvalue()

                return {"image": img_byte_arr}
            elif isinstance(image, Image.Image):
                img_byte_arr = io.BytesIO()

                image = self._resize_image(image)
                image.save(img_byte_arr, format="PNG")

                img_byte_arr = img_byte_arr.getvalue()

                return {"image": img_byte_arr}
            elif isinstance(image, bytes):
                img = Image.open(io.BytesIO(image))
                img_byte_arr = io.BytesIO()

                img = self._resize_image(img)
                img.save(img_byte_arr, format="PNG")

                img_byte_arr = img_byte_arr.getvalue()

                return {"image": img_byte_arr}
            else:
                raise Exception(
                    "Incorrect image type! Accepted: img_string, PIL Image, or bytes"
                )

    def _resize_image(self, image: Image.Image) -> Image.Image:
        max_size = 16384
        min_size = 64
        min_total = 4096

        width, height = image.size

        if width * height < min_total:
            ratio = (width * height) / min_total

            width = int(width * ratio)
            height = int(height * ratio)

            image = image.resize((width, height))

        if max(width, height) > max_size:
            aspect_ratio = width / height
            if width > height:
                new_width = max_size
                new_height = int(new_width / aspect_ratio)
            else:
                new_height = max_size
                new_width = int(new_height * aspect_ratio)

            image = image.resize((new_width, new_height))

        elif min(width, height) < min_size:
            aspect_ratio = width / height
            if width < height:
                new_width = min_size
                new_height = int(new_width / aspect_ratio)
            else:
                new_height = min_size
                new_width = int(new_height * aspect_ratio)

            image = image.resize((new_width, new_height))

        return image

    def call_api(self, prompt: Any, image: Any = None):
        data = self.__build_data(prompt, image)
        files = self.__build_files(image)

        self.response = requests.post(
            self.url,
            headers={"authorization": f"Bearer {self.api_key}", "accept": "image/*"},
            files=files,
            data=data,
        )

        if self.response.status_code == 200:
            self.tokens = self.get_tokens()
            return self.get_image()
        else:
            raise Exception(str(self.response.json()))

    def get_image(self):
        if self.model_type == "creative":
            return self.response.json()["id"]
        return base64.b64encode(self.response.content).decode("utf-8")

    def get_tokens(self):
        return 1


class AudioAPI:
    def __init__(self, api_key: str, model: str, **kwargs):
        self.api_key = api_key
        self.model = model

    def call_api(self, audio: Any):
        _ = audio

        return self.get_output()

    def get_output(self):
        return "Not Implemented"

    def get_tokens(self):
        return {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0}


class SpeechAPI:
    def __init__(self, api_key: str, model: str, **kwargs):
        self.api_key = api_key
        self.model = model

    def call_api(self, text: str) -> bytes:
        _ = text

        return self.get_output()

    def get_output(self):
        return "Not Implemented"

    def get_tokens(self):
        return 0
