import os
from typing import Any, Union
from _utils.custom_exception_handler import custom_exception_handler_without_api_handler
from rest_framework.response import Response
from _utils.gerar_documento_utils.GerarDocumento import (
    GerarDocumento,
)

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,
) -> Response | dict[str, Any]:
    try:
        summarizer = GerarDocumento(serializer, isBubble, axiom_instance)
        summarizer.lista_pdfs = listaPDFs

        await summarizer.get_text_and_pdf_chunks()

        await summarizer.get_response_from_auxiliar_contextual_prompt()

        await summarizer.generate_chunks_processados()

        await summarizer.generate_query_for_vector_store()

        await summarizer.create_enhanced_vector_store()

        structured_summaries = await summarizer.do_last_requests()

        if not isinstance(structured_summaries, list):
            return Response({"erro": structured_summaries})

        await summarizer.generate_complete_text()

        await summarizer.get_document_title()

        if isBubble:
            await summarizer.send_to_bubble()

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