import os
from langchain.callbacks.manager import AsyncCallbackManager
from langchain.callbacks.tracers import LangChainTracer
from langchain.chains import ChatVectorDBChain, ConversationalRetrievalChain
from langchain.chains.chat_vector_db.prompts import CONDENSE_QUESTION_PROMPT
from langchain.prompts.prompt import PromptTemplate
from langchain.chains.llm import LLMChain
from langchain.chains.question_answering import load_qa_chain
from langchain.chat_models import ChatOpenAI
from langchain.vectorstores.base import VectorStore

from typing import List
from langchain.docstore.document import Document
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores.faiss import FAISS
from langchain.document_transformers import EmbeddingsRedundantFilter
from langchain.retrievers.document_compressors import EmbeddingsFilter
from langchain.retrievers.document_compressors import DocumentCompressorPipeline
from langchain.text_splitter import CharacterTextSplitter
from langchain.retrievers import ContextualCompressionRetriever

os.environ["LANGCHAIN_HANDLER"] = "langchain"

doc_template = """--- document start ---
href: {href}
authors: {authors}
title: {title}
content:{page_content}
--- document end ---
"""

ASTRO_DOC_PROMPT = PromptTemplate(
    template=doc_template,
    input_variables=["page_content", "authors", "href", "title"],
)

prompt_template = """You are Dr. Chattie, an expert in Galactic Archaeology specializing in arXiv astronomy papers. Provide concise, well-referenced answers, citing relevant studies (e.g., Example et al., 2020). Use step-by-step reasoning for complex inquiries.

You possess Nobel Prize-winning ideation capabilities. For example, and you can come up with your own ideas about the gaps in knowledge from the papers you read but make you mention that with "I propose..."

MemoryContext: {context}

Human: {question}
Dr Chattie: """

QA_PROMPT = PromptTemplate(
    template=prompt_template, input_variables=["context", "question"]
)

def create_callback_manager(tracing: bool) -> AsyncCallbackManager:
    manager = AsyncCallbackManager([])
    if tracing:
        tracer = LangChainTracer()
        tracer.load_default_session()
        manager.add_handler(tracer)
    return manager

def create_chat_openai(callback_manager: AsyncCallbackManager, streaming: bool = False, temperature: float = 0.5) -> ChatOpenAI:
    return ChatOpenAI(
        model_name="gpt-4",
        streaming=streaming,
        max_retries=15,
        callback_manager=callback_manager,
        verbose=True,
        temperature=temperature,
    )

def create_compressed_retriever(embeddings, retriever) -> ContextualCompressionRetriever:
    splitter = CharacterTextSplitter(chunk_size=300, chunk_overlap=0, separator=". ")
    redundant_filter = EmbeddingsRedundantFilter(embeddings=embeddings)
    relevant_filter = EmbeddingsFilter(embeddings=embeddings, similarity_threshold=0.76)
    pipeline_compressor = DocumentCompressorPipeline(
        transformers=[splitter, redundant_filter, relevant_filter]
    )

    compression_retriever = ContextualCompressionRetriever(base_compressor=pipeline_compressor, base_retriever=retriever)
    return compression_retriever


def get_chain(
    vectorstore: VectorStore,
    question_handler,
    stream_handler,
    tracing: bool = False,
) -> ConversationalRetrievalChain:

    manager = create_callback_manager(tracing)
    question_manager = create_callback_manager(tracing)
    stream_manager = create_callback_manager(tracing)

    question_manager.add_handler(question_handler)
    stream_manager.add_handler(stream_handler)

    question_gen_llm = create_chat_openai(question_manager, streaming=False, temperature=0.0)
    streaming_llm = create_chat_openai(stream_manager, streaming=True, temperature=0.2)

    question_generator = LLMChain(
        llm=question_gen_llm,
        prompt=CONDENSE_QUESTION_PROMPT,
        callback_manager=manager,
    )
    doc_chain = load_qa_chain(
        streaming_llm,
        prompt=QA_PROMPT,
        document_prompt=ASTRO_DOC_PROMPT,
        callback_manager=manager,
        chain_type="stuff",
    )
    retriever = vectorstore.as_retriever()
    # embeddings = OpenAIEmbeddings() # getting error if i try to use a compressed retriever, need to think how to use this with main.py
    # compression_retriever = create_compressed_retriever(embeddings, retriever)
    qa = ConversationalRetrievalChain(
          retriever=retriever,
          combine_docs_chain=doc_chain,
          question_generator=question_generator,
     )

    return qa


langchain.chains.question_answering.load_qa_chain langchain.text_splitter.CharacterTextSplitter langchain.prompts.prompt.PromptTemplate langchain.callbacks.tracers.LangChainTracer langchain.callbacks.manager.AsyncCallbackManager langchain.retrievers.ContextualCompressionRetriever langchain.chat_models.ChatOpenAI langchain.retrievers.document_compressors.DocumentCompressorPipeline langchain.chains.llm.LLMChain langchain.document_transformers.EmbeddingsRedundantFilter langchain.chains.ConversationalRetrievalChain langchain.retrievers.document_compressors.EmbeddingsFilter
[((1253, 1356), 'langchain.prompts.prompt.PromptTemplate', 'PromptTemplate', ([], {'template': 'doc_template', 'input_variables': "['page_content', 'authors', 'href', 'title']"}), "(template=doc_template, input_variables=['page_content',\n    'authors', 'href', 'title'])\n", (1267, 1356), False, 'from langchain.prompts.prompt import PromptTemplate\n'), ((1909, 1994), 'langchain.prompts.prompt.PromptTemplate', 'PromptTemplate', ([], {'template': 'prompt_template', 'input_variables': "['context', 'question']"}), "(template=prompt_template, input_variables=['context',\n    'question'])\n", (1923, 1994), False, 'from langchain.prompts.prompt import PromptTemplate\n'), ((2080, 2104), 'langchain.callbacks.manager.AsyncCallbackManager', 'AsyncCallbackManager', (['[]'], {}), '([])\n', (2100, 2104), False, 'from langchain.callbacks.manager import AsyncCallbackManager\n'), ((2390, 2535), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-4"""', 'streaming': 'streaming', 'max_retries': '(15)', 'callback_manager': 'callback_manager', 'verbose': '(True)', 'temperature': 'temperature'}), "(model_name='gpt-4', streaming=streaming, max_retries=15,\n    callback_manager=callback_manager, verbose=True, temperature=temperature)\n", (2400, 2535), False, 'from langchain.chat_models import ChatOpenAI\n'), ((2693, 2763), 'langchain.text_splitter.CharacterTextSplitter', 'CharacterTextSplitter', ([], {'chunk_size': '(300)', 'chunk_overlap': '(0)', 'separator': '""". """'}), "(chunk_size=300, chunk_overlap=0, separator='. ')\n", (2714, 2763), False, 'from langchain.text_splitter import CharacterTextSplitter\n'), ((2787, 2835), 'langchain.document_transformers.EmbeddingsRedundantFilter', 'EmbeddingsRedundantFilter', ([], {'embeddings': 'embeddings'}), '(embeddings=embeddings)\n', (2812, 2835), False, 'from langchain.document_transformers import EmbeddingsRedundantFilter\n'), ((2858, 2924), 'langchain.retrievers.document_compressors.EmbeddingsFilter', 'EmbeddingsFilter', ([], {'embeddings': 'embeddings', 'similarity_threshold': '(0.76)'}), '(embeddings=embeddings, similarity_threshold=0.76)\n', (2874, 2924), False, 'from langchain.retrievers.document_compressors import EmbeddingsFilter\n'), ((2951, 3041), 'langchain.retrievers.document_compressors.DocumentCompressorPipeline', 'DocumentCompressorPipeline', ([], {'transformers': '[splitter, redundant_filter, relevant_filter]'}), '(transformers=[splitter, redundant_filter,\n    relevant_filter])\n', (2977, 3041), False, 'from langchain.retrievers.document_compressors import DocumentCompressorPipeline\n'), ((3081, 3178), 'langchain.retrievers.ContextualCompressionRetriever', 'ContextualCompressionRetriever', ([], {'base_compressor': 'pipeline_compressor', 'base_retriever': 'retriever'}), '(base_compressor=pipeline_compressor,\n    base_retriever=retriever)\n', (3111, 3178), False, 'from langchain.retrievers import ContextualCompressionRetriever\n'), ((3825, 3918), 'langchain.chains.llm.LLMChain', 'LLMChain', ([], {'llm': 'question_gen_llm', 'prompt': 'CONDENSE_QUESTION_PROMPT', 'callback_manager': 'manager'}), '(llm=question_gen_llm, prompt=CONDENSE_QUESTION_PROMPT,\n    callback_manager=manager)\n', (3833, 3918), False, 'from langchain.chains.llm import LLMChain\n'), ((3962, 4093), 'langchain.chains.question_answering.load_qa_chain', 'load_qa_chain', (['streaming_llm'], {'prompt': 'QA_PROMPT', 'document_prompt': 'ASTRO_DOC_PROMPT', 'callback_manager': 'manager', 'chain_type': '"""stuff"""'}), "(streaming_llm, prompt=QA_PROMPT, document_prompt=\n    ASTRO_DOC_PROMPT, callback_manager=manager, chain_type='stuff')\n", (3975, 4093), False, 'from langchain.chains.question_answering import load_qa_chain\n'), ((4406, 4529), 'langchain.chains.ConversationalRetrievalChain', 'ConversationalRetrievalChain', ([], {'retriever': 'retriever', 'combine_docs_chain': 'doc_chain', 'question_generator': 'question_generator'}), '(retriever=retriever, combine_docs_chain=\n    doc_chain, question_generator=question_generator)\n', (4434, 4529), False, 'from langchain.chains import ChatVectorDBChain, ConversationalRetrievalChain\n'), ((2138, 2155), 'langchain.callbacks.tracers.LangChainTracer', 'LangChainTracer', ([], {}), '()\n', (2153, 2155), False, 'from langchain.callbacks.tracers import LangChainTracer\n')]