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

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

from _utils.models.gerar_documento import DocumentChunk
from gerar_documento.serializer import GerarDocumentoSerializerData
import tiktoken

encoding = tiktoken.get_encoding("cl100k_base")


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],
    }


# 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]):
    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)


def split_text_by_tokens(full_text: str):
    tokens = encoding.encode(full_text)
    max_tokens = 600000

    # Divide os tokens em partes de no máximo max_tokens
    token_chunks = [
        tokens[i : i + max_tokens] for i in range(0, len(tokens), max_tokens)
    ]

    # Decodifica cada pedaço de tokens de volta para texto
    text_chunks = [encoding.decode(chunk) for chunk in token_chunks]

    return text_chunks


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

    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