#
# Copyright (C) Hadad <hadad@linuxmail.org>
# All rights reserved.
#
# This code is made available under the Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0) License.
# You are free to share and adapt the code for non-commercial purposes, as long as you provide appropriate credit,
# do not use it for commercial purposes, and distribute your contributions under the same license.
#
# Contributions can be made by directly submitting pull requests.
#
# For inquiries or permission requests, please contact hadad@linuxmail.org.
#
# License: Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)
#

import gradio as gr
import requests
import json
import os
from dotenv import load_dotenv
import threading
import random
import time

LINUX_SERVER_HOST = os.getenv("LINUX_SERVER_HOST")
LINUX_SERVER_PROVIDER_KEY = [key for key in json.loads(os.getenv("LINUX_SERVER_PROVIDER_KEY", "[]")) if key]

AI_TYPES = {f"AI_TYPE_{i}": os.getenv(f"AI_TYPE_{i}") for i in range(1, 6)}
RESPONSES = {f"RESPONSE_{i}": os.getenv(f"RESPONSE_{i}") for i in range(1, 10)}

MODEL_MAPPING = json.loads(os.getenv("MODEL_MAPPING", "{}"))
MODEL_CONFIG = json.loads(os.getenv("MODEL_CONFIG", "{}"))
MODEL_CHOICES = list(MODEL_MAPPING.values())
DEFAULT_CONFIG = json.loads(os.getenv("DEFAULT_CONFIG", "{}"))

META_TAGS = os.getenv("META_TAGS")

stop_event = threading.Event()
session = requests.Session()

def get_model_key(display_name):
    return next((k for k, v in MODEL_MAPPING.items() if v == display_name), MODEL_CHOICES[0])

def simulate_streaming_response(text):
    for line in text.splitlines():
        if stop_event.is_set():
            return
        yield line + "\n"
        time.sleep(0.05)

def chat_with_model(history, user_input, selected_model_display):
    if stop_event.is_set():
        yield RESPONSES["RESPONSE_1"]
        return

    if not LINUX_SERVER_PROVIDER_KEY or not LINUX_SERVER_HOST:
        yield RESPONSES["RESPONSE_3"]
        return

    selected_model = get_model_key(selected_model_display)
    model_config = MODEL_CONFIG.get(selected_model, DEFAULT_CONFIG)

    messages = [{"role": "user", "content": user} for user, _ in history]
    messages += [{"role": "assistant", "content": assistant} for _, assistant in history if assistant]
    messages.append({"role": "user", "content": user_input})

    data = {"model": selected_model, "messages": messages, **model_config}
    random.shuffle(LINUX_SERVER_PROVIDER_KEY)

    for api_key in LINUX_SERVER_PROVIDER_KEY[:2]:
        if stop_event.is_set():
            yield RESPONSES["RESPONSE_1"]
            return
        try:
            response = session.post(LINUX_SERVER_HOST, json=data, headers={"Authorization": f"Bearer {api_key}"})
            if stop_event.is_set():
                yield RESPONSES["RESPONSE_1"]
                return
            if response.status_code < 400:
                ai_text = response.json().get("choices", [{}])[0].get("message", {}).get("content", RESPONSES["RESPONSE_2"])
                yield from simulate_streaming_response(ai_text)
                return
        except requests.exceptions.RequestException:
            continue

    yield RESPONSES["RESPONSE_3"]

def respond(user_input, history, selected_model_display):
    if not user_input.strip():
        yield history, gr.update(value=""), gr.update(visible=False, interactive=False), gr.update(visible=True)
        return

    stop_event.clear()
    history.append([user_input, RESPONSES["RESPONSE_8"]])

    yield history, gr.update(value=""), gr.update(visible=False), gr.update(visible=True)

    ai_response = ""
    for chunk in chat_with_model(history, user_input, selected_model_display):
        if stop_event.is_set():
            session.close()
            history[-1][1] = RESPONSES["RESPONSE_1"]
            yield history, gr.update(value=""), gr.update(visible=True), gr.update(visible=False)
            return
        ai_response += chunk
        history[-1][1] = ai_response
        yield history, gr.update(value=""), gr.update(visible=False), gr.update(visible=True)

    session.close()
    yield history, gr.update(value=""), gr.update(visible=True), gr.update(visible=False)

def stop_response():
    stop_event.set()
    session.close()

def change_model(new_model_display):
    return [], new_model_display

def check_send_button_enabled(msg):
    return gr.update(visible=bool(msg.strip()), interactive=bool(msg.strip()))

with gr.Blocks(fill_height=True, fill_width=True, title=AI_TYPES["AI_TYPE_4"], head=META_TAGS) as demo:
    user_history = gr.State([])
    selected_model = gr.State(MODEL_CHOICES[0])

    chatbot = gr.Chatbot(label=AI_TYPES["AI_TYPE_1"], show_copy_button=True, show_share_button=False, scale=1, elem_id=AI_TYPES["AI_TYPE_2"])
    model_dropdown = gr.Dropdown(label=AI_TYPES["AI_TYPE_3"], show_label=False, choices=MODEL_CHOICES, value=MODEL_CHOICES[0], interactive=True)
    msg = gr.Textbox(label=RESPONSES["RESPONSE_4"], show_label=False, scale=0, placeholder=RESPONSES["RESPONSE_5"])

    with gr.Row():
        send_btn = gr.Button(RESPONSES["RESPONSE_6"], visible=True, interactive=False)
        stop_btn = gr.Button(RESPONSES["RESPONSE_7"], variant=RESPONSES["RESPONSE_9"], visible=False)

    model_dropdown.change(fn=change_model, inputs=[model_dropdown], outputs=[user_history, selected_model])
    send_btn.click(respond, inputs=[msg, user_history, selected_model], outputs=[chatbot, msg, send_btn, stop_btn])
    msg.change(fn=check_send_button_enabled, inputs=[msg], outputs=[send_btn])
    stop_btn.click(fn=stop_response, outputs=[send_btn, stop_btn])

demo.launch(show_api=False)