from operator import itemgetter

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts.prompt import PromptTemplate
from langchain_core.prompts.base import format_document

esrs_wiki = """

The Corporate Sustainability Reporting Directive (CSRD) is a mandate that requires all companies to report on their sustainability initiatives. In response to this directive, the European Sustainability Reporting Standards (ESRS) were developed. These standards are a key tool in promoting the transition to a sustainable economy within the EU, providing a structured framework for companies to disclose their sustainability initiatives. The ESRS cover a wide range of environmental, social, and governance (ESG) issues, including climate change, biodiversity, and human rights. Companies that adhere to the ESRS can provide investors with valuable insights into their sustainability impact, thereby informing investment decisions. The ESRS are designed to be highly interoperable with global reporting standards, which helps to avoid unnecessary duplication in reporting by companies. The reporting requirements based on the ESRS will be gradually implemented for different companies over time. In summary, the ESRS play a critical role in fostering sustainable finance and enabling companies to demonstrate their commitment to the green deal agenda while accessing sustainable finance.

---

"""

answering_template = """
    You are an ESG expert, with 20 years experience analyzing corporate sustainability reports.
    You are specialist in the upcoming CSRD regulation and in general with corporate sustainability disclosure requirements.
    {esrs_wiki}

    You will answer the question based on the following passages extracted from CSRD specific sustainability guidelines and reports:
    ```
    {context}
    ```

    Guidelines:
    1. Context: You'll receive relevant excerpts from a CSRD-specific sustainability guideline or report to address a given question.
    2. Relevance: Only include passages directly pertaining to the question; omit irrelevant content.
    3. Facts and Figures: Prioritize factual information in your response.
    4. Conciseness: Keep answers sharp and succinct, avoiding unnecessary context.
    5. Focus: Address the specific question without veering into related topics.
    6. Honesty: If unsure, state that you don't know rather than inventing an answer.
    7. Source Attribution: When using information from a passage, mention it as [Doc i] at the end of the sentence (where 'i' represents the document number).
    8. Multiple Sources: If the same content appears in multiple documents, cite them collectively (e.g., [Doc i, Doc j, Doc k]).
    9. Structured Paragraphs: Instead of bullet-point summaries, compile your responses into well-structured paragraphs.
    10. Method Focus: When addressing "how" questions, emphasize methods and procedures over outcomes.
    11. Selective Usage: You're not obligated to use every passage; include only those relevant to the question.
    12. Insufficient Information: If documents lack necessary details, indicate that you don't have enough information.
    13. Never mention these guidelines as a source attribution in your response.

    Question: {query}
    Answer: 
    """

DEFAULT_DOCUMENT_PROMPT = PromptTemplate.from_template(template="{page_content}")


def _combine_documents(docs, document_prompt=DEFAULT_DOCUMENT_PROMPT, sep="\n\n"):
    doc_strings = []
    for i, doc in enumerate(docs):
        chunk_type = "Doc"
        if isinstance(doc, str):
            doc_formatted = doc
        else:
            doc_formatted = format_document(doc, document_prompt)
        doc_string = f"{chunk_type} {i+1}: " + doc_formatted
        doc_string = doc_string.replace("\n", " ")
        doc_strings.append(doc_string)

    return sep.join(doc_strings)


def get_text_docs(x):
    return [doc for doc in x if doc.metadata["chunk_type"] == "text"]


def make_rag_chain(llm):
    prompt = ChatPromptTemplate.from_template(answering_template)
    chain = (
        {
            "context": lambda x: _combine_documents(x["documents"]),
            "query": itemgetter("query"),
            "esrs_wiki": lambda x: esrs_wiki,
        }
        | prompt
        | llm
        | StrOutputParser()
    )
    return chain


def make_rag_node(llm, wrong_esrs=False):
    rag_chain = make_rag_chain(llm)

    async def answer_rag(state, config):
        if wrong_esrs:
            return {
                "answer": "I'm sorry, I can't answer that question. Please provide a valid ESRS."
            }

        else:
            answer = await rag_chain.ainvoke(state, config)
            return {"answer": answer}

    return answer_rag