import os
import json
import requests

import base64
from io import BytesIO
import numpy as np
from server import PromptServer
from PIL import Image

ALL_CODES_LANGS = ['af', 'sq', 'am', 'ar', 'hy', 'as', 'ay', 'az', 'bm', 'eu', 'be', 'bn', 'bho', 'bs', 'bg', 'ca', 'ceb', 'ny', 'zh-CN', 'zh-TW', 'co', 'hr', 'cs', 'da', 'dv', 'doi', 'nl', 'en', 'eo', 'et', 'ee', 'tl', 'fi', 'fr', 'fy', 'gl', 'ka', 'de', 'el', 'gn', 'gu', 'ht', 'ha', 'haw', 'iw', 'hi', 'hmn', 'hu', 'is', 'ig', 'ilo', 'id', 'ga', 'it', 'ja', 'jw', 'kn', 'kk', 'km', 'rw', 'gom', 'ko', 'kri', 'ku', 'ckb', 'ky', 'lo', 'la', 'lv', 'ln', 'lt', 'lg', 'lb', 'mk', 'mai', 'mg', 'ms', 'ml', 'mt', 'mi', 'mr', 'mni-Mtei', 'lus', 'mn', 'my', 'ne', 'no', 'or', 'om', 'ps', 'fa', 'pl', 'pt', 'pa', 'qu', 'ro', 'ru', 'sm', 'sa', 'gd', 'nso', 'sr', 'st', 'sn', 'sd', 'si', 'sk', 'sl', 'so', 'es', 'su', 'sw', 'sv', 'tg', 'ta', 'tt', 'te', 'th', 'ti', 'ts', 'tr', 'tk', 'ak', 'uk', 'ur', 'ug', 'uz', 'vi', 'cy', 'xh', 'yi', 'yo', 'zu']


ENDPOINT_URL = "https://open.bigmodel.cn/api/paas/v4/chat/completions"

# Language models: https://docs.bigmodel.cn/api-reference/%E6%A8%A1%E5%9E%8B-api/%E5%AF%B9%E8%AF%9D%E8%A1%A5%E5%85%A8#%E6%96%87%E6%9C%AC%E6%A8%A1%E5%9E%8B
LIST_LANGUAGE_MODELS = [
    # GLM-4
    "glm-4-plus",
    "glm-4-air-250414",
    "glm-4-airx",
    "glm-4-flashx",
    "glm-4-flashx-250414",
    # GLM-4.5
    "glm-4.5",
    "glm-4.5-air",
    "glm-4.5-x",
    "glm-4.5-airx",
    "glm-4.5-flash",
    # GLM-Z1
    "glm-z1-air",
    "glm-z1-airx",
    "glm-z1-flash",
    "glm-z1-flashx",
]

# Multimodal models: https://docs.bigmodel.cn/api-reference/%E6%A8%A1%E5%9E%8B-api/%E5%AF%B9%E8%AF%9D%E8%A1%A5%E5%85%A8#%E8%A7%86%E8%A7%89%E6%A8%A1%E5%9E%8B
LIST_MULTIMODAL_MODELS = [
    # --- GLM-4v
    "glm-4v-flash",
    "glm-4v",
    "glm-4v-plus-0111",
    # --- GLM-4.1v
    "glm-4.1v-thinking-flashx",
    "glm-4.1v-thinking-flash",
    # --- GLM-4.5v
    "glm-4.5v",
]

def getConfigData():
    # Directory node and config file
    dir_node = os.path.dirname(__file__)
    config_path = os.path.join(os.path.abspath(dir_node), "config.json")
    config = {
                "__comment": "Register on the site https://bigmodel.cn and get a key and add it to the field ZHIPUAI_API_KEY. Change default translate languages ​​'from' and 'to' you use",
                "from_translate": "ru",
                "to_translate": "en",
                "default_language_model": "glm-4.5-flash",
                "default_multimodal_model": "glm-4v-flash",
                "ZHIPUAI_API_KEY": "your_api_key"
            }

    # Load config.js file
    if not os.path.exists(config_path):
        print("[ChatGLMNode] File config.js file not found! Create default config.json...")
        with open(config_path, "w", encoding="utf-8") as f:
            json.dump(config, f, ensure_ascii=False, indent=4)
            return config
    else:
        with open(config_path, "r", encoding="utf-8") as f:
            config = json.load(f)
            return config
        # =====

def checkPropValue(obj, key, not_include = []):
    checkVal = lambda v: v is None or v.strip() == "" or v in not_include

    prop_val = obj.get(key)

    if checkVal(prop_val):
        obj.update(getConfigData())
        return True if checkVal(obj.get(key)) else False

    else:
        return False


CONFIG = getConfigData()

def createRequest(payload):
    global CONFIG

    if checkPropValue(CONFIG, "ZHIPUAI_API_KEY", ["your_api_key"]):
        raise ValueError("ZHIPUAI_API_KEY value is empty or missing")

    ZHIPUAI_API_KEY = CONFIG.get("ZHIPUAI_API_KEY")

    # Headers
    headers = {
        "Authorization": f"Bearer {ZHIPUAI_API_KEY}",
        "Content-Type": "application/json",
    }

    try:
        response = requests.post(ENDPOINT_URL, headers=headers, json=payload)
        response.raise_for_status()

        if response.status_code == 200:
            json_data = response.json()
            response_text = json_data.get("choices")[0]["message"]["content"].strip()

            return response_text

    except requests.HTTPError as e:
        print(f"Error request ChatGLM: {response.status_code}, {response.text}")
        raise e
    except Exception as e:
        print(f"Error ChatGLM: {e}")
        raise e


def translate(prompt, srcTrans, toTrans, model, max_tokens, temperature, top_p):
    # Check prompt exist
    if prompt is None or prompt.strip() == "":
        return ""

    # Create body request
    payload = {
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": f"Translate from {srcTrans} to {toTrans} and return only the translated text: {prompt}",
            },
        ],
        "max_tokens": round(max_tokens, 2),
        "temperature": round(temperature, 2),
        "top_p": round(top_p, 2),
    }

    response_translate_text = createRequest(payload)

    return response_translate_text


class ChatGLM4TranslateCLIPTextEncodeNode:
    @classmethod
    def INPUT_TYPES(self):
        from_lng = CONFIG.get("from_translate") if CONFIG.get("from_translate") in ALL_CODES_LANGS else "ru"
        to_lng = CONFIG.get("to_translate") if CONFIG.get("to_translate") in ALL_CODES_LANGS else "en"
        return {
            "required": {
                "from_translate": (
                    ALL_CODES_LANGS,
                    {"default": from_lng, "tooltip": "Translation from"},
                ),
                "to_translate": (
                    ALL_CODES_LANGS,
                    {"default": to_lng, "tooltip": "Translation to"},
                ),
                "model": (
                    LIST_LANGUAGE_MODELS,
                    {
                        "default": CONFIG.get("default_language_model", "glm-4.5-flash"),
                        "tooltip": "The model code to be called. Models with text 'flash' should be free!",
                    },
                ),
                "max_tokens": (
                    "INT",
                    {
                        "default": 1024,
                        "tooltip": "The maximum number of tokens for model output, maximum output is 4095, default value is 1024.",
                    },
                ),
                "temperature": (
                    "FLOAT",
                    {
                        "default": 0.95,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Sampling temperature, controls the randomness of the output, must be a positive number within the range: [0.0, 1.0], default value is 0.95.",
                    },
                ),
                "top_p": (
                    "FLOAT",
                    {
                        "default": 0.7,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Another method of temperature sampling, value range is: [0.0, 1.0], default value is 0.7.",
                    },
                ),
                "text": ("STRING", {"multiline": True, "placeholder": "Input text"}),
                "clip": ("CLIP",),
            }
        }

    RETURN_TYPES = (
        "CONDITIONING",
        "STRING",
    )
    FUNCTION = "chatglm_translate_text"
    DESCRIPTION = (
        "This is a node that translates the prompt into another language using ChatGLM."
    )
    CATEGORY = "AlekPet Nodes/conditioning"

    def chatglm_translate_text(
        self,
        from_translate,
        to_translate,
        model,
        max_tokens,
        temperature,
        top_p,
        text,
        clip,
    ):

        text = translate(
            text, from_translate, to_translate, model, max_tokens, temperature, top_p
        )
        tokens = clip.tokenize(text)
        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
        return ([[cond, {"pooled_output": pooled}]], text)


class ChatGLM4TranslateTextNode(ChatGLM4TranslateCLIPTextEncodeNode):
    @classmethod
    def INPUT_TYPES(self):
        return_types = super().INPUT_TYPES()
        del return_types["required"]["clip"]
        return return_types

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "chatglm_translate_text"

    CATEGORY = "AlekPet Nodes/text"

    def chatglm_translate_text(
        self, from_translate, to_translate, model, max_tokens, temperature, top_p, text
    ):

        text = translate(
            text, from_translate, to_translate, model, max_tokens, temperature, top_p
        )

        return (text,)


# ChatGLM Instruct Node
class ChatGLM4InstructNode:
    @classmethod
    def INPUT_TYPES(self):
        return {
            "required": {
                "model": (
                    LIST_LANGUAGE_MODELS,
                    {
                        "default": CONFIG.get("default_language_model", "glm-4.5-flash"),
                        "tooltip": "The model code to be called. Models with text 'flash' should be free!",
                    },
                ),
                "max_tokens": (
                    "INT",
                    {
                        "default": 1024,
                        "tooltip": "The maximum number of tokens for model output, maximum output is 4095, default value is 1024.",
                    },
                ),
                "temperature": (
                    "FLOAT",
                    {
                        "default": 0.95,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Sampling temperature, controls the randomness of the output, must be a positive number within the range: [0.0, 1.0], default value is 0.95.",
                    },
                ),
                "top_p": (
                    "FLOAT",
                    {
                        "default": 0.7,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Another method of temperature sampling, value range is: [0.0, 1.0], default value is 0.7.",
                    },
                ),
                "instruct": (
                    "STRING",
                    {
                        "multiline": True,
                        "placeholder": "Input instruct text",
                        "default": "Generate details text, without quotation marks or the word 'prompt' on english: {query}",
                        "tooltip": "Enter the instruction for the neural network to execute and indicate where to insert the query text {query}",
                    },
                ),
                "query": (
                    "STRING",
                    {
                        "multiline": True,
                        "placeholder": "Enter the query text for the instruction",
                        "tooltip": "Query field",
                    },
                ),
            }
        }

    RETURN_TYPES = ("STRING",)
    FUNCTION = "chatglm_instruct"

    CATEGORY = "AlekPet Nodes/Instruct"

    def chatglm_instruct(self, model, max_tokens, temperature, top_p, instruct, query):

        if instruct is None or instruct.strip() == "":
            raise ValueError("Instruct text is empty!")

        if query is None or query.strip() == "":
            raise ValueError("Query text is empty!")

        instruct = instruct.replace("{query}", query)

        # Create body request
        payload = {
            "model": model,
            "messages": [
                {
                    "role": "user",
                    "content": instruct,
                },
            ],
            "max_tokens": round(max_tokens, 2),
            "temperature": round(temperature, 2),
            "top_p": round(top_p, 2),
        }

        answer = createRequest(payload)

        return (answer,)


# ChatGLM Instruct Media Node
def toBase64ImgUrl(img):
    bytesIO = BytesIO()
    img.save(bytesIO, format="PNG")
    img_types = bytesIO.getvalue()
    img_base64 = base64.b64encode(img_types)
    return f"data:image/png;base64,{img_base64.decode('utf-8')}"


class ChatGLM4InstructMediaNode:
    @classmethod
    def INPUT_TYPES(self):
        return {
            "optional": {
                "image": ("IMAGE",),
                # "video": ("STRING", {"forceInput": True, "default": ""}),
            },
            "required": {
                "model": (
                    LIST_MULTIMODAL_MODELS,
                    {
                        "default": CONFIG.get("default_multimodal_model","glm-4v-flash"),
                        "tooltip": "The model code to be called. Models with text 'flash' should be free!",
                    },
                ),
                "max_tokens": (
                    "INT",
                    {
                        "default": 1024,
                        "tooltip": "The maximum number of tokens for model output, maximum output is 4095, default value is 1024.",
                    },
                ),
                "temperature": (
                    "FLOAT",
                    {
                        "default": 0.8,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Sampling temperature, controls the randomness of the output, must be a positive number within the range: [0.0, 1.0], default value is 0.95.",
                    },
                ),
                "top_p": (
                    "FLOAT",
                    {
                        "default": 0.6,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Another method of temperature sampling, value range is: [0.0, 1.0], default value is 0.7.",
                    },
                ),
                "instruct": (
                    "STRING",
                    {
                        "multiline": True,
                        "placeholder": "Input instruct text",
                        "default": "What is shown in the picture?",
                        "tooltip": "Enter the instruction for the neural network",
                    },
                ),
            }
        }

    RETURN_TYPES = ("STRING",)
    FUNCTION = "chatglm_instruct_media"

    CATEGORY = "AlekPet Nodes/Instruct"

    def chatglm_instruct_media(
        self, model, max_tokens, temperature, top_p, instruct, image=None, video=""
        # self, model, max_tokens, temperature, top_p, instruct, image=None, video=""
    ):

        if instruct is None or instruct.strip() == "":
            raise ValueError("Instruct text is empty!")

        # video = video.strip()

        # if image is None and (video is None and video == ""):
        #     raise ValueError("Image or Video path is empty!")

        if image is not None:
            if video != "":
                raise ValueError("You cannot use both an image and a video at the same time!")           

        answer = ""
        payload = {}
        if image is not None:
            img = 255.0 * image.cpu().numpy()
            img = np.squeeze(img)
            img = Image.fromarray(np.clip(img, 0, 255).astype(np.uint8))
            img = toBase64ImgUrl(img)

            # Create body request for image
            payload = {
                "model": model,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "image_url", "image_url": {"url": img}},
                            {"type": "text", "text": "What is shown in the picture?"},
                        ],
                    }
                ],
                "max_tokens": round(max_tokens, 2),
                "temperature": round(temperature, 2),
                "top_p": round(top_p, 2),
            }

        # if video:
        #     # Create body request for video
        #     address = PromptServer.instance.address
        #     port = PromptServer.instance.port
        #     url_video = f"http://{address}:{port}/view?filename={video}&type=input&subfolder="

        #     payload = {
        #         "model": model,
        #         "messages": [
        #             {
        #                 "role": "user",
        #                 "content": [
        #                     {"type": "video_url", "video_url": {"url": url_video}},
        #                     {"type": "text", "text": "Describe this video"},
        #                 ],
        #             }
        #         ],
        #         "max_tokens": round(max_tokens, 2),
        #         "temperature": round(temperature, 2),
        #         "top_p": round(top_p, 2),
        #     }

        answer = createRequest(payload)

        return (answer,)
