import json
import os
from collections.abc import Mapping, Sequence
from enum import StrEnum, auto
from typing import TYPE_CHECKING, Any, TypedDict, cast

from core.app.app_config.entities import PromptTemplateEntity
from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEntity
from core.memory.token_buffer_memory import TokenBufferMemory
from core.prompt.entities.advanced_prompt_entities import MemoryConfig
from core.prompt.prompt_transform import PromptTransform
from core.prompt.utils.prompt_template_parser import PromptTemplateParser
from graphon.file import file_manager
from graphon.model_runtime.entities.message_entities import (
    ImagePromptMessageContent,
    PromptMessage,
    PromptMessageContentUnionTypes,
    SystemPromptMessage,
    TextPromptMessageContent,
    UserPromptMessage,
)
from models.model import AppMode

if TYPE_CHECKING:
    from graphon.file import File


class ModelMode(StrEnum):
    COMPLETION = auto()
    CHAT = auto()


prompt_file_contents: dict[str, Any] = {}


class PromptTemplateConfigDict(TypedDict):
    prompt_template: PromptTemplateParser
    custom_variable_keys: list[str]
    special_variable_keys: list[str]
    prompt_rules: dict[str, Any]


class SimplePromptTransform(PromptTransform):
    """
    Simple Prompt Transform for Chatbot App Basic Mode.
    """

    def get_prompt(
        self,
        app_mode: AppMode,
        prompt_template_entity: PromptTemplateEntity,
        inputs: Mapping[str, str],
        query: str,
        files: Sequence["File"],
        context: str | None,
        memory: TokenBufferMemory | None,
        model_config: ModelConfigWithCredentialsEntity,
        image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
        context_files: list["File"] | None = None,
    ) -> tuple[list[PromptMessage], list[str] | None]:
        inputs = {key: str(value) for key, value in inputs.items()}

        model_mode = ModelMode(model_config.mode)
        if model_mode == ModelMode.CHAT:
            prompt_messages, stops = self._get_chat_model_prompt_messages(
                app_mode=app_mode,
                pre_prompt=prompt_template_entity.simple_prompt_template or "",
                inputs=inputs,
                query=query,
                files=files,
                context=context,
                memory=memory,
                model_config=model_config,
                image_detail_config=image_detail_config,
                context_files=context_files,
            )
        else:
            prompt_messages, stops = self._get_completion_model_prompt_messages(
                app_mode=app_mode,
                pre_prompt=prompt_template_entity.simple_prompt_template or "",
                inputs=inputs,
                query=query,
                files=files,
                context=context,
                memory=memory,
                model_config=model_config,
                image_detail_config=image_detail_config,
                context_files=context_files,
            )

        return prompt_messages, stops

    def _get_prompt_str_and_rules(
        self,
        app_mode: AppMode,
        model_config: ModelConfigWithCredentialsEntity,
        pre_prompt: str,
        inputs: dict[str, Any],
        query: str | None = None,
        context: str | None = None,
        histories: str | None = None,
    ) -> tuple[str, dict[str, Any]]:
        # get prompt template
        prompt_template_config = self.get_prompt_template(
            app_mode=app_mode,
            provider=model_config.provider,
            model=model_config.model,
            pre_prompt=pre_prompt,
            has_context=context is not None,
            query_in_prompt=query is not None,
            with_memory_prompt=histories is not None,
        )

        custom_variable_keys = prompt_template_config["custom_variable_keys"]
        if not isinstance(custom_variable_keys, list):
            raise TypeError(f"Expected list for custom_variable_keys, got {type(custom_variable_keys)}")

        special_variable_keys = prompt_template_config["special_variable_keys"]
        if not isinstance(special_variable_keys, list):
            raise TypeError(f"Expected list for special_variable_keys, got {type(special_variable_keys)}")

        variables = {k: inputs[k] for k in custom_variable_keys if k in inputs}

        for v in special_variable_keys:
            # support #context#, #query# and #histories#
            match v:
                case "#context#":
                    variables["#context#"] = context or ""
                case "#query#":
                    variables["#query#"] = query or ""
                case "#histories#":
                    variables["#histories#"] = histories or ""
                case _:
                    pass

        prompt_template = prompt_template_config["prompt_template"]
        if not isinstance(prompt_template, PromptTemplateParser):
            raise TypeError(f"Expected PromptTemplateParser, got {type(prompt_template)}")

        prompt = prompt_template.format(variables)

        prompt_rules = prompt_template_config["prompt_rules"]
        if not isinstance(prompt_rules, dict):
            raise TypeError(f"Expected dict for prompt_rules, got {type(prompt_rules)}")

        return prompt, prompt_rules

    def get_prompt_template(
        self,
        app_mode: AppMode,
        provider: str,
        model: str,
        pre_prompt: str,
        has_context: bool,
        query_in_prompt: bool,
        with_memory_prompt: bool = False,
    ) -> PromptTemplateConfigDict:
        prompt_rules = self._get_prompt_rule(app_mode=app_mode, provider=provider, model=model)

        custom_variable_keys: list[str] = []
        special_variable_keys: list[str] = []

        prompt = ""
        for order in prompt_rules["system_prompt_orders"]:
            if order == "context_prompt" and has_context:
                prompt += prompt_rules["context_prompt"]
                special_variable_keys.append("#context#")
            elif order == "pre_prompt" and pre_prompt:
                prompt += pre_prompt + "\n"
                pre_prompt_template = PromptTemplateParser(template=pre_prompt)
                custom_variable_keys = pre_prompt_template.variable_keys
            elif order == "histories_prompt" and with_memory_prompt:
                prompt += prompt_rules["histories_prompt"]
                special_variable_keys.append("#histories#")

        if query_in_prompt:
            prompt += prompt_rules.get("query_prompt", "{{#query#}}")
            special_variable_keys.append("#query#")

        result: PromptTemplateConfigDict = {
            "prompt_template": PromptTemplateParser(template=prompt),
            "custom_variable_keys": custom_variable_keys,
            "special_variable_keys": special_variable_keys,
            "prompt_rules": prompt_rules,
        }
        return result

    def _get_chat_model_prompt_messages(
        self,
        app_mode: AppMode,
        pre_prompt: str,
        inputs: dict[str, Any],
        query: str,
        context: str | None,
        files: Sequence["File"],
        memory: TokenBufferMemory | None,
        model_config: ModelConfigWithCredentialsEntity,
        image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
        context_files: list["File"] | None = None,
    ) -> tuple[list[PromptMessage], list[str] | None]:
        prompt_messages: list[PromptMessage] = []

        # get prompt
        prompt, _ = self._get_prompt_str_and_rules(
            app_mode=app_mode,
            model_config=model_config,
            pre_prompt=pre_prompt,
            inputs=inputs,
            query=None,
            context=context,
        )

        if prompt and query:
            prompt_messages.append(SystemPromptMessage(content=prompt))

        if memory:
            prompt_messages = self._append_chat_histories(
                memory=memory,
                memory_config=MemoryConfig(
                    window=MemoryConfig.WindowConfig(
                        enabled=False,
                    )
                ),
                prompt_messages=prompt_messages,
                model_config=model_config,
            )

        if query:
            prompt_messages.append(self._get_last_user_message(query, files, image_detail_config, context_files))
        else:
            prompt_messages.append(self._get_last_user_message(prompt, files, image_detail_config, context_files))

        return prompt_messages, None

    def _get_completion_model_prompt_messages(
        self,
        app_mode: AppMode,
        pre_prompt: str,
        inputs: dict[str, Any],
        query: str,
        context: str | None,
        files: Sequence["File"],
        memory: TokenBufferMemory | None,
        model_config: ModelConfigWithCredentialsEntity,
        image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
        context_files: list["File"] | None = None,
    ) -> tuple[list[PromptMessage], list[str] | None]:
        # get prompt
        prompt, prompt_rules = self._get_prompt_str_and_rules(
            app_mode=app_mode,
            model_config=model_config,
            pre_prompt=pre_prompt,
            inputs=inputs,
            query=query,
            context=context,
        )

        if memory:
            tmp_human_message = UserPromptMessage(content=prompt)

            rest_tokens = self._calculate_rest_token([tmp_human_message], model_config=model_config)
            histories = self._get_history_messages_from_memory(
                memory=memory,
                memory_config=MemoryConfig(
                    window=MemoryConfig.WindowConfig(
                        enabled=False,
                    )
                ),
                max_token_limit=rest_tokens,
                human_prefix=prompt_rules.get("human_prefix", "Human"),
                ai_prefix=prompt_rules.get("assistant_prefix", "Assistant"),
            )

            # get prompt
            prompt, prompt_rules = self._get_prompt_str_and_rules(
                app_mode=app_mode,
                model_config=model_config,
                pre_prompt=pre_prompt,
                inputs=inputs,
                query=query,
                context=context,
                histories=histories,
            )

        stops = prompt_rules.get("stops")
        if stops is not None and len(stops) == 0:
            stops = None

        return [self._get_last_user_message(prompt, files, image_detail_config, context_files)], stops

    def _get_last_user_message(
        self,
        prompt: str,
        files: Sequence["File"],
        image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
        context_files: list["File"] | None = None,
    ) -> UserPromptMessage:
        prompt_message_contents: list[PromptMessageContentUnionTypes] = []
        if files:
            for file in files:
                prompt_message_contents.append(
                    file_manager.to_prompt_message_content(file, image_detail_config=image_detail_config)
                )
        if context_files:
            for file in context_files:
                prompt_message_contents.append(
                    file_manager.to_prompt_message_content(file, image_detail_config=image_detail_config)
                )
        if prompt_message_contents:
            prompt_message_contents.append(TextPromptMessageContent(data=prompt))

            prompt_message = UserPromptMessage(content=prompt_message_contents)
        else:
            prompt_message = UserPromptMessage(content=prompt)

        return prompt_message

    def _get_prompt_rule(self, app_mode: AppMode, provider: str, model: str) -> dict[str, Any]:
        """
        Get simple prompt rule.
        :param app_mode: app mode
        :param provider: model provider
        :param model: model name
        :return:
        """
        prompt_file_name = self._prompt_file_name(app_mode=app_mode, provider=provider, model=model)

        # Check if the prompt file is already loaded
        if prompt_file_name in prompt_file_contents:
            return cast(dict[str, Any], prompt_file_contents[prompt_file_name])

        # Get the absolute path of the subdirectory
        prompt_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "prompt_templates")
        json_file_path = os.path.join(prompt_path, f"{prompt_file_name}.json")

        # Open the JSON file and read its content
        with open(json_file_path, encoding="utf-8") as json_file:
            content = json.load(json_file)

            # Store the content of the prompt file
            prompt_file_contents[prompt_file_name] = content

            return cast(dict[str, Any], content)

    def _prompt_file_name(self, app_mode: AppMode, provider: str, model: str) -> str:
        # baichuan
        is_baichuan = False
        if provider == "baichuan":
            is_baichuan = True
        else:
            baichuan_supported_providers = ["huggingface_hub", "openllm", "xinference"]
            if provider in baichuan_supported_providers and "baichuan" in model.lower():
                is_baichuan = True

        if is_baichuan:
            if app_mode == AppMode.COMPLETION:
                return "baichuan_completion"
            else:
                return "baichuan_chat"

        # common
        if app_mode == AppMode.COMPLETION:
            return "common_completion"
        else:
            return "common_chat"
