import json
import os
import gradio as gr
from huggingface_hub import InferenceClient

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# ---- Load config từ model.json ----
with open(os.path.join(BASE_DIR, "model.json"), "r", encoding="utf-8") as f:
    MODEL_CFG = json.load(f)

MODEL_ID = MODEL_CFG["model_id"]

# ---- Load system prompt mặc định từ systemprompt.md ----
with open(os.path.join(BASE_DIR, "systemprompt.md"), "r", encoding="utf-8") as f:
    DEFAULT_SYSTEM_PROMPT = f.read().strip()

# ---- Load banner HTML ----
with open(os.path.join(BASE_DIR, "public", "index.html"), "r", encoding="utf-8") as f:
    BANNER_HTML = f.read()


def respond(
    message,
    history: list[dict[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    hf_token: gr.OAuthToken,
):
    client = InferenceClient(token=hf_token.token, model=MODEL_ID)

    messages = [{"role": "system", "content": system_message}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})

    response = ""
    for chunk in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        choices = chunk.choices
        token = ""
        if len(choices) and choices[0].delta.content:
            token = choices[0].delta.content
        response += token
        yield response


# ---- CSS tùy chỉnh ----
custom_css = """
#chatbot-container { max-width: 900px; margin: 0 auto; }
.gradio-container { font-family: 'Inter', 'Segoe UI', sans-serif !important; }
.message-wrap { border-radius: 14px !important; }
footer { visibility: hidden }
"""

# ---- JS: đọc localStorage khi load trang ----
LOAD_HISTORY_JS = """
() => {
    try {
        const raw = window.localStorage.getItem("luau_chat_history");
        return raw ? raw : "[]";
    } catch (e) {
        return "[]";
    }
}
"""

# ---- JS: lưu lịch sử vào localStorage sau mỗi lượt chat ----
SAVE_HISTORY_JS = """
(history) => {
    try {
        window.localStorage.setItem("luau_chat_history", JSON.stringify(history));
    } catch (e) {
        console.error("Không thể lưu lịch sử chat:", e);
    }
    return history;
}
"""

# ---- JS: xoá lịch sử ----
CLEAR_HISTORY_JS = """
() => {
    try {
        window.localStorage.removeItem("luau_chat_history");
    } catch (e) {}
    return [];
}
"""


def parse_stored_history(raw_json):
    try:
        data = json.loads(raw_json)
        if isinstance(data, list):
            return data
    except Exception:
        pass
    return []


chatbot_ui = gr.Chatbot(
    label=MODEL_CFG.get("display_name", "Roblox Luau Assistant"),
    height=550,
    show_copy_button=True,
    avatar_images=(None, MODEL_CFG.get("avatar_url")),
    bubble_full_width=False,
    render_markdown=True,
    type="messages",
)

with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="violet")) as demo:
    gr.HTML(BANNER_HTML)

    with gr.Sidebar():
        gr.Markdown("### 🔑 Tài khoản")
        gr.LoginButton()

    # State ẩn để chứa dữ liệu lịch sử đọc từ localStorage lúc khởi động
    stored_history_json = gr.Textbox(visible=False)

    with gr.Column(elem_id="chatbot-container"):
        chat = gr.ChatInterface(
            respond,
            chatbot=chatbot_ui,
            type="messages",
            additional_inputs=[
                gr.Textbox(
                    value=DEFAULT_SYSTEM_PROMPT,
                    label="System message",
                    lines=4,
                ),
                gr.Slider(
                    minimum=1,
                    maximum=MODEL_CFG.get("max_tokens_limit", 4096),
                    value=MODEL_CFG.get("default_max_tokens", 1024),
                    step=1,
                    label="Max new tokens",
                ),
                gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=MODEL_CFG.get("default_temperature", 0.7),
                    step=0.1,
                    label="Temperature",
                ),
                gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=MODEL_CFG.get("default_top_p", 0.95),
                    step=0.05,
                    label="Top-p (nucleus sampling)",
                ),
            ],
            additional_inputs_accordion=gr.Accordion(label="⚙️ Cài đặt nâng cao", open=False),
            examples=[
                ["Viết một script Luau spawn item ngẫu nhiên mỗi 10 giây"],
                ["Giải thích RemoteEvent và cách dùng an toàn trong Roblox"],
                ["Sửa lỗi trong đoạn code Luau sau: ..."],
            ],
            cache_examples=False,
        )

        clear_btn = gr.Button("🗑️ Xoá lịch sử chat (localStorage)")

    # Khi trang load: đọc localStorage -> đổ vào chatbot
    def _load_into_chatbot(raw_json):
        return parse_stored_history(raw_json)

    demo.load(
        fn=None,
        inputs=None,
        outputs=stored_history_json,
        js=LOAD_HISTORY_JS,
    ).then(
        fn=_load_into_chatbot,
        inputs=stored_history_json,
        outputs=chatbot_ui,
    )

    # Sau mỗi lần chatbot_ui thay đổi (có phản hồi mới) -> lưu xuống localStorage
    chatbot_ui.change(
        fn=None,
        inputs=chatbot_ui,
        outputs=None,
        js=SAVE_HISTORY_JS,
    )

    # Nút xoá lịch sử: xoá localStorage + xoá chatbot trên UI
    clear_btn.click(
        fn=None,
        inputs=None,
        outputs=chatbot_ui,
        js=CLEAR_HISTORY_JS,
    )

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