import logging
import sys
import os
import re
import base64
import nest_asyncio
import pandas as pd
from pathlib import Path
from typing import Any, Dict, List, Optional
from PIL import Image
import streamlit as st
import torch
from llama_index.core import Settings, SimpleDirectoryReader, StorageContext, Document
from llama_index.core.storage.docstore import SimpleDocumentStore
# from llama_index.llms.ollama import Ollama
# from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.core.node_parser import LangchainNodeParser
from langchain.text_splitter import RecursiveCharacterTextSplitter
from llama_index.core.storage.chat_store import SimpleChatStore
from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.chat_engine import CondensePlusContextChatEngine
#from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import VectorStoreIndex
# from llama_index.llms.huggingface import HuggingFaceLLM
# from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
# from llama_index.embeddings.huggingface import HuggingFaceEmbedding
import chromadb

###############################################################################
#                           MONKEY PATCH EM bm25s                             #
###############################################################################
import bm25s

# Guardamos a referência da função original
orig_find_newline_positions = bm25s.utils.corpus.find_newline_positions

def patched_find_newline_positions(path, show_progress=True, leave_progress=True):
    """
    Versão 'gambiarra' da função original, forçando uso de encoding='utf-8'
    e ignorando erros de decodificação. Assim, evitamos UnicodeDecodeError
    mesmo que o arquivo contenha caracteres fora da faixa UTF-8.
    
    (Esta referência é real, baseada em ajustes de leitura de arquivos do Python.)
    """
    path = str(path)
    indexes = []

    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        indexes.append(f.tell())
        file_size = os.path.getsize(path)

        try:
            from tqdm.auto import tqdm
            pbar = tqdm(
                total=file_size,
                desc="Finding newlines for mmindex",
                unit="B",
                unit_scale=True,
                leave=leave_progress,
                disable=not show_progress,
            )
        except ImportError:
            pbar = None

        while True:
            line = f.readline()
            if not line:
                break
            t = f.tell()
            indexes.append(t)
            if pbar is not None:
                pbar.update(t - indexes[-2])

        if pbar is not None:
            pbar.close()

    return indexes[:-1]

# Aplicamos nosso patch
bm25s.utils.corpus.find_newline_positions = patched_find_newline_positions
###############################################################################
#                   CLASSE BM25Retriever (AJUSTADA PARA ENCODING)             #
###############################################################################
import json
import Stemmer

from llama_index.core.base.base_retriever import BaseRetriever
from llama_index.core.callbacks.base import CallbackManager
from llama_index.core.constants import DEFAULT_SIMILARITY_TOP_K
from llama_index.core.schema import (
    BaseNode,
    IndexNode,
    NodeWithScore,
    QueryBundle,
    MetadataMode,
)
from llama_index.core.vector_stores.utils import (
    node_to_metadata_dict,
    metadata_dict_to_node,
)
from typing import cast

logger = logging.getLogger(__name__)

DEFAULT_PERSIST_ARGS = {"similarity_top_k": "similarity_top_k", "_verbose": "verbose"}
DEFAULT_PERSIST_FILENAME = "retriever.json"


class BM25Retriever(BaseRetriever):
    """
    Implementação customizada do algoritmo BM25 com a lib bm25s, incluindo um
    'monkey patch' para contornar problemas de decodificação de caracteres.
    """

    def __init__(
        self,
        nodes: Optional[List[BaseNode]] = None,
        stemmer: Optional[Stemmer.Stemmer] = None,
        language: str = "en",
        existing_bm25: Optional[bm25s.BM25] = None,
        similarity_top_k: int = DEFAULT_SIMILARITY_TOP_K,
        callback_manager: Optional[CallbackManager] = None,
        objects: Optional[List[IndexNode]] = None,
        object_map: Optional[dict] = None,
        verbose: bool = False,
    ) -> None:
        self.stemmer = stemmer or Stemmer.Stemmer("english")
        self.similarity_top_k = similarity_top_k

        if existing_bm25 is not None:
            # Usa instância BM25 existente
            self.bm25 = existing_bm25
            self.corpus = existing_bm25.corpus
        else:
            # Cria uma nova instância BM25 a partir de 'nodes'
            if nodes is None:
                raise ValueError("É preciso fornecer 'nodes' ou um 'existing_bm25'.")

            self.corpus = [node_to_metadata_dict(node) for node in nodes]
            corpus_tokens = bm25s.tokenize(
                [node.get_content(metadata_mode=MetadataMode.EMBED) for node in nodes],
                stopwords=language,
                stemmer=self.stemmer,
                show_progress=verbose,
            )
            self.bm25 = bm25s.BM25()
            self.bm25.index(corpus_tokens, show_progress=verbose)

        super().__init__(
            callback_manager=callback_manager,
            object_map=object_map,
            objects=objects,
            verbose=verbose,
        )

    @classmethod
    def from_defaults(
        cls,
        index: Optional[VectorStoreIndex] = None,
        nodes: Optional[List[BaseNode]] = None,
        docstore: Optional["BaseDocumentStore"] = None,
        stemmer: Optional[Stemmer.Stemmer] = None,
        language: str = "en",
        similarity_top_k: int = DEFAULT_SIMILARITY_TOP_K,
        verbose: bool = False,
        tokenizer: Optional[Any] = None,
    ) -> "BM25Retriever":
        if tokenizer is not None:
            logger.warning(
                "O parâmetro 'tokenizer' foi descontinuado e será removido "
                "no futuro. Use um Stemmer do PyStemmer para melhor controle."
            )

        if sum(bool(val) for val in [index, nodes, docstore]) != 1:
            raise ValueError("Passe exatamente um entre 'index', 'nodes' ou 'docstore'.")

        if index is not None:
            docstore = index.docstore

        if docstore is not None:
            nodes = cast(List[BaseNode], list(docstore.docs.values()))

        assert nodes is not None, (
            "Não foi possível determinar os nodes. Verifique seus parâmetros."
        )

        return cls(
            nodes=nodes,
            stemmer=stemmer,
            language=language,
            similarity_top_k=similarity_top_k,
            verbose=verbose,
        )

    def get_persist_args(self) -> Dict[str, Any]:
        """Dicionário com os parâmetros de persistência a serem salvos."""
        return {
            DEFAULT_PERSIST_ARGS[key]: getattr(self, key)
            for key in DEFAULT_PERSIST_ARGS
            if hasattr(self, key)
        }

    def persist(self, path: str, **kwargs: Any) -> None:
        """
        Persiste o retriever em um diretório, incluindo
        a estrutura do BM25 e o corpus em JSON.
        """
        self.bm25.save(path, corpus=self.corpus, **kwargs)
        with open(
            os.path.join(path, DEFAULT_PERSIST_FILENAME),
            "wt",
            encoding="utf-8",
            errors="ignore",
        ) as f:
            json.dump(self.get_persist_args(), f, indent=2, ensure_ascii=False)

    @classmethod
    def from_persist_dir(cls, path: str, **kwargs: Any) -> "BM25Retriever":
        """
        Carrega o retriever de um diretório, incluindo o BM25 e o corpus.
        Devido ao nosso patch, ignoramos qualquer erro de decodificação
        que eventualmente apareça.
        """
        bm25_obj = bm25s.BM25.load(path, load_corpus=True, **kwargs)
        with open(
            os.path.join(path, DEFAULT_PERSIST_FILENAME),
            "rt",
            encoding="utf-8",
            errors="ignore",
        ) as f:
            retriever_data = json.load(f)

        return cls(existing_bm25=bm25_obj, **retriever_data)

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """Recupera nós relevantes a partir do BM25."""
        query = query_bundle.query_str
        tokenized_query = bm25s.tokenize(
            query, stemmer=self.stemmer, show_progress=self._verbose
        )
        indexes, scores = self.bm25.retrieve(
            tokenized_query, k=self.similarity_top_k, show_progress=self._verbose
        )

        # bm25s retorna lista de listas, pois suporta batched queries
        indexes = indexes[0]
        scores = scores[0]

        nodes: List[NodeWithScore] = []
        for idx, score in zip(indexes, scores):
            if isinstance(idx, dict):
                node = metadata_dict_to_node(idx)
            else:
                node_dict = self.corpus[int(idx)]
                node = metadata_dict_to_node(node_dict)

            nodes.append(NodeWithScore(node=node, score=float(score)))

        return nodes

#Configuração da imagem da aba
im = Image.open("pngegg.png")
st.set_page_config(page_title = "Chatbot Carômetro", page_icon=im, layout = "wide")

#Removido loop e adicionado os.makedirs
os.makedirs("bm25_retriever", exist_ok=True)
os.makedirs("chat_store", exist_ok=True)
os.makedirs("chroma_db", exist_ok=True)
os.makedirs("documentos", exist_ok=True)

# Configuração do Streamlit
st.sidebar.title("Configuração de LLM")
sidebar_option = st.sidebar.radio("Selecione o LLM", ["gpt-3.5-turbo"])

# logo_url = 'app\logos\logo-sicoob.jpg'
# st.sidebar.image(logo_url)
import base64

#Configuração da imagem da sidebar
with open("sicoob-logo.png", "rb") as f:
    data = base64.b64encode(f.read()).decode("utf-8")

    st.sidebar.markdown(
        f"""
        <div style="display:table;margin-top:-80%;margin-left:0%;">
            <img src="data:image/png;base64,{data}" width="250" height="70">
        </div>
        """,
        unsafe_allow_html=True,
    )

#if sidebar_option == "Ollama":
   # Settings.llm = Ollama(model="llama3.2:latest", request_timeout=500.0, num_gpu=1)
   # Settings.embed_model = OllamaEmbedding(model_name="nomic-embed-text:latest")
if sidebar_option == "gpt-3.5-turbo":
    from llama_index.llms.openai import OpenAI
    from llama_index.embeddings.openai import OpenAIEmbedding
    Settings.llm = OpenAI(model="gpt-3.5-turbo")
    Settings.embed_model = OpenAIEmbedding(model_name="text-embedding-ada-002")
# elif sidebar_option == 'NuExtract-1.5':
#     #Embedding do huggingface
#     Settings.embed_model = HuggingFaceEmbedding(
#         model_name="BAAI/bge-small-en-v1.5"
#     )
#     #Carregamento do modelo local, descomentar o modelo desejado

#     llm = HuggingFaceLLM(
#         context_window=2048,
#         max_new_tokens=2048,
#         generate_kwargs={"do_sample": False},
#         #query_wrapper_prompt=query_wrapper_prompt,
#         #model_name="Qwen/Qwen2.5-Coder-32B-Instruct",
#         #model_name="Qwen/Qwen2.5-14B-Instruct",
#         # model_name="meta-llama/Llama-3.2-3B",
#         #model_name="HuggingFaceH4/zephyr-7b-beta",
#         # model_name="meta-llama/Meta-Llama-3-8B",
#         model_name="numind/NuExtract-1.5",
#         #model_name="meta-llama/Llama-3.2-3B",
#         tokenizer_name="numind/NuExtract-1.5",
#         device_map="auto",
#         tokenizer_kwargs={"max_length": 512},
#         # uncomment this if using CUDA to reduce memory usage
#         model_kwargs={"torch_dtype": torch.bfloat16},
#     )
#     chat = [
#         {"role": "user", "content": "Hello, how are you?"},
#         {"role": "assistant", "content": "I'm doing great. How can I help you today?"},
#         {"role": "user", "content": "I'd like to show off how chat templating works!"},
#     ]

#     from transformers import AutoTokenizer

#     tokenizer = AutoTokenizer.from_pretrained("numind/NuExtract-1.5")
#     tokenizer.apply_chat_template(chat, tokenize=False)

#     Settings.chunk_size = 512
#     Settings.llm = llm

else:
    raise Exception("Opção de LLM inválida!")

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))


# Diretórios configurados pelo usuário
chat_store_path = os.path.join("chat_store", "chat_store.json")
documents_path = os.path.join("documentos")
chroma_storage_path = os.path.join("chroma_db")  # Diretório para persistência do Chroma
bm25_persist_path = os.path.join("bm25_retriever")

# Classe CSV Customizada (novo código)
class CustomPandasCSVReader:
    """PandasCSVReader modificado para incluir cabeçalhos nos documentos."""
    def __init__(
        self,
        *args: Any,
        concat_rows: bool = True,
        col_joiner: str = ", ",
        row_joiner: str = "\n",
        pandas_config: dict = {},
        **kwargs: Any
    ) -> None:
        self._concat_rows = concat_rows
        self._col_joiner = col_joiner
        self._row_joiner = row_joiner
        self._pandas_config = pandas_config

    def load_data(
        self,
        file: Path,
        extra_info: Optional[Dict] = None,
    ) -> List[Document]:
        df = pd.read_csv(file, **self._pandas_config)
        text_list = [" ".join(df.columns.astype(str))]
        text_list += (
            df.astype(str)
            .apply(lambda row: self._col_joiner.join(row.values), axis=1)
            .tolist()
        )

        metadata = {"filename": file.name, "extension": file.suffix}
        if extra_info:
            metadata.update(extra_info)

        if self._concat_rows:
            return [Document(text=self._row_joiner.join(text_list), metadata=metadata)]
        else:
            return [
                Document(text=text, metadata=metadata)
                for text in text_list
            ]

def clean_documents(documents):
    """Remove caracteres não desejados diretamente nos textos dos documentos."""
    cleaned_documents = []
    for doc in documents:
        cleaned_text = re.sub(r"[^0-9A-Za-zÀ-ÿ ]", "", doc.get_content())
        doc.text = cleaned_text
        cleaned_documents.append(doc)
    return cleaned_documents

from llama_index.readers.google import GoogleDriveReader
import json

credentials_json = os.getenv('GOOGLE_CREDENTIALS')
token_json = os.getenv('GOOGLE_TOKEN')

if credentials_json is None:
    raise ValueError("The GOOGLE_CREDENTIALS environment variable is not set.")

# Write the credentials to a file
credentials_path = "credentials.json"
token_path = "token.json"
with open(credentials_path, 'w') as credentials_file:
    credentials_file.write(credentials_json)

with open(token_path, 'w') as credentials_file:
    credentials_file.write(token_json)
    
google_drive_reader = GoogleDriveReader(credentials_path=credentials_path)
google_drive_reader._creds = google_drive_reader._get_credentials()

def are_docs_downloaded(directory_path: str) -> bool:
    return os.path.isdir(directory_path) and any(os.scandir(directory_path))

def download_original_files_from_folder(greader: GoogleDriveReader, pasta_documentos_drive: str, local_path: str):
    os.makedirs(local_path, exist_ok=True)
    files_meta = greader._get_fileids_meta(folder_id=pasta_documentos_drive)
    if not files_meta:
        logging.info("Nenhum arquivo encontrado na pasta especificada.")
        return
    for fmeta in files_meta:
        file_id = fmeta[0]
        file_name = os.path.basename(fmeta[2])
        local_file_path = os.path.join(local_path, file_name)

        if os.path.exists(local_file_path):
            logging.info(f"Arquivo '{file_name}' já existe localmente, ignorando download.")
            continue

        downloaded_file_path = greader._download_file(file_id, local_file_path)
        if downloaded_file_path:
            logging.info(f"Arquivo '{file_name}' baixado com sucesso em: {downloaded_file_path}")
        else:
            logging.warning(f"Não foi possível baixar '{file_name}'")

#DADOS/QA_database/Documentos CSV/documentos
pasta_documentos_drive = "1xVzo8s1D0blzR5ZB3m5k4dVWHuRmKUu-"

# Verifica e baixa arquivos se necessário (novo código)
if not are_docs_downloaded(documents_path):
    logging.info("Baixando arquivos originais do Drive para 'documentos'...")
    download_original_files_from_folder(google_drive_reader, pasta_documentos_drive, documents_path)
else:
    logging.info("'documentos' já contém arquivos, ignorando download.")

# Configuração de leitura de documentos
file_extractor = {".csv": CustomPandasCSVReader()}
documents = SimpleDirectoryReader(
    input_dir=documents_path,
    file_extractor=file_extractor,
    filename_as_id=True,
    recursive=True
    #Recursive caso tenha varias pastas no drive
).load_data()

documents = clean_documents(documents)

# Configuração do Chroma e BM25 com persistência
docstore = SimpleDocumentStore()
docstore.add_documents(documents)

db = chromadb.PersistentClient(path=chroma_storage_path)
chroma_collection = db.get_or_create_collection("dense_vectors")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

# Configuração do StorageContext
storage_context = StorageContext.from_defaults(
    docstore=docstore, vector_store=vector_store
)

# Criação/Recarregamento do índice com embeddings
if os.path.exists(chroma_storage_path):
    index = VectorStoreIndex.from_vector_store(vector_store)
else:
    splitter = LangchainNodeParser(
        RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=128)
    )
    index = VectorStoreIndex.from_documents(
        documents,
        storage_context=storage_context,
        transformations=[splitter]
    )
    vector_store.persist()

# Criação/Recarregamento do BM25 Retriever
if os.path.exists(os.path.join(bm25_persist_path, "params.index.json")):
    bm25_retriever = BM25Retriever.from_persist_dir(bm25_persist_path)
else:
    bm25_retriever = BM25Retriever.from_defaults(
        docstore=docstore,
        similarity_top_k=2,
        language="portuguese",  # Idioma ajustado para seu caso
    )
    os.makedirs(bm25_persist_path, exist_ok=True)
    bm25_retriever.persist(bm25_persist_path)

# Combinação de Retrievers (Embeddings + BM25)
vector_retriever = index.as_retriever(similarity_top_k=2)
retriever = QueryFusionRetriever(
    [vector_retriever, bm25_retriever],
    similarity_top_k=3,
    num_queries=0,
    mode="reciprocal_rerank",
    use_async=True,
    verbose=True,
    query_gen_prompt=(
        "Gere {num_queries} perguntas de busca relacionadas à seguinte pergunta. "
        "Priorize o significado da pergunta sobre qualquer histórico de conversa. "
        "Se o histórico não for relevante para a pergunta, ignore-o. "
        "Não adicione explicações, notas ou introduções. Apenas escreva as perguntas. "
        "Pergunta: {query}\n\n"
        "Perguntas:\n"
    ),
)

# Configuração do chat engine
nest_asyncio.apply()
memory = ChatMemoryBuffer.from_defaults(token_limit=3900)
query_engine = RetrieverQueryEngine.from_args(retriever)
chat_engine = CondensePlusContextChatEngine.from_defaults(
    query_engine,
    memory=memory,
    context_prompt=(
        "Você é um assistente virtual capaz de interagir normalmente, além de"
        " fornecer informações sobre organogramas e listar funcionários."
        " Aqui estão os documentos relevantes para o contexto:\n"
        "{context_str}"
        "\nInstrução: Use o histórico da conversa anterior, ou o contexto acima, para responder."
    ),
    verbose=True,
)

# Armazenamento do chat
chat_store = SimpleChatStore()
if os.path.exists(chat_store_path):
    chat_store = SimpleChatStore.from_persist_path(persist_path=chat_store_path)
else:
    chat_store.persist(persist_path=chat_store_path)

# Interface do Chatbot
st.title("Chatbot Carômetro")
st.write("Este chatbot pode te ajudar a conseguir informações relevantes sobre os carômetros da Sicoob.")

if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []

for message in st.session_state.chat_history:
    role, text = message.split(":", 1)
    with st.chat_message(role.strip().lower()):
        st.write(text.strip())

user_input = st.chat_input("Digite sua pergunta")
if user_input:
    # Exibir a mensagem do usuário e adicionar ao histórico
    with st.chat_message('user'):
        st.write(user_input)
    st.session_state.chat_history.append(f"user: {user_input}")

    # Placeholder para a mensagem do assistente
    with st.chat_message('assistant'):
        message_placeholder = st.empty()
        assistant_message = ''

    # Obter a resposta em streaming do chat_engine
    response = chat_engine.stream_chat(user_input)
    for token in response.response_gen:
        assistant_message += token
        # Atualizar o placeholder da mensagem
        message_placeholder.markdown(assistant_message + "▌")

    # Remover o cursor após a conclusão
    message_placeholder.markdown(assistant_message)
    st.session_state.chat_history.append(f"assistant: {assistant_message}")