import os
import argparse
import logging
import gradio as gr
from huggingface_hub import login
import spaces


parser = argparse.ArgumentParser("cvchatbot")
parser.add_argument("--debug", help="Enable debug logging", action='store_true')
parser.add_argument("--prod", help="Use production environment", action='store_true')
parser.add_argument("--generate", dest="generate_cache", help="Generate new cache", action='store_true')
parser.add_argument("--download-cv-data", dest="download_cv_data", help="Download CV data from Cloud Storage", action='store_true')
args = parser.parse_args()
if args.prod:
    os.environ["ENV"] = "PROD"
    
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.DEBUG if args.debug else logging.INFO)


from config import CV_DATA_DIRNAME, EMBEDDING_MODEL_ID, HF_TOKEN, LLM_ID, \
    WELCOME_MESSAGE_TEMPLATE, NO_MATCH_RESPONSE, RESET_QUERY, \
    CV_DATA_FILE_EDUCATION, CV_DATA_FILE_JOBS, CV_DATA_FILE_PAPERS, CV_DATA_FILE_SKILLS, \
    EDUCATION_DESCRIPTION, JOBS_DESCRIPTION, PAPERS_DESCRIPTION, SKILLS_DESCRIPTION
from helpers import get_education_documents, get_job_documents, get_paper_documents, get_skill_documents
from knowledge import ChatContext, DomainKnowledgeConfig, KnowledgeBase, KnowledgeConfig


logger.info("Logging in to Hugging Face Hub...")
login(token=HF_TOKEN)

logger.info("Initializing CV Knowledge Base...")
cv_knowledge_base = KnowledgeBase(
    KnowledgeConfig(
        WELCOME_MESSAGE_TEMPLATE,
        NO_MATCH_RESPONSE,
        RESET_QUERY,
        [
            DomainKnowledgeConfig(JOBS_DESCRIPTION, CV_DATA_DIRNAME + CV_DATA_FILE_JOBS, get_job_documents),
            DomainKnowledgeConfig(SKILLS_DESCRIPTION, CV_DATA_DIRNAME + CV_DATA_FILE_SKILLS, get_skill_documents),
            DomainKnowledgeConfig(EDUCATION_DESCRIPTION, CV_DATA_DIRNAME + CV_DATA_FILE_EDUCATION, get_education_documents),
            DomainKnowledgeConfig(PAPERS_DESCRIPTION, CV_DATA_DIRNAME + CV_DATA_FILE_PAPERS, get_paper_documents),
        ],
    ),
    args.generate_cache,
    args.download_cv_data,
)


@spaces.GPU
def respond(
    message: str,
    history: list[gr.MessageDict],
    chat_context: ChatContext,
):
    yield from cv_knowledge_base.answer_query(message, history, chat_context)


with gr.Blocks() as demo:
    # login button
    with gr.Sidebar():
        gr.LoginButton()

    # chat context
    context_state = gr.State(ChatContext())
    
    # chat interface
    gr.ChatInterface(
        respond,
        chatbot=gr.Chatbot(
            label=f"{LLM_ID} + {EMBEDDING_MODEL_ID}",
            value=[gr.MessageDict(role="assistant", content=cv_knowledge_base.knowledge.welcome_message)],
            avatar_images=(None, "./assets/chatbot_icon.png"),
        ),
        textbox=gr.Textbox(
            submit_btn=True,
        ),
        additional_inputs=[context_state], # Pass the updated state as an input
        additional_outputs=[context_state], # Receive the updated state as an output
    )


if __name__ == "__main__":
    demo.launch()