import pinecone
from pathlib import Path
from environs import Env
from observability import span
from langchain.chains import LLMChain
from langchain.chains.base import Chain
from typing import Callable, Dict, List
from example import EmbeddingExampleSelector
from langchain.embeddings.base import Embeddings
from langchain.base_language import BaseLanguageModel
from langchain.prompts import StringPromptTemplate, PromptTemplate
from langchain.prompts.example_selector.base import BaseExampleSelector


class RevisePromptTemplate(StringPromptTemplate):
    example_selector: BaseExampleSelector = None
    example_prompt: PromptTemplate = None
        
    def __init__(self, embedding_endpoint: str, namespace: str):
        super(RevisePromptTemplate, self).__init__(
            input_variables=["messages"]
        )
        self.example_selector = EmbeddingExampleSelector(embedding_endpoint, namespace)
        self.example_prompt = PromptTemplate.from_template(Path("prompts/revise_example.tpl").read_text())
    
    def get_context_from_messages(self, messages: List[dict]) -> str:
        return '\n'.join([f'{m["role"]}: {m["content"]}' for m in messages])
    
    @span(name="revise prompt")
    def format(self, **kwargs) -> str:
        messages = kwargs['messages']
        query = messages[-1]['content']
        prefix_path = "prompts/revise_prefix.tpl"
        suffix_path = "prompts/revise_suffix.tpl"
        context = self.get_context_from_messages(messages)
        if 'prefix' not in kwargs:
            kwargs['prefix'] = Path(prefix_path).read_text()
        if 'suffix' not in kwargs:
            kwargs['suffix'] = PromptTemplate.from_template(Path(suffix_path).read_text()).format(context=context)
        examples = self.example_selector.select_examples({'query': query})
        return (
            kwargs['prefix'] + '\n' + 
            '\n\n'.join([self.example_prompt.format(
                context=e['C'],
                instruction=e['I'],
            ) for e in examples]) + '\n\n' + 
            kwargs['suffix']
        )
    
def load_revise_chain(env: Env, llm: BaseLanguageModel) -> Chain:
    prompt = RevisePromptTemplate(env.str('EMBEDDING_SERVICE_ENDPOINT'), 'chatbot_revise')
    return LLMChain(llm=llm, prompt=prompt)