from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate
from langchain.chains.qa_with_sources import load_qa_with_sources_chain
from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage
import tiktoken


def num_token(string: str) -> int:
    """Returns the number of tokens in a text string."""
    encoding = tiktoken.get_encoding('cl100k_base')
    num_tokens = len(encoding.encode(string))
    return num_tokens


def retrive_doc_on_token(rephrased_query, db):
    top_docs = db.similarity_search(rephrased_query, k=10)
    num = [num_token(doc.page_content) for doc in top_docs]
    cum_sum = [abs(sum(num[:i+1]) - 2700) for i in range(len(num))]
    idx = cum_sum.index(min(cum_sum))

    return top_docs[:idx+1]




def multi_docs_qa(query, rephrased_query, db, company, language, temperature):
    """
    Return an answer to the query based on multiple documents limited
    by total token of 3000.
    """
    print('temperature: ', temperature)
    template = """<|im_start|>
    Manulife's assistant helps the user with their questions about products and services. Be brief in your answers with appropriate tone and emotion. Answer ONLY with the facts listed in the list of sources below. If there isn't enough information below, say you don't know. Do not generate answers that don't use the sources below. If asking a clarifying question to the user would help, ask the question. For tabular information return it as an html table. Do not return markdown format. Each source has a name followed by colon and the actual information, ALWAYS include the source name for each fact you use in the response. Use square brakets to reference the source, e.g. [info1.txt]. Don't combine sources, list each source separately, e.g. [info1.txt][info2.pdf].

    Sources:
    {sources}
    <|im_end|>

    {chat_history}
    """

    docs =  retrive_doc_on_token(query, db)
    sources = []
    for i in docs:
        source_txt = i.metadata['source']
        source_content = i.page_content
        add = f"{source_txt}: {source_content}"
        sources.append(add)

    s_txt = '\n\n'.join(sources)
    # print('this is docs: ', docs)
    ch = rephrased_query + f'\nUser: {query}' + '\nAssistant: '
    final_template = template.format(sources=s_txt, chat_history=ch)
    print('\n\n', final_template, '\n\n')
    # PROMPT = PromptTemplate(template=template, input_variables=["summaries", "question"])
    # llm = OpenAI(temperature=temperature, model_name='gpt-3.5-turbo')
    # chain = load_qa_with_sources_chain(llm, chain_type="stuff", prompt=PROMPT)
    # response = chain({"input_documents": docs, "question": query})

    llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=temperature)
    response = llm([HumanMessage(content=final_template)]).content


    # return response['output_text'], docs
    return response, ''



def multi_docs_qa_hkas(query, rephrased_query, db, language, temperature):
    """
    Return an answer to the query based on multiple documents limited
    by total token of 3000.
    """

    template = (
        f"Create a comprehensive and truthful final response in {language}. "
        "Ask for clearification before answering if the QUESTION is not clear.\n\n"

        "Context (may or may not be useful)"
        "===\n"
        "{summaries}\n"
        "===\n\n"

        "Query:  "
        "===\n"
        "{question} "
        f"({rephrased_query})\n"
        "===\n\n"

        f"FINAL RESPONSE (in complete sentence):"
    )
    docs =  retrive_doc_on_token(query+ f" ({rephrased_query})", db)


    PROMPT = PromptTemplate(template=template, input_variables=["summaries", "question"])
    llm = OpenAI(temperature=temperature, model_name='gpt-3.5-turbo')
    chain = load_qa_with_sources_chain(llm, chain_type="stuff", prompt=PROMPT)
    response = chain({"input_documents": docs, "question": query})

    return response['output_text'], docs
    # return response['output_text'].replace('Manulife', 'Company A').replace('manulife', 'Company A'), docs



def summary(query, context):
    template = (
        "Use the following portion of a long document to see if any of the text is relevant to answer the question. "
        "Return any relevant text verbatim and 'SOURCE'.\n\n"

        "===\n"
        "{context}"
        "===\n\n"

        "Question: {question}\n"
        "Relevant text, if any:"
    )
    return template.format(context=context,question=query)

import asyncio

async def multi_reponse(temperature, messages, docs):
    chat = ChatOpenAI(temperature=temperature, model_name="gpt-3.5-turbo")
    responses = await chat.agenerate(messages=messages)

    text = ''

    for i, r in enumerate(responses.generations):
        if 'N/A' not in r[0].text:
            text += f"{r[0].text}\nSOURCE: {docs[i].metadata['source']}\n\n"

    print(text)

    return text


from typing import Any, Dict, List, Optional, Union
from types import GeneratorType
from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler
from langchain.schema import AgentAction, AgentFinish, LLMResult

class SyncStreamingLLMCallbackHandler(BaseCallbackHandler):
    """Callback handler for streaming LLM responses to a queue."""

    def __init__(self, q):
        self.q = q

    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        self.q.put(token)

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
        """Do nothing."""
        pass

    def on_llm_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> None:
        """Do nothing."""
        pass

    def on_chain_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        **kwargs: Any,
    ) -> None:
        """Do nothing."""
        pass

    def on_tool_end(
        self,
        output: str,
        color: Optional[str] = None,
        observation_prefix: Optional[str] = None,
        llm_prefix: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        """Do nothing."""
        pass

    def on_tool_error(
        self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> None:
        """Do nothing."""
        pass

    def on_agent_action(self, action: AgentAction, **kwargs: Any) -> Any:
        """Run on agent action."""
        pass

    def on_agent_finish(
        self, finish: AgentFinish, color: Optional[str] = None, **kwargs: Any
    ) -> None:
        """Run on agent end."""
        pass