import os
import json
import gradio as gr
import streamlit as st
from huggingface_hub import HfApi, login
from dotenv import load_dotenv

from llm import get_groq_llm
from vectorstore import get_chroma_vectorstore
from embeddings import get_SFR_Code_embedding_model
from kadi_apy_bot import KadiAPYBot

# Load environment variables from .env file
load_dotenv()

vectorstore_path = "data/vectorstore"

GROQ_API_KEY = os.environ["GROQ_API_KEY"]
HF_TOKEN = os.environ["HF_Token"]


with open('config.json', 'r') as file:
    config = json.load(file)


login(HF_TOKEN)
hf_api = HfApi()

# Access the values
LLM_MODEL_NAME = config['llm_model_name']
LLM_MODEL_TEMPERATURE = float(config['llm_model_temperature'])


def initialize():
    global kadiAPY_bot

    vectorstore = get_chroma_vectorstore(get_SFR_Code_embedding_model(), vectorstore_path)
    llm = get_groq_llm(LLM_MODEL_NAME, LLM_MODEL_TEMPERATURE, GROQ_API_KEY)

    kadiAPY_bot = KadiAPYBot(llm, vectorstore)

initialize()




def bot_kadi(history, session_state):
    user_query = history[-1][0]   

    response = kadiAPY_bot.process_query({
        "query": user_query,
        "history": session_state["conversation"]  # Pass full conversation history
    })

    # Update the session history
    history[-1] = (user_query, response)
    session_state["conversation"].append({"query": user_query, "response": response})

    yield history

# Gradio utils with session state
def main():
    with gr.Blocks() as demo:
        gr.Markdown("## KadiAPY - AI Coding-Assistant")
        gr.Markdown("AI assistant for KadiAPY based on RAG architecture powered by LLM")
        
        session_state = gr.State({"conversation": []}) 

        with gr.Tab("KadiAPY - AI Assistant"):
            with gr.Row():
                with gr.Column(scale=10):
                    chatbot = gr.Chatbot([], elem_id="chatbot", label="Kadi Bot", bubble_full_width=False, show_copy_button=True, height=600)
                    user_txt = gr.Textbox(label="Question", placeholder="Type in your question and press Enter or click Submit")

                    with gr.Row():
                        with gr.Column(scale=1):
                            submit_btn = gr.Button("Submit", variant="primary")
                        with gr.Column(scale=1):
                            clear_btn = gr.Button("Clear", variant="stop")

                    gr.Examples(
                        examples=[
                            "Write me a python script which can convert plain JSON to a Kadi4Mat-compatible extra metadata structure",
                            "I need a method to upload a file to a record. The id of the record is 3",
                        ],
                        inputs=user_txt,
                        outputs=chatbot,
                        fn=add_text,
                        label="Try asking...",
                        cache_examples=False,
                        examples_per_page=3,
                    )

            # Bind input and button to modified bot_kadi
            user_txt.submit(check_input_text, user_txt, None).success(add_text, [chatbot, user_txt], [chatbot, user_txt]).then(
                bot_kadi, [chatbot, session_state], [chatbot]
            )
            submit_btn.click(check_input_text, user_txt, None).success(add_text, [chatbot, user_txt], [chatbot, user_txt]).then(
                bot_kadi, [chatbot, session_state], [chatbot]
            )
            clear_btn.click(lambda: None, None, chatbot, queue=False).then(
                lambda: {"conversation": []}, None, session_state, queue=False  # Clear session state
            )

    demo.launch()

if __name__ == "__main__":
    main()