import pinecone
from environs import Env
from observability import span
from vectorstore import PineconeStore
from revise import load_revise_chain
from response import load_response_chain
from langchain.chains.base import Chain
from typing import List, Dict, Optional, Any
from toxicity import load_toxicity_check_chain
from langchain.embeddings.base import Embeddings
from langchain.base_language import BaseLanguageModel
from langchain.callbacks.manager import (
    AsyncCallbackManagerForChainRun,
    CallbackManagerForChainRun,
)


class PingChatChain(Chain):
    revise_chain: Chain
    toxicity_chain: Chain
    response_chain: Chain
    
    @property
    def input_keys(self) -> List[str]:
        return ['messages']
    
    @property
    def output_keys(self) -> List[str]:
        return [
            'content',
            'links',
            'recalls',
            'instruction',
        ]

    @span(name="pingchat chain", input_fn=lambda args, kwargs: args[0])
    def _call(
        self,
        inputs: Dict[str, Any], /,
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, str]:
        messages = inputs['messages']
        revised_query = span(
            name="revise chain",
        )(self.revise_chain.run)(
            messages=messages,
            callbacks=run_manager.get_child() if run_manager else None,
        )

        toxicity = span(
            name="toxicity chain",
        )(self.toxicity_chain.predict_and_parse)(
            instruction=revised_query,
            stop=['\n'],
            callbacks=run_manager.get_child() if run_manager else None,
        )
        
        r = self.response_chain(inputs={
            'revised_query': revised_query,
            'toxicity': toxicity,
            'messages': messages,
        }, callbacks=run_manager.get_child() if run_manager else None)
        
        return {
            'content': r['content'],
            'links': r['links'],
            'recalls': r['recalls'],
            'instruction': revised_query,
        }

    @span(name="pingchat chain", input_fn=lambda args, kwargs: args[0])
    async def _acall(
        self,
        inputs: Dict[str, Any], /,
        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
    ) -> Dict[str, str]:
        messages = inputs['messages']
        revised_query = await span(
            name="revise chain",
            input_fn = lambda args, kwargs: kwargs["messages"],
        )(self.revise_chain.arun)(
            messages=messages,
            callbacks=run_manager.get_child() if run_manager else None,
        )
        
        toxicity = await span(
            name="toxicity chain",
            input_fn = lambda args, kwargs: kwargs["instruction"],
        )(self.toxicity_chain.apredict_and_parse)(
            instruction=revised_query,
            stop=['\n'],
            callbacks=run_manager.get_child() if run_manager else None,
        )

        r = await self.response_chain.acall({
            'revised_query': revised_query,
            'toxicity': toxicity,
            'messages': messages,
        }, callbacks=run_manager.get_child() if run_manager else None)
        
        return {
            'content': r['content'],
            'links': r['links'],
            'recalls': r['recalls'],
            'instruction': revised_query,
        }
    
def load_pingchat_chain(env: Env, llm: BaseLanguageModel, embedding: Embeddings, doc_ns: List[str]) -> Chain:
    return PingChatChain(
        revise_chain=load_revise_chain(env, llm),
        toxicity_chain=load_toxicity_check_chain(env, llm),
        response_chain=load_response_chain(
            env,
            llm,
            PineconeStore(env, embedding, doc_ns).as_retriever(),
        ),
    )