import os
from langchain_core.messages import HumanMessage
from typing import Any, Union, cast
from _utils.langchain_utils.LLM_class import LLM
from _utils.bubble_integrations.enviar_resposta_final import enviar_resposta_final
from _utils.custom_exception_handler import custom_exception_handler_wihout_api_handler
from _utils.gerar_relatorio_modelo_usuario.prompts import (
    prompt_gerar_query_dinamicamente,
)
from _utils.gerar_relatorio_modelo_usuario.GerarDocumento import (
    GerarDocumento,
)
from _utils.gerar_relatorio_modelo_usuario.contextual_retriever import (
    ContextualRetriever,
)
from _utils.gerar_relatorio_modelo_usuario.utils import (
    generate_document_title,
    gerar_resposta_compilada,
    get_full_text_and_all_PDFs_chunks,
    get_response_from_auxiliar_contextual_prompt,
)
from _utils.models.gerar_relatorio import (
    RetrievalConfig,
)
import markdown

from _utils.langchain_utils.Prompt_class import Prompt
from _utils.utils import convert_markdown_to_HTML
from gerar_documento.serializer import (
    GerarDocumentoComPDFProprioSerializer,
    GerarDocumentoComPDFProprioSerializerData,
    GerarDocumentoSerializerData,
)
from setup.logging import Axiom


os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_ENDPOINT"] = "https://api.smith.langchain.com"
os.environ.get("LANGCHAIN_API_KEY")
os.environ["LANGCHAIN_PROJECT"] = "VELLA"


async def gerar_documento(
    serializer: Union[
        GerarDocumentoSerializerData, GerarDocumentoComPDFProprioSerializerData, Any
    ],
    listaPDFs,
    axiom_instance: Axiom,
    isBubble=False,
):
    """Parâmetro "contexto" só deve ser passado quando quiser utilizar o teste com ragas, e assim, não quiser passar PDFs"""
    try:
        contextual_retriever = ContextualRetriever(serializer)

        # Initialize enhanced summarizer
        summarizer = GerarDocumento(serializer)

        all_PDFs_chunks, full_text_as_array = await get_full_text_and_all_PDFs_chunks(
            listaPDFs,
            summarizer.splitter,
            serializer.should_use_llama_parse,
            isBubble,
        )

        is_contextualized_chunk = serializer.should_have_contextual_chunks

        if is_contextualized_chunk:
            response_auxiliar_summary = (
                await get_response_from_auxiliar_contextual_prompt(full_text_as_array)
            )
            axiom_instance.send_axiom(
                f"RESUMO INICIAL DO PROCESSO: {response_auxiliar_summary}"
            )

            axiom_instance.send_axiom("COMEÇANDO A FAZER AS REQUISIÇÕES DO CONTEXTUAL")
            contextualized_chunks = await contextual_retriever.contextualize_all_chunks(
                all_PDFs_chunks, response_auxiliar_summary
            )
            axiom_instance.send_axiom(
                "TERMINOU DE FAZER TODAS AS REQUISIÇÕES DO CONTEXTUAL"
            )
            chunks_processados = contextualized_chunks
            axiom_instance.send_axiom(
                f"CHUNKS PROCESSADOS INICIALMENTE: {chunks_processados}"
            )
        else:
            chunks_processados = all_PDFs_chunks

        llm = LLM()
        prompt_para_gerar_query_dinamico = prompt_gerar_query_dinamicamente(
            cast(str, response_auxiliar_summary)
        )

        axiom_instance.send_axiom(
            "COMEÇANDO REQUISIÇÃO PARA GERAR O QUERY DINAMICAMENTE DO VECTOR STORE"
        )
        query_gerado_dinamicamente_para_o_vector_store = await llm.google_gemini(
            "gemini-2.5-pro-exp-03-25"
        ).ainvoke([HumanMessage(content=prompt_para_gerar_query_dinamico)])

        axiom_instance.send_axiom(
            f"query_gerado_dinamicamente_para_o_vector_store: {query_gerado_dinamicamente_para_o_vector_store.content}",
        )

        # Create enhanced vector store and BM25 index
        vector_store, bm25, chunk_ids = (
            summarizer.vector_store.create_enhanced_vector_store(
                chunks_processados, is_contextualized_chunk, axiom_instance
            )
        )

        llm_ultimas_requests = serializer.llm_ultimas_requests
        axiom_instance.send_axiom("COMEÇANDO A FAZER ÚLTIMA REQUISIÇÃO")
        structured_summaries = await summarizer.gerar_documento_final(
            vector_store,
            bm25,
            chunk_ids,
            llm_ultimas_requests,
            cast(
                str, query_gerado_dinamicamente_para_o_vector_store.content
            ),  # prompt_auxiliar_SEM_CONTEXT,
        )
        axiom_instance.send_axiom("TERMINOU DE FAZER A ÚLTIMA REQUISIÇÃO")

        if not isinstance(structured_summaries, list):
            from rest_framework.response import Response

            return Response({"erro": structured_summaries})

        texto_completo = summarizer.resumo_gerado + "\n\n"

        for x in structured_summaries:
            texto_completo = texto_completo + x["content"] + "\n"
            x["source"]["text"] = x["source"]["text"][0:200]
            x["source"]["context"] = x["source"]["context"][0:200]

        texto_completo_como_html = convert_markdown_to_HTML(texto_completo)
        axiom_instance.send_axiom(
            f"texto_completo_como_html: {texto_completo_como_html}"
        )

        if is_contextualized_chunk:
            prompt_titulo_do_documento = response_auxiliar_summary
        else:
            prompt_titulo_do_documento = texto_completo_como_html
        titulo_do_documento = await generate_document_title(
            cast(str, prompt_titulo_do_documento)
        )

        if isBubble:
            axiom_instance.send_axiom("COMEÇANDO A REQUISIÇÃO FINAL PARA O BUBBLE")
            enviar_resposta_final(
                serializer.doc_id,  # type: ignore
                serializer.form_response_id,  # type: ignore
                serializer.version,  # type: ignore
                texto_completo_como_html,
                False,
                cast(str, titulo_do_documento),
            )
            axiom_instance.send_axiom("TERMINOU A REQUISIÇÃO FINAL PARA O BUBBLE")

        return {
            "texto_completo": texto_completo_como_html,
            "titulo_do_documento": titulo_do_documento,
            "resultado": structured_summaries,
            "parametros-utilizados": gerar_resposta_compilada(serializer),
        }
    except Exception as e:
        custom_exception_handler_wihout_api_handler(e, serializer, axiom_instance)
        raise