import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
from threading import Thread
import gc
import os
import shutil
import torch
import psutil
import time
from orbit import orbit

# Define path for HF cache to clean
HF_CACHE_DIR = os.path.expanduser("~/.cache/huggingface/hub")

# List of models for autocomplete
MODELS = [
    'HuggingFaceTB/SmolLM2-135M', 'AxiomicLabs/GPT-X2-125M', 'Qwen/Qwen3-0.6B',
    'facebook/MobileLLM-R1-140M-base', 'SupraLabs/Supra-50M-Base', 'CompactAI-O/Shard-1',
    'SupraLabs/Supra-50M-Instruct', 'HuggingFaceTB/SmolLM-135M', 'facebook/opt-125m',
    'AxiomicLabs/GPT-S-5M', 'openai-community/gpt2', 'LH-Tech-AI/Spark-5M-Base-v4',
    'SupraLabs/Supra-Mini-v5-8M', 'EleutherAI/pythia-70m', 'SupraLabs/Supra-Mini-v4-2M',
    'EleutherAI/pythia-31m', 'StentorLabs/Stentor3-50M', 'StentorLabs/Stentor3-20M',
    'StentorLabs/Portimbria-150M', 'HuggingFaceTB/nanowhale-100m-base', 'EleutherAI/pythia-14m',
    'Harley-ml/Tenete-8M', 'Harley-ml/Dillion-1.2M', 'MihaiPopa-1/CinnabarLM-1.4M-Base',
    'MihaiPopa-1/CinnabarLM-4M-Base', 'MihaiPopa-1/PotentSulfurLM-500K-Base',
    'MihaiPopa-1/CinnabarLM-1.5M-Base', 'Harley-ml/Dillionv2-1.3M', 'Eclipse-Senpai/KeyLM-75M',
    'SupraLabs/Supra-Mini-v6-1M', 'AxiomicLabs/GPT-S-1.4M', 'GODELEV/Archaea-74M',
    'Sandroeth/cali-0.1B', 'veyra-ai/veyra3-5m-base', 'veyra-ai/veyra-30m-base-5b-tokens',
    'ThingAI/Quark-50m', 'ThingAI/Quark-135m', 'HuggingFaceTB/SmolLM2-135M-Instruct',
    'Aravindan/awesome-gpt-2-coder', 'Qwen/Qwen2.5-Coder-0.5B', 'SupraLabs/Supra-50M-Reasoning'
]

ACTIVE_SESSIONS = {}
SESSION_TIMEOUT = 60


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"


model_manager = ModelManager()


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


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).to(model_manager.device)

        model_manager.tokenizer = tokenizer
        model_manager.model = model
        model_manager.model_id = model_id

        yield f"Successfully loaded {model_id} on {model_manager.device.upper()}"

    except Exception as e:
        yield f"Error loading model: {str(e)}"


def run_inference(user_prompt, max_tokens, temperature, top_k, top_p, rep_penalty, ngram_size, do_sample):
    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
    model_id = model_manager.model_id

    # ORBIT GOVERNOR INTERCEPT
    original_prompt = user_prompt
    decision = orbit.as_dict(user_prompt)

    if not decision["allow"]:
        yield decision["message"], "Orbit blocked prompt"
        return

    user_prompt = decision["prompt"]

    is_supra_reasoning = "Supra-50M-Reasoning" in model_id if model_id else False

    if is_supra_reasoning:
        SYSTEM_PROMPT = (
            "Your role as an assistant involves thoroughly exploring questions through "
            "a systematic long thinking process before providing the final precise and accurate solutions."
        )
        prompt_to_encode = (
            f"[SYSTEM]: {SYSTEM_PROMPT}\n\n"
            f"[USER]: {user_prompt}\n\n"
            f"[ASSISTANT]: <|begin_of_thought|>\n"
        )
        skip_special = False
    else:
        prompt_to_encode = user_prompt
        skip_special = True

    inputs = tokenizer([prompt_to_encode], return_tensors="pt").to(model_manager.device)

    streamer = TextIteratorStreamer(
        tokenizer,
        timeout=60.0,
        skip_prompt=True,
        skip_special_tokens=skip_special
    )

    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 is_supra_reasoning:
        base_display = f"Prompt: {original_prompt}\n\n----------------------------------------\n\n"
        generated_text = ""
    else:
        base_display = ""
        generated_text = ""

    token_count = 0

    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 is_supra_reasoning:
            display_text = display_text.replace("<s>", "").replace("</s>", "")

            if not display_text.startswith("🧠 Thinking Process:"):
                display_text = "🧠 Thinking Process:\n" + display_text

            display_text = display_text.replace("<|begin_of_thought|>", "🧠 Thinking Process:\n")
            display_text = display_text.replace("<|end_of_thought|>", "\n\n")
            display_text = display_text.replace("<|begin_of_solution|>", "✅ Final Answer:\n\n")
            display_text = display_text.replace("<|end_of_solution|>", "")

        yield base_display + display_text, f"Speed: {tps:.2f} tokens/sec"


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="Small MF Model Tester + Orbit Governor", theme=gr.themes.Soft()) as app:

    gr.Markdown("# 🚀 Small Model Evaluation Hub with Orbit Governor")

    with gr.Row():
        with gr.Column(scale=1):

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

            with gr.Group():
                gr.Markdown("### Select or paste custom model ID here")
                with gr.Row():
                    model_id_input = gr.Dropdown(
                        choices=MODELS,
                        label="Model",
                        allow_custom_value=True,
                        show_label=False,
                        scale=3
                    )
                    load_btn = gr.Button("Load", variant="secondary", scale=1)

                clean_btn = gr.Button("Clean HF Cache", variant="stop", size="sm")

            with gr.Accordion("Generation Configuration", open=False):
                do_sample_input = gr.Checkbox(
                    label="Enable Sampling (do_sample)",
                    value=True,
                    info="Uncheck for greedy decoding"
                )
                max_tokens_input = gr.Slider(
                    minimum=10,
                    maximum=2048,
                    value=256,
                    step=1,
                    label="Max Output Tokens"
                )
                temperature_input = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=0.7,
                    step=0.1,
                    label="Temperature",
                    info="Higher = more creative"
                )
                top_k_input = gr.Slider(
                    minimum=0,
                    maximum=100,
                    value=50,
                    step=1,
                    label="Top-K",
                    info="0 = disabled"
                )
                top_p_input = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.9,
                    step=0.05,
                    label="Top-P (Nucleus)",
                    info="1.0 = disabled"
                )
                rep_penalty_input = gr.Slider(
                    minimum=1.0,
                    maximum=2.0,
                    value=1.1,
                    step=0.05,
                    label="Repetition Penalty",
                    info="1.0 = disabled"
                )
                ngram_size_input = gr.Slider(
                    minimum=0,
                    maximum=10,
                    value=0,
                    step=1,
                    label="No Repeat N-Gram Size",
                    info="0 = disabled"
                )

        with gr.Column(scale=2):
            user_prompt = gr.Textbox(
                label="Prompt",
                value="Once upon a time in a digital kingdom,",
                placeholder="Enter your prompt here...",
                lines=5
            )

            run_btn = gr.Button("Generate text", variant="primary", size="lg")
            status_output = gr.Markdown("Status: *Waiting to load model...*")
            output_text = gr.Textbox(
                label="Result",
                lines=15,
                buttons=["copy"],
                autoscroll=True
            )

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

    run_btn.click(
        fn=run_inference,
        inputs=[
            user_prompt,
            max_tokens_input,
            temperature_input,
            top_k_input,
            top_p_input,
            rep_penalty_input,
            ngram_size_input,
            do_sample_input
        ],
        outputs=[output_text, status_output]
    )

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


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