import argparse
import os
import subprocess
import time

import gradio as gr
from huggingface_hub import snapshot_download
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import (
    Docx2txtLoader,
    PyPDFLoader,
    TextLoader,
    YoutubeLoader,
)
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from openai import OpenAI

from chat_with_mlx import __version__
from chat_with_mlx.models.utils import model_info
from chat_with_mlx.rag.utils import get_prompt

os.environ["TOKENIZERS_PARALLELISM"] = "False"
SUPPORTED_LANG = [
    "default",
    "English",
    "Spanish",
    "Chinese",
    "Vietnamese",
    "Japanese",
    "Korean",
    "Indian",
    "Turkish",
    "German",
    "French",
    "Italian",
]
openai_api_base = "http://127.0.0.1:8080/v1"
model_dicts, yml_path, cfg_list, mlx_config = model_info()
model_list = list(cfg_list.keys())
client = OpenAI(api_key="EMPTY", base_url=openai_api_base)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=50)
emb = HuggingFaceEmbeddings(
    model_name="nomic-ai/nomic-embed-text-v1.5",
    model_kwargs={"trust_remote_code": True},
)
vectorstore = None


def load_model(model_name, lang):
    global process, rag_prompt, rag_his_prompt, sys_prompt, default_lang
    default_lang = "default"
    prompts, sys_prompt = get_prompt(f"{yml_path[cfg_list[model_name]]}", lang)
    rag_prompt, rag_his_prompt = prompts[0], prompts[1]
    model_name_list = cfg_list[model_name].split("/")
    directory_path = os.path.dirname(os.path.abspath(__file__))
    local_model_dir = os.path.join(
        directory_path, "models", "download", model_name_list[1]
    )

    if not os.path.exists(local_model_dir):
        snapshot_download(repo_id=mlx_config[model_name], local_dir=local_model_dir)

    command = ["python3", "-m", "mlx_lm.server", "--model", local_model_dir]

    try:
        process = subprocess.Popen(
            command, stdin=subprocess.PIPE, stderr=subprocess.PIPE, text=True
        )
        process.stdin.write("y\n")
        process.stdin.flush()
        return {model_status: "Model Loaded"}
    except Exception as e:
        return {model_status: f"Exception occurred: {str(e)}"}


def kill_process():
    global process
    process.terminate()
    time.sleep(2)
    if process.poll() is None:  # Check if the process has indeed terminated
        process.kill()  # Force kill if still running

    print("Model Killed")
    return {model_status: "Model Unloaded"}


def check_file_type(file_path):
    # Check for document file extensions
    if (
        file_path.endswith(".pdf")
        or file_path.endswith(".txt")
        or file_path.endswith(".doc")
        or file_path.endswith(".docx")
    ):
        return True
        # Check for YouTube link formats
    elif (
        file_path.startswith("https://www.youtube.com/")
        or file_path.startswith("https://youtube.com/")
        or file_path.startswith("https://youtu.be/")
    ):
        return True
    else:
        return False


def upload(files):
    supported = check_file_type(files)
    if supported:
        return {url: files, index_status: "Not Done"}
    else:
        return {url: "File type not supported", index_status: "Not Done"}


def indexing(mode, url):
    global vectorstore

    try:
        if mode == "Files (docx, pdf, txt)":
            if url.endswith(".pdf"):
                loader = PyPDFLoader(url)
            elif url.endswith(".docx"):
                loader = Docx2txtLoader(url)
            elif url.endswith(".txt"):
                loader = TextLoader(url)
            splits = loader.load_and_split(text_splitter)
        elif mode == "YouTube (url)":
            loader = YoutubeLoader.from_youtube_url(
                url, add_video_info=False, language=["en", "vi"]
            )
            splits = loader.load_and_split(text_splitter)

        vectorstore = Chroma.from_documents(documents=splits, embedding=emb)
        return {index_status: "Indexing Done"}
    except Exception as e:
        # Print the error message or return it as part of the response
        print(f"Error: {e}")  # This will print the error to the console or log
        return {"index_status": "Indexing Error", "error_message": str(e)}


def kill_index():
    global vectorstore
    vectorstore = None
    return {index_status: "Indexing Undone"}


def build_rag_context(docs):
    context = ""
    for doc in docs:
        context += doc.page_content + "\n"

    return context


def chatbot(query, history, temp, max_tokens, freq_penalty, k_docs):
    global chat_history, sys_prompt

    if "vectorstore" in globals() and vectorstore is not None:
        if len(history) == 0:
            chat_history = []
            if sys_prompt is not None:
                chat_history.append({"role": "system", "content": sys_prompt})
            docs = vectorstore.similarity_search(query, k=k_docs)
        else:
            history_str = ""
            for i, message in enumerate(history):
                history_str += f"User: {message[0]}\n"
                history_str += f"AI: {message[1]}\n"

            if sys_prompt is not None:
                chat_history.append({"role": "system", "content": sys_prompt})
            chat_history.append({"role": "user", "content": history_str})
            docs = vectorstore.similarity_search(history_str)

        context = build_rag_context(docs)

        if len(history) == 0:
            prompt = rag_prompt.format(context=context, question=query)
        else:
            prompt = rag_his_prompt.format(
                chat_history=history_str, context=context, question=query
            )
        messages = [{"role": "user", "content": prompt}]
    else:
        if len(history) == 0:
            chat_history = []
            if sys_prompt is not None:
                chat_history.append({"role": "system", "content": sys_prompt})
        else:
            chat_history = []
            if sys_prompt is not None:
                chat_history.append({"role": "system", "content": sys_prompt})
            for i, message in enumerate(history):
                chat_history.append({"role": "user", "content": message[0]})
                chat_history.append({"role": "assistant", "content": message[1]})
        chat_history.append({"role": "user", "content": query})
        messages = chat_history

    # Uncomment for debugging
    # print(messages)

    response = client.chat.completions.create(
        model="gpt",
        messages=messages,
        temperature=temp,
        frequency_penalty=freq_penalty,
        max_tokens=max_tokens,
        stream=True,
    )
    stop = ["<|im_end|>", "<|endoftext|>"]
    partial_message = ""
    for chunk in response:
        if len(chunk.choices) != 0:
            if chunk.choices[0].delta.content not in stop:
                partial_message = partial_message + chunk.choices[0].delta.content
            else:
                partial_message = partial_message + ""
            yield partial_message


with gr.Blocks(fill_height=True, theme=gr.themes.Soft()) as demo:
    model_name = gr.Dropdown(
        label="Model",
        info="Select your model",
        choices=sorted(model_list),
        interactive=True,
        render=False,
    )
    temp_slider = gr.State(0.2)
    max_gen_token = gr.State(512)
    freq_penalty = gr.State(1.05)
    retrieve_docs = gr.State(3)
    language = gr.State("default")
    gr.ChatInterface(
        chatbot=gr.Chatbot(height=600, render=False),
        fn=chatbot,  # Function to call on user input
        title="Chat with MLX🍎",  # Title of the web page
        description="Chat with your data using Apple MLX Backend",  # Description
        additional_inputs=[temp_slider, max_gen_token, freq_penalty, retrieve_docs],
    )
    with gr.Accordion("Advanced Setting", open=False):
        with gr.Row():
            with gr.Column(scale=2):
                temp_slider = gr.Slider(
                    label="Temperature",
                    value=0.2,
                    minimum=0.0,
                    maximum=1.0,
                    step=0.05,
                    interactive=True,
                )
                max_gen_token = gr.Slider(
                    label="Max Tokens",
                    value=512,
                    minimum=512,
                    maximum=4096,
                    step=256,
                    interactive=True,
                )
            with gr.Column(scale=2):
                freq_penalty = gr.Slider(
                    label="Frequency Penalty",
                    value=1.05,
                    minimum=-2,
                    maximum=2,
                    step=0.05,
                    interactive=True,
                )
                retrieve_docs = gr.Slider(
                    label="No. Retrieval Docs",
                    value=3,
                    minimum=1,
                    maximum=10,
                    step=1,
                    interactive=True,
                )

    with gr.Row():
        with gr.Column(scale=2):
            model_name.render()
            language = gr.Dropdown(
                label="Language",
                choices=sorted(SUPPORTED_LANG),
                value="default",
                interactive=True,
            )
            btn1 = gr.Button("Load Model", variant="primary")
            btn3 = gr.Button("Unload Model", variant="stop")
        with gr.Column(scale=4):
            with gr.Row():
                with gr.Column(scale=9):
                    mode = gr.Dropdown(
                        label="Dataset",
                        info="Choose your dataset type",
                        choices=["Files (docx, pdf, txt)", "YouTube (url)"],
                        scale=5,
                    )
                    url = gr.Textbox(
                        label="URL",
                        info="Enter your filepath (URL for Youtube)",
                        interactive=True,
                    )
                    upload_button = gr.UploadButton(
                        label="Upload File", variant="primary"
                    )

                # data = gr.Textbox(visible=lambda mode: mode == 'YouTube')
                with gr.Column(scale=1):
                    model_status = gr.Textbox("Model Not Loaded", label="Model Status")
                    index_status = gr.Textbox("Not Index", label="Index Status")
                    btn1.click(
                        load_model,
                        inputs=[model_name, language],
                        outputs=[model_status],
                    )
                    btn3.click(kill_process, outputs=[model_status])
                    upload_button.upload(
                        upload, inputs=upload_button, outputs=[url, index_status]
                    )

                    index_button = gr.Button("Start Indexing", variant="primary")
                    index_button.click(
                        indexing, inputs=[mode, url], outputs=[index_status]
                    )
                    stop_index_button = gr.Button("Stop Indexing")
                    stop_index_button.click(kill_index, outputs=[index_status])


def app(port, share):
    print(f"Starting MLX Chat on port {port}")
    print(f"Sharing: {share}")
    demo.launch(inbrowser=True, share=share, server_port=port)


def main():
    parser = argparse.ArgumentParser(
        description="Chat with MLX \n"
        "Native RAG on MacOS and Apple Silicon with MLX 🧑‍💻"
    )
    parser.add_argument(
        "--version", action="version", version=f"Chat with MLX {__version__}"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=7860,
        help="Port number to run the app",
    )
    parser.add_argument(
        "--share",
        default=False,
        help="Enable sharing the app",
    )
    args = parser.parse_args()
    app(port=args.port, share=args.share)

langchain_community.document_loaders.PyPDFLoader langchain_community.document_loaders.Docx2txtLoader langchain_community.document_loaders.YoutubeLoader.from_youtube_url langchain.text_splitter.RecursiveCharacterTextSplitter langchain_community.vectorstores.Chroma.from_documents langchain_community.document_loaders.TextLoader langchain_community.embeddings.HuggingFaceEmbeddings
[((934, 946), 'chat_with_mlx.models.utils.model_info', 'model_info', ([], {}), '()\n', (944, 946), False, 'from chat_with_mlx.models.utils import model_info\n'), ((991, 1040), 'openai.OpenAI', 'OpenAI', ([], {'api_key': '"""EMPTY"""', 'base_url': 'openai_api_base'}), "(api_key='EMPTY', base_url=openai_api_base)\n", (997, 1040), False, 'from openai import OpenAI\n'), ((1057, 1121), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(512)', 'chunk_overlap': '(50)'}), '(chunk_size=512, chunk_overlap=50)\n', (1087, 1121), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((1128, 1240), 'langchain_community.embeddings.HuggingFaceEmbeddings', 'HuggingFaceEmbeddings', ([], {'model_name': '"""nomic-ai/nomic-embed-text-v1.5"""', 'model_kwargs': "{'trust_remote_code': True}"}), "(model_name='nomic-ai/nomic-embed-text-v1.5',\n    model_kwargs={'trust_remote_code': True})\n", (1149, 1240), False, 'from langchain_community.embeddings import HuggingFaceEmbeddings\n'), ((1431, 1484), 'chat_with_mlx.rag.utils.get_prompt', 'get_prompt', (['f"""{yml_path[cfg_list[model_name]]}"""', 'lang'], {}), "(f'{yml_path[cfg_list[model_name]]}', lang)\n", (1441, 1484), False, 'from chat_with_mlx.rag.utils import get_prompt\n'), ((1681, 1751), 'os.path.join', 'os.path.join', (['directory_path', '"""models"""', '"""download"""', 'model_name_list[1]'], {}), "(directory_path, 'models', 'download', model_name_list[1])\n", (1693, 1751), False, 'import os\n'), ((2378, 2391), 'time.sleep', 'time.sleep', (['(2)'], {}), '(2)\n', (2388, 2391), False, 'import time\n'), ((7394, 7407), 'gradio.State', 'gr.State', (['(0.2)'], {}), '(0.2)\n', (7402, 7407), True, 'import gradio as gr\n'), ((7428, 7441), 'gradio.State', 'gr.State', (['(512)'], {}), '(512)\n', (7436, 7441), True, 'import gradio as gr\n'), ((7461, 7475), 'gradio.State', 'gr.State', (['(1.05)'], {}), '(1.05)\n', (7469, 7475), True, 'import gradio as gr\n'), ((7496, 7507), 'gradio.State', 'gr.State', (['(3)'], {}), '(3)\n', (7504, 7507), True, 'import gradio as gr\n'), ((7523, 7542), 'gradio.State', 'gr.State', (['"""default"""'], {}), "('default')\n", (7531, 7542), True, 'import gradio as gr\n'), ((11507, 11620), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Chat with MLX \nNative RAG on MacOS and Apple Silicon with MLX 🧑\u200d💻"""'}), '(description=\n    """Chat with MLX \nNative RAG on MacOS and Apple Silicon with MLX 🧑\u200d💻""")\n', (11530, 11620), False, 'import argparse\n'), ((1632, 1657), 'os.path.abspath', 'os.path.abspath', (['__file__'], {}), '(__file__)\n', (1647, 1657), False, 'import os\n'), ((1778, 1809), 'os.path.exists', 'os.path.exists', (['local_model_dir'], {}), '(local_model_dir)\n', (1792, 1809), False, 'import os\n'), ((1819, 1895), 'huggingface_hub.snapshot_download', 'snapshot_download', ([], {'repo_id': 'mlx_config[model_name]', 'local_dir': 'local_model_dir'}), '(repo_id=mlx_config[model_name], local_dir=local_model_dir)\n', (1836, 1895), False, 'from huggingface_hub import snapshot_download\n'), ((2002, 2089), 'subprocess.Popen', 'subprocess.Popen', (['command'], {'stdin': 'subprocess.PIPE', 'stderr': 'subprocess.PIPE', 'text': '(True)'}), '(command, stdin=subprocess.PIPE, stderr=subprocess.PIPE,\n    text=True)\n', (2018, 2089), False, 'import subprocess\n'), ((4005, 4059), 'langchain_community.vectorstores.Chroma.from_documents', 'Chroma.from_documents', ([], {'documents': 'splits', 'embedding': 'emb'}), '(documents=splits, embedding=emb)\n', (4026, 4059), False, 'from langchain_community.vectorstores import Chroma\n'), ((7912, 7956), 'gradio.Accordion', 'gr.Accordion', (['"""Advanced Setting"""'], {'open': '(False)'}), "('Advanced Setting', open=False)\n", (7924, 7956), True, 'import gradio as gr\n'), ((9136, 9144), 'gradio.Row', 'gr.Row', ([], {}), '()\n', (9142, 9144), True, 'import gradio as gr\n'), ((7172, 7188), 'gradio.themes.Soft', 'gr.themes.Soft', ([], {}), '()\n', (7186, 7188), True, 'import gradio as gr\n'), ((7581, 7617), 'gradio.Chatbot', 'gr.Chatbot', ([], {'height': '(600)', 'render': '(False)'}), '(height=600, render=False)\n', (7591, 7617), True, 'import gradio as gr\n'), ((7971, 7979), 'gradio.Row', 'gr.Row', ([], {}), '()\n', (7977, 7979), True, 'import gradio as gr\n'), ((9159, 9177), 'gradio.Column', 'gr.Column', ([], {'scale': '(2)'}), '(scale=2)\n', (9168, 9177), True, 'import gradio as gr\n'), ((9429, 9471), 'gradio.Button', 'gr.Button', (['"""Load Model"""'], {'variant': '"""primary"""'}), "('Load Model', variant='primary')\n", (9438, 9471), True, 'import gradio as gr\n'), ((9491, 9532), 'gradio.Button', 'gr.Button', (['"""Unload Model"""'], {'variant': '"""stop"""'}), "('Unload Model', variant='stop')\n", (9500, 9532), True, 'import gradio as gr\n'), ((9546, 9564), 'gradio.Column', 'gr.Column', ([], {'scale': '(4)'}), '(scale=4)\n', (9555, 9564), True, 'import gradio as gr\n'), ((3514, 3530), 'langchain_community.document_loaders.PyPDFLoader', 'PyPDFLoader', (['url'], {}), '(url)\n', (3525, 3530), False, 'from langchain_community.document_loaders import Docx2txtLoader, PyPDFLoader, TextLoader, YoutubeLoader\n'), ((3813, 3898), 'langchain_community.document_loaders.YoutubeLoader.from_youtube_url', 'YoutubeLoader.from_youtube_url', (['url'], {'add_video_info': '(False)', 'language': "['en', 'vi']"}), "(url, add_video_info=False, language=['en', 'vi']\n    )\n", (3843, 3898), False, 'from langchain_community.document_loaders import Docx2txtLoader, PyPDFLoader, TextLoader, YoutubeLoader\n'), ((7998, 8016), 'gradio.Column', 'gr.Column', ([], {'scale': '(2)'}), '(scale=2)\n', (8007, 8016), True, 'import gradio as gr\n'), ((8048, 8149), 'gradio.Slider', 'gr.Slider', ([], {'label': '"""Temperature"""', 'value': '(0.2)', 'minimum': '(0.0)', 'maximum': '(1.0)', 'step': '(0.05)', 'interactive': '(True)'}), "(label='Temperature', value=0.2, minimum=0.0, maximum=1.0, step=\n    0.05, interactive=True)\n", (8057, 8149), True, 'import gradio as gr\n'), ((8316, 8416), 'gradio.Slider', 'gr.Slider', ([], {'label': '"""Max Tokens"""', 'value': '(512)', 'minimum': '(512)', 'maximum': '(4096)', 'step': '(256)', 'interactive': '(True)'}), "(label='Max Tokens', value=512, minimum=512, maximum=4096, step=\n    256, interactive=True)\n", (8325, 8416), True, 'import gradio as gr\n'), ((8568, 8586), 'gradio.Column', 'gr.Column', ([], {'scale': '(2)'}), '(scale=2)\n', (8577, 8586), True, 'import gradio as gr\n'), ((8619, 8723), 'gradio.Slider', 'gr.Slider', ([], {'label': '"""Frequency Penalty"""', 'value': '(1.05)', 'minimum': '(-2)', 'maximum': '(2)', 'step': '(0.05)', 'interactive': '(True)'}), "(label='Frequency Penalty', value=1.05, minimum=-2, maximum=2,\n    step=0.05, interactive=True)\n", (8628, 8723), True, 'import gradio as gr\n'), ((8891, 8991), 'gradio.Slider', 'gr.Slider', ([], {'label': '"""No. Retrieval Docs"""', 'value': '(3)', 'minimum': '(1)', 'maximum': '(10)', 'step': '(1)', 'interactive': '(True)'}), "(label='No. Retrieval Docs', value=3, minimum=1, maximum=10, step=\n    1, interactive=True)\n", (8900, 8991), True, 'import gradio as gr\n'), ((9583, 9591), 'gradio.Row', 'gr.Row', ([], {}), '()\n', (9589, 9591), True, 'import gradio as gr\n'), ((3596, 3615), 'langchain_community.document_loaders.Docx2txtLoader', 'Docx2txtLoader', (['url'], {}), '(url)\n', (3610, 3615), False, 'from langchain_community.document_loaders import Docx2txtLoader, PyPDFLoader, TextLoader, YoutubeLoader\n'), ((9614, 9632), 'gradio.Column', 'gr.Column', ([], {'scale': '(9)'}), '(scale=9)\n', (9623, 9632), True, 'import gradio as gr\n'), ((9661, 9789), 'gradio.Dropdown', 'gr.Dropdown', ([], {'label': '"""Dataset"""', 'info': '"""Choose your dataset type"""', 'choices': "['Files (docx, pdf, txt)', 'YouTube (url)']", 'scale': '(5)'}), "(label='Dataset', info='Choose your dataset type', choices=[\n    'Files (docx, pdf, txt)', 'YouTube (url)'], scale=5)\n", (9672, 9789), True, 'import gradio as gr\n'), ((9930, 10021), 'gradio.Textbox', 'gr.Textbox', ([], {'label': '"""URL"""', 'info': '"""Enter your filepath (URL for Youtube)"""', 'interactive': '(True)'}), "(label='URL', info='Enter your filepath (URL for Youtube)',\n    interactive=True)\n", (9940, 10021), True, 'import gradio as gr\n'), ((10149, 10204), 'gradio.UploadButton', 'gr.UploadButton', ([], {'label': '"""Upload File"""', 'variant': '"""primary"""'}), "(label='Upload File', variant='primary')\n", (10164, 10204), True, 'import gradio as gr\n'), ((10349, 10367), 'gradio.Column', 'gr.Column', ([], {'scale': '(1)'}), '(scale=1)\n', (10358, 10367), True, 'import gradio as gr\n'), ((10404, 10456), 'gradio.Textbox', 'gr.Textbox', (['"""Model Not Loaded"""'], {'label': '"""Model Status"""'}), "('Model Not Loaded', label='Model Status')\n", (10414, 10456), True, 'import gradio as gr\n'), ((10492, 10537), 'gradio.Textbox', 'gr.Textbox', (['"""Not Index"""'], {'label': '"""Index Status"""'}), "('Not Index', label='Index Status')\n", (10502, 10537), True, 'import gradio as gr\n'), ((10982, 11028), 'gradio.Button', 'gr.Button', (['"""Start Indexing"""'], {'variant': '"""primary"""'}), "('Start Indexing', variant='primary')\n", (10991, 11028), True, 'import gradio as gr\n'), ((11208, 11234), 'gradio.Button', 'gr.Button', (['"""Stop Indexing"""'], {}), "('Stop Indexing')\n", (11217, 11234), True, 'import gradio as gr\n'), ((3680, 3695), 'langchain_community.document_loaders.TextLoader', 'TextLoader', (['url'], {}), '(url)\n', (3690, 3695), False, 'from langchain_community.document_loaders import Docx2txtLoader, PyPDFLoader, TextLoader, YoutubeLoader\n')]