import spaces
import gradio as gr
from transformers import pipeline, AutoTokenizer, TextIteratorStreamer
from huggingface_hub import snapshot_download
import torch
import gc
from threading import Thread
import os

model_choices = ["Locutusque/Esmeralda-Qwen3.6-35B-A3B"]

CURRENT = {"name": None, "pipe": None}

def predownload_all_models():
    for model_name in model_choices:
        print(f"Downloading {model_name}...")
        snapshot_download(repo_id=model_name, token=os.environ["token"])
    print("All models downloaded to cache.")

def _unload_current():
    if CURRENT["pipe"] is not None:
        print(f"Unloading {CURRENT['name']}...")
        del CURRENT["pipe"]
        CURRENT["pipe"] = None
        CURRENT["name"] = None
        gc.collect()

def get_model(model_name):
    if CURRENT["name"] != model_name:
        _unload_current()
        task = "image-text-to-text" if model_name == "Locutusque/Esmeralda-31B" else "text-generation"
        from transformers import BitsAndBytesConfig
        quant = BitsAndBytesConfig(
            load_in_8bit=True,
        )
        pipe = pipeline(
            task,
            model=model_name,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            token=os.environ["token"],
            use_fast=True,
            device_map="auto",
            model_kwargs={"quantization_config": quant},
        )
        tokenizer = AutoTokenizer.from_pretrained(
            model_name, trust_remote_code=True, token=os.environ["token"]
        )
        tokenizer.eos_token = "<|im_end|>"
        pipe.tokenizer = tokenizer
        CURRENT["name"] = model_name
        CURRENT["pipe"] = pipe
    return CURRENT["pipe"]

@spaces.GPU(duration=45)
def generate(
    message,
    history,
    model_name,
    system,
    temperature=0.4,
    top_p=0.95,
    min_p=0.1,
    top_k=50,
    max_new_tokens=256,
):
    outputs = []
    try:
        pipe = get_model(model_name)
        prompt = f"<|im_start|>system\n{system}<|im_end|>\n"
        for (user_turn, assistant_turn) in history:
            prompt += f"<|im_start|>user\n{user_turn}<|im_end|>\n<|im_start|>assistant\n{assistant_turn}<|im_end|>\n"
        prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
        streamer = TextIteratorStreamer(pipe.tokenizer, timeout=240.0, skip_prompt=True, skip_special_tokens=True)
        generation_kwargs = dict(
            text_inputs=prompt,
            streamer=streamer,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            top_p=top_p,
            min_p=min_p,
            top_k=top_k,
            temperature=temperature,
            num_beams=1,
            repetition_penalty=1.1,
        )
        t = Thread(target=pipe.__call__, kwargs=generation_kwargs)
        t.start()
        for chunk in streamer:
            outputs.append(chunk)
            yield "".join(outputs)
    except StopAsyncIteration:
        print("Stream stopped unexpectedly.")
        yield "".join(outputs)
    except Exception as e:
        print(f"An error occurred: {e}")
        yield "An error occurred during generation."

g = gr.ChatInterface(
    fn=generate,
    additional_inputs=[
        gr.components.Dropdown(choices=model_choices, label="Model", value=model_choices[0], interactive=True),
        gr.components.Textbox(lines=2, label="System Prompt", value="You are a sentient, super intelligent AI developed by a man named Locutusque."),
        gr.components.Slider(minimum=0, maximum=2, value=0.8, label="Temperature"),
        gr.components.Slider(minimum=0, maximum=1, value=0.95, label="Top p"),
        gr.components.Slider(minimum=0, maximum=1, value=0.1, label="Min P"),
        gr.components.Slider(minimum=0, maximum=100, step=1, value=15, label="Top k"),
        gr.components.Slider(minimum=1, maximum=8192, step=1, value=1024, label="Max tokens"),
    ],
    title="Locutusque's Language Models",
    description="Try out Locutusque's language models here! Credit goes to Mediocreatmybest for this space. You may also find some experimental preview models that have not been made public here.",
)

if __name__ == "__main__":
    predownload_all_models()
    g.launch()