import json
import re

import langchain.chains
from langchain.base_language import BaseLanguageModel
from langchain.chains.openai_functions import create_openai_fn_chain
from langchain.chat_models.base import BaseChatModel
from langchain.prompts import StringPromptTemplate
from langchain.prompts.chat import BaseChatPromptTemplate
from langchain.schema import AIMessage, HumanMessage

from classes.derived_str import JsonString
from scheduler.block import BaseBlock
from utils.usage import with_openai_stats


class LLMChain(BaseBlock):
    def __init__(self, model: BaseLanguageModel, prompt_template: StringPromptTemplate):
        self.chain = langchain.chains.LLMChain(llm=model, prompt=prompt_template)

    @with_openai_stats
    def __call__(self, text: str, **kwargs) -> str:
        return self.chain.run(text=text, **kwargs)


class ChatLLMChain(BaseBlock):
    def __init__(self, model: BaseChatModel, prompt_template: BaseChatPromptTemplate):
        self.chat = model
        self.prompt = prompt_template

    @with_openai_stats
    def __call__(self, messages: list, **kwargs) -> str:
        ms = []
        for i, m in enumerate(messages):
            role = "asistant"
            if i % 2 == 0:
                ms.append(HumanMessage(content=m))
            else:
                ms.append(AIMessage(content=m))
        prompt_value = self.prompt.format_prompt(messages=ms, **kwargs)
        response = self.chat.generate_prompt([prompt_value])
        return response.generations[0][0].text


class RawPromptChain(BaseBlock):
    def __init__(self, prompt_template: StringPromptTemplate):
        self.prompt = prompt_template

    def __call__(self, text: str, **kwargs) -> str:
        return self.prompt.format(text=text, **kwargs)


class LLMFunctionCall(BaseBlock):
    def __init__(
        self,
        model: BaseChatModel,
        prompt_template: StringPromptTemplate,
        functions: str,
    ):  # JsonString):
        self.chat = model
        self.prompt = prompt_template
        # self.functions = get_functions_from_str(functions)
        self.functions = get_functions_from_str(JsonString(functions))
        self.chain = create_openai_fn_chain(
            self.functions, self.chat, self.prompt, verbose=False
        )

    @with_openai_stats
    def __call__(self, text: str, **kwargs) -> JsonString:
        return JsonString(json.dumps(self.chain.run(text=text, **kwargs)))


class VisionChain(BaseBlock):
    def __init__(self, model: BaseChatModel, prompt_template: StringPromptTemplate):
        if not "vision" in model.model_name:
            raise ValueError("Model must be a vision model.")
        self.chat = model
        self.prompt = prompt_template

    @with_openai_stats
    def __call__(self, text: str, image: str, **kwargs) -> str:
        pattern = re.compile(r"data:\s*image\s*/\s*(\w+)\s*;\s*base64,")
        if not (image.startswith("http") or pattern.match(image)):
            raise ValueError(f"Invalid image. {image[:50]}...")
        req = self.prompt.format(text=text, **kwargs)
        response = self.chat.invoke(
            [
                HumanMessage(
                    content=[
                        {"type": "text", "text": req},
                        {
                            "type": "image_url",
                            "image_url": {"url": image, "detail": "auto"},
                        },
                    ]
                )
            ]
        )
        return response.content


def validate_function(func: dict) -> bool:
    def validate_name(name):
        if not isinstance(name, str):
            return False
        if len(name) > 64:
            return False
        return bool(re.match(r"^[\w-]+$", name))

    def validate_parameters(parameters):
        if not isinstance(parameters, dict):
            return False
        if parameters.get("type") != "object":
            return False
        if not isinstance(parameters.get("properties"), dict):
            return False
        if not isinstance(parameters.get("required"), list):
            return False
        if not set(parameters["required"]).issubset(
            set(parameters["properties"].keys())
        ):
            return False
        # It is not necessary to check the keys and values of properties.values()
        return True

    keys = func.keys()
    if "name" not in keys or "parameters" not in keys:
        raise ValueError(f"Invalid OpenAI function: Name and parameters are required.")
    if not validate_name(func["name"]):
        raise ValueError(
            f"Invalid OpenAI function name. Please check https://platform.openai.com/docs/api-reference/chat/create#chat-create-tools"
        )
    if not validate_parameters(func["parameters"]):
        raise ValueError(
            f"Invalid OpenAI function parameters. Please check https://platform.openai.com/docs/api-reference/chat/create#chat-create-tools"
        )
    return True


def get_functions_from_str(funcs: JsonString):
    functions = funcs.parsed()
    if isinstance(functions, list):
        for func in functions:
            if not validate_function(func):
                raise ValueError(f"Invalid OpenAI function: {func}")
        return functions
    else:
        raise ValueError("Functions must be an array of object(s).")