import os
from abc import abstractmethod, ABC
from typing import Literal, Optional

from langchain_core.prompts import ChatPromptTemplate
from langfuse import get_client
from openkosmos_core.common.model import BaseConfig
from openkosmos_core.templatex import TemplateLoader

TEMPLATE_SOURCE = Literal["local", "langfuse", "langsmith"]
TEMPLATE_FORMAT = Literal["f-string", "mustache", "jinja2", "templatex"]

PROMPT_ROLE = Literal["system", "ai", "human"]


class AITemplateConfig(BaseConfig, ABC):
    source: TEMPLATE_SOURCE = "local"
    format: TEMPLATE_FORMAT = "mustache"

    @abstractmethod
    def get_template(self, name: str, role: PROMPT_ROLE):
        pass

    @abstractmethod
    def save_template(self, name: str, content: str):
        pass

    @abstractmethod
    def get_template_content(self, name: str):
        pass


class LocalTemplateConfig(AITemplateConfig):
    base_dir: str = "./templates"
    ext: str = ".tpl"

    def get_local_template(self, name: str):
        with open(f"{self.base_dir}/{name}{self.ext}", "r", encoding="utf8") as template_file:
            return template_file.read()

    def get_template_content(self, name: str):
        return self.get_local_template(name)

    def get_template(self, name: str, role: PROMPT_ROLE):
        prompt = self.get_local_template(name)
        chat_prompt = ChatPromptTemplate.from_messages(
            messages=[(role, prompt)],
            template_format=self.format)

        return chat_prompt

    def save_template(self, name: str, content: str):
        os.makedirs(f"{self.base_dir}/{name[0:name.rindex("/")]}", exist_ok=True)
        with open(f"{self.base_dir}/{name}{self.ext}", "w",
                  encoding="utf-8") as output_file:
            output_file.write(content)


class LangfuseTemplateConfig(AITemplateConfig):
    type: Literal["chat", "text"] = "text"
    label: Optional[str] = "production"
    cache: Optional[int] = 86400
    version: Optional[int] = None

    def get_langfuse_template(self, name: str):
        return get_client().get_prompt(name, type=self.type, label=self.label,
                                       cache_ttl_seconds=self.cache,
                                       version=self.version)

    def get_template_content(self, name: str):
        return self.get_langfuse_template(name).prompt

    def get_template(self, name: str, role: PROMPT_ROLE):
        langfuse_prompt = self.get_langfuse_template(name)
        chat_prompt = ChatPromptTemplate.from_messages(
            messages=[(role, langfuse_prompt.prompt)],
            template_format=self.format
        )
        chat_prompt.metadata = {"langfuse_prompt": langfuse_prompt}

        return chat_prompt

    def save_template(self, name: str, content: str):
        get_client().create_prompt(name=name, prompt=content, type=self.type, labels=[self.label])


class LangsmithTemplateConfig(AITemplateConfig, ABC):
    pass


class AITemplateLoader(TemplateLoader):
    DEFAULT_TEMPLATE_SOURCE: TEMPLATE_SOURCE = "local"

    def __init__(self, config: TEMPLATE_SOURCE | AITemplateConfig = DEFAULT_TEMPLATE_SOURCE):
        if isinstance(config, str):
            match config:
                case "local":
                    self._config = LocalTemplateConfig()
                case "langfuse":
                    self._config = LangfuseTemplateConfig()
                case "langsmith":
                    self._config = LangsmithTemplateConfig()
                case _:
                    raise NotImplementedError(config)
        else:
            self._config = config

    def config(self) -> AITemplateConfig:
        return self._config

    def get_prompt(self, name: str, role: PROMPT_ROLE = "system") -> ChatPromptTemplate:
        return self.config().get_template(name, role)

    def transfer(self, name: str, target: AITemplateConfig):
        target.save_template(name, self.config().get_template_content(name))
