from typing import Dict, List

import requests

# from langchain import PromptTemplate is deprecated.
from langchain.prompts import MessagesPlaceholder, PromptTemplate, StringPromptTemplate
from langchain.prompts.chat import ChatPromptTemplate, SystemMessagePromptTemplate
from langchain.prompts.example_selector.base import BaseExampleSelector

from exceptions import EmbeddingServiceError
from embedding import Namespace


class EmbeddingExampleSelector(BaseExampleSelector):
    def __init__(self, embedding_endpoint: str, namespace: str):
        super(EmbeddingExampleSelector, self).__init__()
        self.embedding_endpoint = embedding_endpoint
        self.namespace = namespace

    def add_example(self, example: Dict[str, str]) -> None:
        """Add new example to store for a key."""

        with requests.put(
            f"{self.embedding_endpoint}/{self.namespace}/upsert", json=example
        ) as r:
            if r.status_code != 200:
                raise EmbeddingServiceError(f"add example error: {r.text}")

    def select_examples(self, input_variables: Dict[str, str]) -> List[dict]:
        with requests.post(
            f"{self.embedding_endpoint}/{self.namespace}/query",
            json=input_variables,
        ) as r:
            x = r.json()
            if x.get("status") != "success":
                raise EmbeddingServiceError("embedding service call not success")
            return x["matches"]

class FewShotPromptTemplate(StringPromptTemplate):
    namespace: Namespace = None
    prefix: str = ''
    suffix: str = ''
    example_prompt: str = ''

    def __init__(
        self,
        namespace: Namespace,
        prefix: str,
        suffix: str,
        example_prompt: str,
    ):
        super(FewShotPromptTemplate, self).__init__(
            input_variables=list(
                set(
                    PromptTemplate.from_template(prefix).input_variables
                    + PromptTemplate.from_template(suffix).input_variables
                )
            ),
        )
        self.namespace = namespace
        self.prefix = prefix
        self.suffix = suffix
        self.example_prompt = example_prompt

    def format(self, text: str, **kwargs) -> str:
        kwargs["text"] = text
        examples = self.namespace.retrieve(text)
        examples.reverse()
        return self.prefix.format(**kwargs) + "\n" "\n".join(
            [self.example_prompt.format(**e) for e in examples]
        ) + "\n" + self.suffix.format(**kwargs)

class FewShotPromptTemplateOld(StringPromptTemplate):
    example_selector: BaseExampleSelector = None
    prefix: str
    suffix: str
    example_prompt: str

    def __init__(
        self,
        embedding_endpoint: str,
        namespace: str,
        prefix: str,
        suffix: str,
        example_prompt: str,
    ):
        super(FewShotPromptTemplateOld, self).__init__(
            input_variables=list(
                set(
                    PromptTemplate.from_template(prefix).input_variables
                    + PromptTemplate.from_template(suffix).input_variables
                )
            ),
            prefix=prefix,
            suffix=suffix,
            example_prompt=example_prompt,
        )
        self.example_selector = EmbeddingExampleSelector(embedding_endpoint, namespace)
        self.prefix = prefix
        self.suffix = suffix
        self.example_prompt = example_prompt

    def format(self, text: str, **kwargs) -> str:
        kwargs["text"] = text
        examples = self.example_selector.select_examples({"query": text})
        examples.reverse()
        return self.prefix.format(**kwargs) + "\n" "\n".join(
            [self.example_prompt.format(**e) for e in examples]
        ) + "\n" + self.suffix.format(**kwargs)


class ChatMessagePrompt(ChatPromptTemplate):
    def __init__(self, system_prompt: str):
        system_prompt = SystemMessagePromptTemplate.from_template(system_prompt)
        super(ChatMessagePrompt, self).__init__(
            input_variables=system_prompt.input_variables + ["messages"],
            messages=[
                system_prompt,
                MessagesPlaceholder(variable_name="messages"),
            ],
        )


class ZeroShotPromptTemplate(StringPromptTemplate):
    prompt_template: str

    def __init__(self, prompt_template: str):
        super(ZeroShotPromptTemplate, self).__init__(
            input_variables=PromptTemplate.from_template(
                prompt_template
            ).input_variables,
            prompt_template=prompt_template,
        )
        self.prompt_template = prompt_template

    def format(self, text: str = "", **kwargs) -> str:
        kwargs["text"] = text
        prompt = self.prompt_template.format(**kwargs)
        return prompt