from src.config import setup_logging, Config
from src.embeddings_model import GEmbeddings
from src.text_generation_model import GLLM
from src.pinecone_index import PineconeIndex
from src.llamaindex_backend import GLlamaIndex

import gradio as gr
import google.generativeai as genai
from llama_index.core import Settings

from typing import List
import time

# import dotenv
# dotenv.load_dotenv(".env")

logger = setup_logging()

# Google Generative AI
genai.configure(api_key=Config.GAI_API_KEY)

# Llama-Index LLM
embed_model = GEmbeddings(model_name=Config.EMB_MODEL_NAME)
llm = GLLM(model_name=Config.TEXT_MODEL_NAME, system_instruction=None)

Settings.embed_model = embed_model
Settings.llm = llm

index = PineconeIndex(api_key=Config.PINECONE_API_KEY, index_name=Config.PC_INDEX_NAME, index_namespace=Config.PC_INDEX_NAMESPACE)
backend = GLlamaIndex(logger, embed_model, llm, index, Config.SIMILARITY_THRESHOLD)

# Gradio
chat_history = []

def clear_chat() -> None:
    global chat_history 
    chat_history = []
    return None

def get_chat_history(chat_history: List[str]) -> str:
    ind = 0
    formatted_chat_history = ""
    for message in chat_history:
        formatted_chat_history += f"User: \n{message}\n" if ind % 2 == 0 else f"Bot: \n{message}\n"
        ind += 1
    return formatted_chat_history

def generate_text(prompt: str, backend: GLlamaIndex):
    global chat_history

    logger.info("Generating Message...")
    logger.info(f"User Message:\n{prompt}\n")

    result = backend.generate_text(prompt, chat_history)
    chat_history.append(prompt)
    chat_history.append(result)
    
    logger.info(f"Replied Message:\n{result}\n")
    return result

if __name__ == "__main__":
    try:
        with gr.Blocks(css=".input textarea {font-size: 16px !important}") as app:
            chatbot = gr.Chatbot(
                bubble_full_width=False, 
                container=True, 
                show_share_button=False, 
                avatar_images=[None, './asset/akag-g-only.png']
            )
            msg = gr.Textbox(
                show_label=False,
                label="Type your message...", 
                placeholder="Hi Gerard, can you introduce yourself?",
                container=False, 
                elem_classes="input"
            )
            with gr.Row():
                clear = gr.Button("Clear", scale=1)
                send = gr.Button(
                    value="", 
                    variant="primary",
                    icon="./asset/send-message.png", 
                    scale=1
                )

            def user(user_message, history):
                return "", history + [[user_message, None]]

            def bot(history):
                bot_message = generate_text(history[-1][0], backend)
                history[-1][1] = ""
                for character in bot_message:
                    history[-1][1] += character
                    time.sleep(0.01)
                    yield history

            msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot, chatbot, chatbot
            )
            send.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot, chatbot, chatbot
            )
            clear.click(clear_chat, None, chatbot, queue=False)

            gr.HTML("""
                <p><center><i>Disclaimer: This RAG app is for demostration only. Hallucination might occur.</i></center></p>
                <p><center>Hosted on 🤗 Spaces | Built with Google Gemini & 🦙 LlamaIndex | Last updated 2025</center></p>
            """)
            
        app.queue()
        app.launch()
        
    except Exception as e:
        logger.exception(e)