from typing import Any, List, Tuple, Union
from langchain_core.documents import Document
from langchain_core.messages import HumanMessage

from _utils.gerar_relatorio_modelo_usuario.llm_calls import agemini_answer
from _utils.langchain_utils.Splitter_class import Splitter
from _utils.langchain_utils.LLM_class import LLM
from _utils.gerar_relatorio_modelo_usuario.prompts import (
    create_prompt_auxiliar_do_contextual_prompt,
)
import re

from gerar_documento.serializer import GerarDocumentoSerializerData


def gerar_resposta_compilada(serializer: Union[GerarDocumentoSerializerData, Any]):
    return {
        "num_chunks_retrieval": serializer.num_chunks_retrieval,
        "embedding_weight": serializer.embedding_weight,
        "bm25_weight": serializer.bm25_weight,
        "context_window": serializer.context_window,
        "chunk_overlap": serializer.chunk_overlap,
        "num_k_rerank": serializer.num_k_rerank,
        "model_cohere_rerank": serializer.model_cohere_rerank,
        "more_initial_chunks_for_reranking": serializer.more_initial_chunks_for_reranking,
        "claude_context_model": serializer.claude_context_model,
        "gpt_temperature": serializer.gpt_temperature,
        "user_message": serializer.user_message,
        "model": serializer.model,
        "hf_embedding": serializer.hf_embedding,
        "chunk_size": serializer.chunk_size,
        "chunk_overlap": serializer.chunk_overlap,
        # "prompt_auxiliar": serializer.prompt_auxiliar,
        "prompt_gerar_documento": serializer.prompt_gerar_documento[0:200],
    }


def check_regex_patterns(context: str, lista_de_document_ids: List[int]):
    patterns = [
        # r"\[*([\d.\-]+)\]*\s*---\s*\[*([^]]+)\]*\s*---\s*\[*([^]]+)\]*\s*</chunk_context>", # PRIMEIRO DE TODOS
        # r"<chunk_context>\s*([\d.\-]+)\s*---\s*([^<]+)\s*---\s*([^<]+)\s*</chunk_context>",
        r"<chunk_context>\s*(\d+)(?:\s*-\s*Pág\.\s*\d+)?\s*---\s*([^-\n]+)\s*---\s*([^<]+)</chunk_context>",
        r"<chunk_context>\s*(?:\[*([\d]+)\]*\s*[-–]*\s*(?:Pág\.\s*\d+\s*[-–]*)?)?\s*\[*([^\]]+)\]*\s*[-–]*\s*\[*([^\]]+)\]*\s*[-–]*\s*\[*([^\]]+)\]*\s*</chunk_context>",
        # r"\[([\d.\-]+)\]\s*---\s*\[([^]]+)\]\s*---\s*\[([^]]+)\]\s*</chunk_context>",
        # r"<chunk_context>\s*\[?([\d.\-]+)\]?\s*---\s*\[?([^\]\[]+?)\]?\s*---\s*\[?([^<]+?)\]?\s*</chunk_context>",
        # r"<chunk_context>\s*\[([\d.\-]+)\]\s*---\s*\[([^\]]+)\]\s*---\s*\[([^\]]+)\]\s*</chunk_context>"
        # r"<chunk_context>\s*\[?([\d.\-\s]+)\]?\s*---\s*\[?([^\]\[]+?)\]?\s*---\s*\[?([\s\S]+?)\]?\s*</chunk_context>",
    ]

    for pattern in patterns:
        matches = re.findall(pattern, context, re.DOTALL)
        if len(matches) == len(lista_de_document_ids):
            print("\n--------------- REGEX DO CONTEXTUAL FUNCIONOU")
            break

    return matches


def validate_many_chunks_in_one_request(
    response: str, lista_de_document_ids: List[int]
):
    context = (
        response.replace("document_id: ", "")
        .replace("document_id:", "")
        .replace("DOCUMENT_ID: ", "")
        .replace("DOCUMENT_ID: ", "")
    )

    # pattern = r"\[(\d+|[-.]+)\] --- (.+?) --- (.+?)</chunk_context>"  # Funciona para quando a resposta do LLM não vem com "document_id" escrito
    matches = check_regex_patterns(context, lista_de_document_ids)

    matches_as_list = []

    for index, match in enumerate(list(matches)):
        if index >= 20:
            break
        resultado = match[0].replace(".", "").replace("-", "")

        resultado = lista_de_document_ids[index]

        matches_as_list.append((resultado, match[1], match[2]))

    if len(matches) == 0:
        print(
            "----------- ERROU NA TENTATIVA ATUAL DE FORMATAR O CONTEXTUAL -----------"
        )
        return False
    return matches_as_list


# Esta função gera a resposta que será usada em cada um das requisições de cada chunk
async def get_response_from_auxiliar_contextual_prompt(full_text_as_array: List[str]):
    import tiktoken

    encoding = tiktoken.get_encoding("cl100k_base")
    llms = LLM()
    responses = []

    current_chunk = []
    current_token_count = 0
    chunk_counter = 1

    for part in full_text_as_array:
        part_tokens = len(encoding.encode(part))

        # Check if adding this part would EXCEED the limit
        if current_token_count + part_tokens > 600000:
            # Process the accumulated chunk before it exceeds the limit
            chunk_text = "".join(current_chunk)
            print(
                f"\nProcessing chunk {chunk_counter} with {current_token_count} tokens"
            )

            prompt = create_prompt_auxiliar_do_contextual_prompt(chunk_text)
            response = await llms.google_gemini().ainvoke(
                [HumanMessage(content=prompt)]
            )
            responses.append(response.content)

            # Start new chunk with current part
            current_chunk = [part]
            current_token_count = part_tokens
            chunk_counter += 1
        else:
            # Safe to add to current chunk
            current_chunk.append(part)
            current_token_count += part_tokens

    # Process the final remaining chunk
    if current_chunk:
        chunk_text = "".join(current_chunk)
        print(
            f"\nProcessing final chunk {chunk_counter} with {current_token_count} tokens"
        )
        prompt = create_prompt_auxiliar_do_contextual_prompt(chunk_text)
        response = await llms.google_gemini().ainvoke([HumanMessage(content=prompt)])
        responses.append(response.content)

    return "".join(responses)


async def get_full_text_and_all_PDFs_chunks(
    listaPDFs: List[str],
    splitterObject: Splitter,
    should_use_llama_parse: bool,
    isBubble: bool,
):
    all_PDFs_chunks = []

    pages: List[str] = []

    # Load and process document
    for pdf_path in listaPDFs:
        chunks, pages = await splitterObject.load_and_split_document(
            pdf_path, should_use_llama_parse, isBubble
        )
        all_PDFs_chunks = all_PDFs_chunks + chunks

    return all_PDFs_chunks, pages


async def generate_document_title(resumo_para_gerar_titulo: str):
    prompt = f"Você é um assistente jurídico e irá receber abaixo o resumo de um documento jurídico. Quero que você gere um título para este documento. Mande como resposta apenas o título gerado, nada mais. Aqui está um título de exemplo pra você se basear ao criar um novo: <titulo_de_exemplo>Ação Penal por Furto Qualificado nº 0002269-86.2009.805.0032<titulo_de_exemplo>\n\nSegue abaixo o resumo do documento jurídico:\n{resumo_para_gerar_titulo}"
    response = await agemini_answer(prompt, "gemini-2.0-flash-lite")
    return response