import spaces
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
from huggingface_hub import HfApi
from threading import Thread
import gc
import os
import shutil
import torch
import psutil
import time
import queue as _queue

HF_CACHE_DIR = os.path.expanduser("~/.cache/huggingface/hub")
DEFAULT_MODEL = "HuggingFaceTB/SmolLM2-135M-Instruct"

MODELS = [
    "HuggingFaceTB/SmolLM2-135M-Instruct",
    "HuggingFaceTB/SmolLM2-135M",
    "HuggingFaceTB/SmolLM2-360M-Instruct",
    "HuggingFaceTB/SmolLM2-1.7B-Instruct",
    "HuggingFaceTB/SmolLM-135M",
    "Qwen/Qwen3-0.6B",
    "Qwen/Qwen2.5-Coder-0.5B",
    "Qwen/Qwen2.5-0.5B",
    "Qwen/Qwen2.5-1.5B",
    "Qwen/Qwen2.5-3B",
    "Qwen/Qwen3-1.7B",
    "facebook/MobileLLM-R1-140M-base",
    "facebook/opt-125m",
    "facebook/opt-350m",
    "microsoft/phi-2",
    "microsoft/Phi-3.5-mini-instruct",
    "microsoft/Phi-3-mini-4k-instruct",
    "openai-community/gpt2",
    "openai-community/gpt2-medium",
    "EleutherAI/pythia-70m",
    "EleutherAI/pythia-160m",
    "EleutherAI/pythia-410m",
    "EleutherAI/gpt-neo-125M",
    "EleutherAI/gpt-neo-1.3B",
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    "stabilityai/StableLM-3b-4e1t",
    "stabilityai/StableLM-Zephyr-3B",
    "NousResearch/Hermes-3-Llama-3.1-8B",
    "meta-llama/Llama-3.2-1B",
    "SupraLabs/Supra-50M-Base",
    "SupraLabs/Supra-50M-Instruct",
    "SupraLabs/Supra-50M-Reasoning",
    "GODELEV/Archaea-74M",
    "Sandroeth/cali-0.1B",
    "ThingAI/Quark-50m",
    "ThingAI/Quark-135m",
    "Aravindan/awesome-gpt-2-coder",
    "LiquidAI/LFM2-1.2B",
    "LiquidAI/LFM2-2.6B",
    "LiquidAI/LFM2.5-230M",
    "LiquidAI/LFM2.5-350M",
    "LiquidAI/LFM2.5-1.2B-Instruct",
    "LiquidAI/LFM2.5-1.2B-Thinking",
    "LiquidAI/LFM2.5-8B-A1B",
]

TASK_MODES = ["Completion", "Chat", "Q&A", "Translation"]

LANGUAGES = [
    "English", "Spanish", "French", "German", "Italian", "Portuguese",
    "Chinese", "Japanese", "Korean", "Arabic", "Russian", "Hindi",
    "Dutch", "Turkish", "Polish", "Czech", "Romanian", "Greek",
    "Thai", "Vietnamese", "Indonesian", "Malay", "Finnish", "Swedish",
    "Norwegian", "Danish"
]

ACTIVE_SESSIONS = {}
SESSION_TIMEOUT = 60

THINKING_PATTERNS = [
    # (start_token, end_token, thinking_label, answer_label)
    ("<|begin_of_thought|>", "<|end_of_thought|>", "Thinking Process:", "Final Answer:"),
    ("<think>", "</think>", "Thinking:", "Answer:"),
    ("<thinking>", "</thinking>", "Thinking:", "Answer:"),
]


def live_count(request: gr.Request):
    current_time = time.time()
    if request:
        ACTIVE_SESSIONS[request.session_hash] = current_time
    expired = [s for s, t in ACTIVE_SESSIONS.items() if current_time - t > SESSION_TIMEOUT]
    for s in expired:
        ACTIVE_SESSIONS.pop(s, None)
    return len(ACTIVE_SESSIONS)


class ModelManager:
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.model_id = None
        self.stop_generation = False
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.thinking = None  # (start, end, think_label, answer_label) or None


model_manager = ModelManager()


def detect_thinking(tokenizer):
    if tokenizer is None:
        return None
    try:
        special = set()
        if hasattr(tokenizer, "additional_special_tokens"):
            special.update(tokenizer.additional_special_tokens)
        if hasattr(tokenizer, "added_tokens_decoder"):
            special.update(str(v) for v in tokenizer.added_tokens_decoder.values())
        for start, end, think_label, answer_label in THINKING_PATTERNS:
            if start in special or start in tokenizer.get_vocab():
                return (start, end, think_label, answer_label)
    except Exception:
        pass
    return None


class StopOnFlag(StoppingCriteria):
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        return model_manager.stop_generation


_search_cache = {}
_last_search_time = 0

def search_hf_models(query):
    global _last_search_time
    import time as _time
    _last_search_time = _time.time()

    if not query or len(query) < 2:
        return gr.update(choices=MODELS)
    q = query.strip().lower()
    if q in _search_cache:
        return gr.update(choices=_search_cache[q])
    try:
        api = HfApi()
        models = list(api.list_models(search=query, limit=15, sort="downloads"))
        model_ids = [m.id for m in models if m.id]
        if model_ids:
            _search_cache[q] = model_ids
            return gr.update(choices=model_ids)
        return gr.update(choices=MODELS)
    except Exception:
        return gr.update(choices=MODELS)


def get_system_stats(request: gr.Request = None):
    mem = psutil.virtual_memory()
    disk = psutil.disk_usage('/')
    return (
        f"CPU\t\t: \t{psutil.cpu_percent(interval=1)}%\n"
        f"Mem\t\t: \t{round(mem.used / (1024**3), 2)} / {round(mem.total / (1024**3), 2)} GB\n"
        f"Disk\t\t: \t{round(disk.used / (1024**3), 2)} / {round(disk.total / (1024**3), 2)} GB\n"
        f"Active\t: \t{len(ACTIVE_SESSIONS) if request is None else live_count(request)} session(s)"
    )


def load_new_model(model_id):
    model_manager.stop_generation = True
    model_manager.model = None
    model_manager.tokenizer = None
    model_manager.model_id = None
    yield f"Loading {model_id}..."
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(
            model_id, trust_remote_code=True, dtype=torch.float16
        )
        model_manager.tokenizer = tokenizer
        model_manager.model = model
        model_manager.model_id = model_id
        model_manager.thinking = detect_thinking(tokenizer)
        tag = " (thinking model)" if model_manager.thinking else ""
        yield f"Loaded **{model_id}** on {model_manager.device.upper()}{tag}"
    except Exception as e:
        yield f"Error loading model: {str(e)}"


def update_mode_ui(mode):
    show_sys = mode == "Chat"
    show_ctx = mode == "Q&A"
    show_src = mode == "Translation"
    show_tgt = mode == "Translation"

    defaults = {
        "Completion": ("Prompt", "Enter your prompt here...", "Once upon a time in a digital kingdom,"),
        "Chat": ("User Message", "Type your message...", "What is the capital of France?"),
        "Q&A": ("Question", "Enter your question...", "How does photosynthesis work?"),
        "Translation": ("Text to Translate", "Enter text to translate...", "Hello, how are you today?"),
    }
    label, placeholder, default = defaults.get(mode, defaults["Completion"])

    return (
        gr.update(visible=show_sys),
        gr.update(label=label, placeholder=placeholder, value=default),
        gr.update(visible=show_ctx),
        gr.update(visible=show_src),
        gr.update(visible=show_tgt),
    )


def format_prompt(mode, prompt, system_prompt="", context="", src_lang="", tgt_lang=""):
    if mode == "Completion":
        return prompt

    elif mode == "Chat":
        if model_manager.tokenizer and hasattr(model_manager.tokenizer, "apply_chat_template"):
            try:
                messages = []
                if system_prompt:
                    messages.append({"role": "system", "content": system_prompt})
                messages.append({"role": "user", "content": prompt})
                return model_manager.tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True,
                    enable_thinking=False
                )
            except Exception:
                pass
        parts = []
        if system_prompt:
            parts.append(f"[SYSTEM]: {system_prompt}")
        parts.append(f"[USER]: {prompt}")
        parts.append("[ASSISTANT]:")
        return "\n\n".join(parts)

    elif mode == "Q&A":
        if context and context.strip():
            return f"Context:\n{context.strip()}\n\nQuestion: {prompt.strip()}\n\nAnswer:"
        return f"Question: {prompt.strip()}\n\nAnswer:"

    elif mode == "Translation":
        src = src_lang or "English"
        tgt = tgt_lang or "Spanish"
        return f"Translate the following text from {src} to {tgt}:\n\n{prompt.strip()}\n\nTranslation:"

    return prompt


def estimate_duration(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, show_prompt=False):
    return min(max(int(max_tokens) // 20, 15), 60)


def run_inference(mode, prompt, system_prompt, context, src_lang, tgt_lang,
                  max_tokens, temperature, top_k, top_p, rep_penalty,
                  ngram_size, do_sample, gpu, stream_timeout):
    formatted = format_prompt(mode, prompt, system_prompt, context, src_lang, tgt_lang)
    show_prompt = mode == "Completion"
    if gpu:
        try:
            yield from run_inference_gpu(formatted, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, show_prompt, stream_timeout)
        except Exception as e:
            yield f"GPU error: {e}\n\nClick **Generate** to retry.", "GPU Unavailable"
    else:
        yield from run_inference_raw(formatted, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, show_prompt=show_prompt, stream_timeout=stream_timeout)


@spaces.GPU(duration=estimate_duration)
def run_inference_gpu(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, show_prompt=False, stream_timeout=120):
    if model_manager.model is not None:
        model_manager.model = model_manager.model.to("cuda")

    yield from run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, use_cuda=True, show_prompt=show_prompt, stream_timeout=stream_timeout)


def run_inference_raw(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample, use_cuda=False, show_prompt=False, stream_timeout=120):
    if model_manager.model is None or model_manager.tokenizer is None:
        yield "Please load a model first.", "Model not loaded"
        return

    model_manager.stop_generation = False

    tokenizer = model_manager.tokenizer
    model = model_manager.model
    thinking = model_manager.thinking

    # NOTE: Jangan paksa thinking dengan append start-token. Kalau
    # enable_thinking=False sudah dipakai di format_prompt (Chat mode),
    # append ini justru membatalkannya dan memaksa model berpikir.
    inputs = tokenizer([user_prompt], return_tensors="pt")

    if use_cuda:
        inputs = {k: v.to("cuda") for k, v in inputs.items()}
    else:
        model = model.to("cpu")
        inputs = {k: v.to("cpu") for k, v in inputs.items()}

    streamer = TextIteratorStreamer(tokenizer, timeout=float(stream_timeout), skip_prompt=True, skip_special_tokens=True)

    if not do_sample:
        temperature = 1.0

    generate_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=int(max_tokens),
        temperature=float(temperature),
        top_k=int(top_k),
        top_p=float(top_p),
        repetition_penalty=float(rep_penalty),
        no_repeat_ngram_size=int(ngram_size),
        do_sample=do_sample,
        pad_token_id=tokenizer.eos_token_id,
        stopping_criteria=StoppingCriteriaList([StopOnFlag()])
    )

    start_time = time.time()
    thread = Thread(target=model.generate, kwargs=generate_kwargs)
    thread.start()

    if thinking:
        base_display = ""
        generated_text = ""
    elif show_prompt:
        base_display = user_prompt
        generated_text = ""
    else:
        base_display = ""
        generated_text = ""

    token_count = 0
    try:
        for new_text in streamer:
            if model_manager.stop_generation:
                break

            generated_text += new_text
            token_count += 1
            duration = time.time() - start_time
            tps = token_count / duration if duration > 0 else 0

            display_text = generated_text

            if thinking:
                start_tok, end_tok, think_label, answer_label = thinking
                # Hanya bungkus sebagai thinking kalau output BENAR-BENAR
                # mengandung token thinking. Kalau thinking dimatikan
                # (enable_thinking=False), output adalah teks biasa (mis. JSON)
                # dan harus ditampilkan apa adanya tanpa prefix "> " / "*...*".
                if start_tok in generated_text or end_tok in generated_text:
                    clean = generated_text.replace("<s>", "").replace("</s>", "")
                    clean = clean.replace(start_tok, "").replace(end_tok, "")
                    clean = clean.replace("<|begin_of_solution|>", "").replace("<|end_of_solution|>", "")
                    if end_tok in generated_text:
                        parts = generated_text.split(end_tok, 1)
                        think_raw = parts[0].replace(start_tok, "").strip()
                        answer_raw = parts[1].replace("<|begin_of_solution|>", "").replace("<|end_of_solution|>", "").strip()
                        think_block = "\n".join("> " + line for line in think_raw.splitlines()) if think_raw else "> _thinking..._"
                        display_text = f"{think_block}\n\n**{answer_raw}**"
                    else:
                        think_raw = clean.strip()
                        think_block = "\n".join("> " + line for line in think_raw.splitlines()) if think_raw else "> _thinking..._"
                        display_text = f"{think_block}\n\n*...*"

            device_label = "CUDA" if use_cuda else "CPU"
            yield base_display + display_text, f"Speed: {tps:.2f} tokens/sec ({device_label})"
    except _queue.Empty:
        device_label = "CUDA" if use_cuda else "CPU"
        yield base_display + generated_text, f"Stream timed out after {stream_timeout}s ({device_label})"


def clean_cache():
    if os.path.exists(HF_CACHE_DIR):
        shutil.rmtree(HF_CACHE_DIR)
        os.makedirs(HF_CACHE_DIR)
        return "Cache cleaned successfully!"
    return "Cache directory not found."


with gr.Blocks(title="SLM Model Tester", css="""
#output-box {
    background: var(--background-fill-secondary);
    border: 1px solid var(--border-color-primary);
    border-radius: var(--radius-lg);
    padding: 16px;
    min-height: 200px;
    font-size: 15px;
    line-height: 1.6;
}
""") as app:

    gr.Markdown("# SLM Model Evaluation Hub")

    with gr.Row():
        with gr.Column(scale=1, min_width=300):

            with gr.Accordion("System", open=False):
                stats_output = gr.Textbox(label="Stats", show_label=False, max_lines=5)
                gr.Timer(2).tick(get_system_stats, None, stats_output)

            with gr.Group():
                model_input = gr.Dropdown(
                    choices=MODELS, label="Model", value=DEFAULT_MODEL,
                    allow_custom_value=True,
                )
                search_btn = gr.Button("Search HuggingFace", variant="secondary")
                with gr.Row():
                    load_btn = gr.Button("Load Model", variant="primary", scale=2)
                    clean_btn = gr.Button("Clear Cache", variant="stop", scale=1)

            with gr.Group():
                mode_input = gr.Dropdown(
                    choices=TASK_MODES, value="Completion", label="Mode"
                )
                use_gpu = gr.Checkbox(label="Use GPU", value=True)

            with gr.Accordion("Parameters", open=False):
                do_sample_input = gr.Checkbox(label="Sampling", value=True, info="Uncheck for greedy")
                max_tokens_input = gr.Slider(minimum=10, maximum=102400, value=256, step=1, label="Max Tokens")
                temperature_input = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature")
                top_k_input = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top-K")
                top_p_input = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-P")
                rep_penalty_input = gr.Slider(minimum=1.0, maximum=2.0, value=1.1, step=0.05, label="Rep. Penalty")
                ngram_size_input = gr.Slider(minimum=0, maximum=10, value=0, step=1, label="N-Gram Size")
                stream_timeout_input = gr.Slider(minimum=10, maximum=600, value=120, step=10, label="Stream Timeout (sec)")

        with gr.Column(scale=3):
            system_prompt_input = gr.Textbox(
                label="System Prompt", value="You are a helpful assistant.",
                lines=2, visible=False
            )

            user_prompt = gr.Textbox(
                label="Prompt",
                value="Once upon a time in a digital kingdom,",
                placeholder="Enter your prompt here...",
                lines=3
            )

            context_input = gr.Textbox(
                label="Context", placeholder="Paste reference context here...",
                lines=5, visible=False
            )

            src_lang_input = gr.Dropdown(
                choices=LANGUAGES, value="English", label="Translate from",
                visible=False
            )
            tgt_lang_input = gr.Dropdown(
                choices=LANGUAGES, value="Spanish", label="Translate to",
                visible=False
            )

            run_btn = gr.Button("Generate", variant="primary", size="lg")
            status_output = gr.Markdown("*Ready*")
            output_text = gr.Markdown(label="Output", elem_id="output-box")

    search_btn.click(
        fn=search_hf_models,
        inputs=[model_input],
        outputs=[model_input],
    )

    load_btn.click(
        fn=load_new_model,
        inputs=[model_input],
        outputs=[status_output]
    )

    mode_input.change(
        fn=update_mode_ui,
        inputs=[mode_input],
        outputs=[system_prompt_input, user_prompt, context_input, src_lang_input, tgt_lang_input]
    )

    run_btn.click(
        fn=run_inference,
        inputs=[
            mode_input,
            user_prompt,
            system_prompt_input,
            context_input,
            src_lang_input,
            tgt_lang_input,
            max_tokens_input,
            temperature_input,
            top_k_input,
            top_p_input,
            rep_penalty_input,
            ngram_size_input,
            do_sample_input,
            use_gpu,
            stream_timeout_input
        ],
        outputs=[output_text, status_output]
    )

    clean_btn.click(fn=clean_cache, outputs=[status_output])

if __name__ == "__main__":
    app.launch(theme=gr.themes.Soft())