import gradio as gr
from huggingface_hub import InferenceClient

# ── Cơ chế token theo IOAI ────────────────────────────────────────────────
# Theo Technical Appendix của IOAI 2026 (mục "Language Models and AI Assistance"):
#   - Individual Contest: Gemma 3, tối đa 1000 token ĐẦU RA mỗi câu hỏi.
#   - GAITE Contest:      Gemma 4, tối đa 2000 token ĐẦU RA mỗi câu hỏi.
# => Đây là giới hạn token đầu ra cho TỪNG câu hỏi (per query), không cộng dồn.
#    Về kỹ thuật chính là tham số max_tokens: câu trả lời dài hơn sẽ bị cắt.
# --------------------------------------------------------------------------

# Danh sách model (mặc định Gemma 3 để giống Individual Contest của IOAI).
# Dropdown bật allow_custom_value=True nên bạn cũng có thể gõ model ID bất kỳ.
MODELS = [
    "google/gemma-3-27b-it",          # Gemma 3 (IOAI Individual)
    "google/gemma-3-12b-it",
    "google/gemma-2-9b-it",
    "meta-llama/Llama-3.3-70B-Instruct",
    "meta-llama/Llama-3.1-8B-Instruct",
    "Qwen/Qwen2.5-72B-Instruct",
    "Qwen/Qwen2.5-7B-Instruct",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "deepseek-ai/DeepSeek-V3",
    "HuggingFaceH4/zephyr-7b-beta",
]

# Giới hạn token đầu ra mỗi câu hỏi (IOAI Individual = 1000, GAITE = 2000)
DEFAULT_OUTPUT_LIMIT = 1000

# Dấu nhận biết dòng thống kê token (để cắt khỏi history trước khi gửi lại model)
TOKEN_FOOTER_MARKER = "📊 Token đầu ra:"


def extract_text(content):
    """Lấy phần text từ 'content' của một message.

    Gradio 6 trả content dạng list block, ví dụ [{"type": "text", "text": "..."}],
    còn các bản cũ trả thẳng một chuỗi. Hàm này xử lý được cả hai trường hợp.
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for block in content:
            if isinstance(block, dict):
                parts.append(block.get("text", ""))
            elif isinstance(block, str):
                parts.append(block)
        return "".join(parts)
    return str(content)


def strip_token_footer(text):
    """Bỏ dòng thống kê token ở cuối câu trả lời (nếu có)."""
    idx = text.rfind(TOKEN_FOOTER_MARKER)
    if idx == -1:
        return text
    sep = text.rfind("\n\n", 0, idx)
    return text[:sep].rstrip() if sep != -1 else text[:idx].rstrip()


def estimate_tokens(text):
    """Ước lượng số token khi provider không trả về 'usage' (xấp xỉ 4 ký tự/token)."""
    if not isinstance(text, str):
        text = str(text)
    return len(text) // 4


def respond(
    message,
    history: list[dict],
    model_name,
    output_limit,
    system_message,
    temperature,
    top_p,
    hf_token: gr.OAuthToken,
):
    """
    For more information on `huggingface_hub` Inference API support, please check the docs:
    https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
    """
    # Giới hạn token đầu ra cho CHÍNH câu hỏi này (giống IOAI: per query)
    max_output = int(output_limit) if output_limit else DEFAULT_OUTPUT_LIMIT

    client = InferenceClient(token=hf_token.token, model=model_name)

    # Dựng messages, cắt bỏ dòng thống kê token khỏi history
    messages = [{"role": "system", "content": system_message}]
    for turn in history:
        content = extract_text(turn["content"])
        if turn["role"] == "assistant":
            content = strip_token_footer(content)
        messages.append({"role": turn["role"], "content": content})
    messages.append({"role": "user", "content": extract_text(message)})

    response = ""
    usage = None
    finish_reason = None
    # max_tokens = giới hạn token đầu ra cho câu hỏi này -> model tự cắt khi đạt
    for chunk in client.chat_completion(
        messages,
        max_tokens=max_output,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        # Chunk cuối cùng mang 'usage' (token thật), choices rỗng
        if getattr(chunk, "usage", None) is not None:
            usage = chunk.usage

        choices = chunk.choices
        if len(choices):
            if choices[0].finish_reason:
                finish_reason = choices[0].finish_reason
            if choices[0].delta.content:
                response += choices[0].delta.content
                yield response

    # Số token ĐẦU RA đã dùng cho câu hỏi này
    if usage is not None and getattr(usage, "completion_tokens", None) is not None:
        out_tokens = usage.completion_tokens
        prefix = ""
    else:
        out_tokens = max(1, estimate_tokens(response))
        prefix = "~"

    footer = f"\n\n`{TOKEN_FOOTER_MARKER} {prefix}{out_tokens}/{max_output} (mỗi câu hỏi)`"
    # finish_reason == "length" nghĩa là câu trả lời bị cắt vì chạm giới hạn token
    if finish_reason == "length":
        footer += "\n`✂️ Câu trả lời đã chạm giới hạn token và bị cắt (giống IOAI).`"

    yield response + footer


"""
For information on how to customize the ChatInterface, peruse the gradio docs:
https://www.gradio.app/docs/chatinterface
"""
chatbot = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Dropdown(
            choices=MODELS,
            value=MODELS[0],
            label="Model",
            allow_custom_value=True,
            info="Mặc định Gemma 3 (như IOAI Individual). Có thể gõ model ID bất kỳ.",
        ),
        gr.Slider(
            minimum=1,
            maximum=4096,
            value=DEFAULT_OUTPUT_LIMIT,
            step=1,
            label="Giới hạn token đầu ra / câu hỏi",
            info="Như IOAI: Individual = 1000, GAITE = 2000. Câu trả lời dài hơn sẽ bị cắt.",
        ),
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)

with gr.Blocks() as demo:
    with gr.Sidebar():
        gr.LoginButton()
    chatbot.render()

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