from typing import List, Optional

import gradio as gr
from langchain_core.vectorstores import VectorStore

from config import (
    LLM_MODEL_REPOS,
    EMBED_MODEL_REPOS,
    SUBTITLES_LANGUAGES,
    GENERATE_KWARGS,
)

from utils import (
    load_llm_model,
    load_embed_model,
    load_documents_and_create_db,
    user_message_to_chatbot,
    update_user_message_with_context,
    get_llm_response,
    get_gguf_model_names,
    add_new_model_repo,
    clear_llm_folder,
    clear_embed_folder,
    get_memory_usage,
)


# ============ INTERFACE COMPONENT INITIALIZATION FUNCS ============

def get_rag_settings(rag_mode: bool, render: bool = True):
    k = gr.Radio(
        choices=[1, 2, 3, 4, 5, 'all'],
        value=2,
        label='Number of relevant documents for search',
        visible=rag_mode,
        render=render,
        )
    score_threshold = gr.Slider(
        minimum=0,
        maximum=1,
        value=0.5,
        step=0.05,
        label='relevance_scores_threshold',
        visible=rag_mode,
        render=render,
        )
    return k, score_threshold


def get_user_message_with_context(text: str, rag_mode: bool) -> gr.component:
    num_lines = len(text.split('\n'))
    max_lines = 10
    num_lines = max_lines if num_lines > max_lines else num_lines
    return gr.Textbox(
        text,
        visible=rag_mode,
        interactive=False,
        label='User Message With Context',
        lines=num_lines,
        )


def get_system_prompt_component(interactive: bool) -> gr.Textbox:
    value = '' if interactive else 'System prompt is not supported by this model'
    return gr.Textbox(value=value, label='System prompt', interactive=interactive)


def get_generate_args(do_sample: bool) -> List[gr.component]:
    generate_args = [
        gr.Slider(minimum=0.1, maximum=3, value=GENERATE_KWARGS['temperature'], step=0.1, label='temperature', visible=do_sample),
        gr.Slider(minimum=0.1, maximum=1, value=GENERATE_KWARGS['top_p'], step=0.01, label='top_p', visible=do_sample),
        gr.Slider(minimum=1, maximum=50, value=GENERATE_KWARGS['top_k'], step=1, label='top_k', visible=do_sample),
        gr.Slider(minimum=1, maximum=5, value=GENERATE_KWARGS['repeat_penalty'], step=0.1, label='repeat_penalty', visible=do_sample),
    ]
    return generate_args


def get_rag_mode_component(db: Optional[VectorStore]) -> gr.Checkbox:
    value = visible = db is not None
    return gr.Checkbox(value=value, label='RAG Mode', scale=1, visible=visible)


# ================ LOADING AND INITIALIZING MODELS ========================

start_llm_model, start_support_system_role, load_log = load_llm_model(LLM_MODEL_REPOS[0], 'gemma-2-2b-it-Q8_0.gguf')
start_embed_model, load_log = load_embed_model(EMBED_MODEL_REPOS[0])



# ================== APPLICATION WEB INTERFACE ============================

theme = gr.themes.Base(primary_hue='green', secondary_hue='yellow', neutral_hue='zinc').set(
    loader_color='rgb(0, 255, 0)',
    slider_color='rgb(0, 200, 0)',
    body_text_color_dark='rgb(0, 200, 0)',
    button_secondary_background_fill_dark='green',
)
css = '''.gradio-container {width: 60% !important}'''

with gr.Blocks(theme=theme, css=css) as interface:

    # ==================== GRADIO STATES ===============================

    documents = gr.State([])
    db = gr.State(None)
    user_message_with_context = gr.State('')
    support_system_role = gr.State(start_support_system_role)
    llm_model_repos = gr.State(LLM_MODEL_REPOS)
    embed_model_repos = gr.State(EMBED_MODEL_REPOS)
    llm_model = gr.State(start_llm_model)
    embed_model = gr.State(start_embed_model)



    # ==================== BOT PAGE =================================

    with gr.Tab(label='Chatbot'):
        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    show_copy_button=True,
                    bubble_full_width=False,
                    height=480,
                )
                user_message = gr.Textbox(label='User')

                with gr.Row():
                    user_message_btn = gr.Button('Send')
                    stop_btn = gr.Button('Stop')
                    clear_btn = gr.Button('Clear')

            # ------------- GENERATION PARAMETERS -------------------

            with gr.Column(scale=1, min_width=80):
                with gr.Group():
                    gr.Markdown('History size')
                    history_len = gr.Slider(
                        minimum=0,
                        maximum=5,
                        value=0,
                        step=1,
                        info='Number of previous messages taken into account in history',
                        label='history_len',
                        show_label=False,
                        )

                    with gr.Group():
                        gr.Markdown('Generation parameters')
                        do_sample = gr.Checkbox(
                            value=False,
                            label='do_sample',
                            info='Activate random sampling',
                            )
                        generate_args = get_generate_args(do_sample.value)
                        do_sample.change(
                            fn=get_generate_args,
                            inputs=do_sample,
                            outputs=generate_args,
                            show_progress=False,
                            )

        rag_mode = get_rag_mode_component(db=db.value)
        k, score_threshold = get_rag_settings(rag_mode=rag_mode.value, render=False)
        rag_mode.change(
            fn=get_rag_settings,
            inputs=[rag_mode],
            outputs=[k, score_threshold],
            )
        with gr.Row():
            k.render()
            score_threshold.render()

        # ---------------- SYSTEM PROMPT AND USER MESSAGE -----------

        with gr.Accordion('Prompt', open=True):
            system_prompt = get_system_prompt_component(interactive=support_system_role.value)
            user_message_with_context = get_user_message_with_context(text='', rag_mode=rag_mode.value)

        # ---------------- SEND, CLEAR AND STOP BUTTONS ------------

        generate_event = gr.on(
            triggers=[user_message.submit, user_message_btn.click],
            fn=user_message_to_chatbot,
            inputs=[user_message, chatbot],
            outputs=[user_message, chatbot],
            queue=False,
        ).then(
            fn=update_user_message_with_context,
            inputs=[chatbot, rag_mode, db, k, score_threshold],
            outputs=[user_message_with_context],
        ).then(
            fn=get_user_message_with_context,
            inputs=[user_message_with_context, rag_mode],
            outputs=[user_message_with_context],
        ).then(
            fn=get_llm_response,
            inputs=[chatbot, llm_model, user_message_with_context, rag_mode, system_prompt,
                    support_system_role, history_len, do_sample, *generate_args],
            outputs=[chatbot],
        )

        stop_btn.click(
            fn=None,
            inputs=None,
            outputs=None,
            cancels=generate_event,
            queue=False,
        )

        clear_btn.click(
            fn=lambda: (None, ''),
            inputs=None,
            outputs=[chatbot, user_message_with_context],
            queue=False,
            )



    # ================= FILE DOWNLOAD PAGE =========================

    with gr.Tab(label='Load documents'):
        with gr.Row(variant='compact'):
            upload_files = gr.File(file_count='multiple', label='Loading text files')
            web_links = gr.Textbox(lines=6, label='Links to Web sites or YouTube')

        with gr.Row(variant='compact'):
            chunk_size = gr.Slider(50, 2000, value=500, step=50, label='Chunk size')
            chunk_overlap = gr.Slider(0, 200, value=20, step=10, label='Chunk overlap')

            subtitles_lang = gr.Radio(
                SUBTITLES_LANGUAGES,
                value=SUBTITLES_LANGUAGES[0],
                label='YouTube subtitle language',
                )

        load_documents_btn = gr.Button(value='Upload documents and initialize database')
        load_docs_log = gr.Textbox(label='Status of loading and splitting documents', interactive=False)

        load_documents_btn.click(
            fn=load_documents_and_create_db,
            inputs=[upload_files, web_links, subtitles_lang, chunk_size, chunk_overlap, embed_model],
            outputs=[documents, db, load_docs_log],
        ).success(
            fn=get_rag_mode_component,
            inputs=[db],
            outputs=[rag_mode],
        )

        gr.HTML("""<h3 style='text-align: center'>
        <a href="https://github.com/sergey21000/chatbot-rag" target='_blank'>GitHub Repository</a></h3>
        """)



    # ================= VIEW PAGE FOR ALL DOCUMENTS =================

    with gr.Tab(label='View documents'):
        view_documents_btn = gr.Button(value='Show downloaded text chunks')
        view_documents_textbox = gr.Textbox(
            lines=1,
            placeholder='To view chunks, load documents in the Load documents tab',
            label='Uploaded chunks',
            )
        sep = '=' * 20
        view_documents_btn.click(
            lambda documents: f'\n{sep}\n\n'.join([doc.page_content for doc in documents]),
            inputs=[documents],
            outputs=[view_documents_textbox],
        )


    # ============== GGUF MODELS DOWNLOAD PAGE =====================

    with gr.Tab('Load LLM model'):
        new_llm_model_repo = gr.Textbox(
            value='',
            label='Add repository',
            placeholder='Link to repository of HF models in GGUF format',
            )
        new_llm_model_repo_btn = gr.Button('Add repository')
        curr_llm_model_repo = gr.Dropdown(
            choices=LLM_MODEL_REPOS,
            value=None,
            label='HF Model Repository',
            )
        curr_llm_model_path = gr.Dropdown(
            choices=[],
            value=None,
            label='GGUF model file',
            )
        load_llm_model_btn = gr.Button('Loading and initializing model')
        load_llm_model_log = gr.Textbox(
            value=f'Model {LLM_MODEL_REPOS[0]} loaded at application startup',
            label='Model loading status',
            lines=6,
            )

        with gr.Group():
            gr.Markdown('Free up disk space by deleting all models except the currently selected one')
            clear_llm_folder_btn = gr.Button('Clear folder')

        new_llm_model_repo_btn.click(
            fn=add_new_model_repo,
            inputs=[new_llm_model_repo, llm_model_repos],
            outputs=[curr_llm_model_repo, load_llm_model_log],
        ).success(
            fn=lambda: '',
            inputs=None,
            outputs=[new_llm_model_repo],
        )

        curr_llm_model_repo.change(
            fn=get_gguf_model_names,
            inputs=[curr_llm_model_repo],
            outputs=[curr_llm_model_path],
        )

        load_llm_model_btn.click(
            fn=load_llm_model,
            inputs=[curr_llm_model_repo, curr_llm_model_path],
            outputs=[llm_model, support_system_role, load_llm_model_log],
            queue=True,
        ).success(
            fn=lambda log: log + get_memory_usage(),
            inputs=[load_llm_model_log],
            outputs=[load_llm_model_log],
        ).then(
            fn=get_system_prompt_component,
            inputs=[support_system_role],
            outputs=[system_prompt],
        )

        clear_llm_folder_btn.click(
            fn=clear_llm_folder,
            inputs=[curr_llm_model_path],
            outputs=None,
        ).success(
            fn=lambda model_path: f'Models other than {model_path} removed',
            inputs=[curr_llm_model_path],
            outputs=None,
        )


    # ============== EMBEDDING MODELS DOWNLOAD PAGE =============

    with gr.Tab('Load embed model'):
        new_embed_model_repo = gr.Textbox(
            value='',
            label='Add repository',
            placeholder='Link to HF model repository',
            )
        new_embed_model_repo_btn = gr.Button('Add repository')
        curr_embed_model_repo = gr.Dropdown(
            choices=EMBED_MODEL_REPOS,
            value=None,
            label='HF model repository',
            )

        load_embed_model_btn = gr.Button('Loading and initializing model')
        load_embed_model_log = gr.Textbox(
            value=f'Model {EMBED_MODEL_REPOS[0]} loaded at application startup',
            label='Model loading status',
            lines=7,
            )
        with gr.Group():
            gr.Markdown('Free up disk space by deleting all models except the currently selected one')
            clear_embed_folder_btn = gr.Button('Clear folder')

        new_embed_model_repo_btn.click(
            fn=add_new_model_repo,
            inputs=[new_embed_model_repo, embed_model_repos],
            outputs=[curr_embed_model_repo, load_embed_model_log],
        ).success(
            fn=lambda: '',
            inputs=None,
            outputs=new_embed_model_repo,
        )

        load_embed_model_btn.click(
            fn=load_embed_model,
            inputs=[curr_embed_model_repo],
            outputs=[embed_model, load_embed_model_log],
        ).success(
            fn=lambda log: log + get_memory_usage(),
            inputs=[load_embed_model_log],
            outputs=[load_embed_model_log],
        )

        clear_embed_folder_btn.click(
            fn=clear_embed_folder,
            inputs=[curr_embed_model_repo],
            outputs=None,
        ).success(
            fn=lambda model_repo: f'Models other than {model_repo} removed',
            inputs=[curr_embed_model_repo],
            outputs=None,
        )


interface.launch(server_name='0.0.0.0', server_port=7860)  # debug=True