import os
import gc
import torch
import gradio as gr

from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "devoppro/FastLLM"

# ---------------------------------------------------------
# Environment
# ---------------------------------------------------------

os.environ.setdefault("HF_HOME", "/tmp/huggingface")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

print("=" * 60)
print("FastLLM")
print("=" * 60)
print(f"Model:  {MODEL_ID}")
print(f"Device: {DEVICE}")

# ---------------------------------------------------------
# Load tokenizer
# ---------------------------------------------------------

print("Loading tokenizer...")

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print("Tokenizer loaded.")

# ---------------------------------------------------------
# Load model
# ---------------------------------------------------------

print("Loading FastLLM...")

model_kwargs = {
    "trust_remote_code": True,
    "low_cpu_mem_usage": True,
}

if DEVICE == "cuda":
    model_kwargs["torch_dtype"] = torch.float16
else:
    model_kwargs["torch_dtype"] = torch.float32

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    **model_kwargs,
)

model.eval()
model.to(DEVICE)

print("FastLLM loaded successfully.")
print("=" * 60)


# ---------------------------------------------------------
# Generation
# ---------------------------------------------------------

def generate_response(
    message,
    history,
    temperature,
    top_p,
    max_new_tokens,
    repetition_penalty,
):
    if not message or not message.strip():
        return history, ""

    try:
        # Convert Gradio history to a simple conversation
        prompt_parts = []

        for item in history:
            if isinstance(item, dict):
                role = item.get("role")
                content = item.get("content", "")

                if role == "user":
                    prompt_parts.append(f"User: {content}")

                elif role == "assistant":
                    prompt_parts.append(f"Assistant: {content}")

        prompt_parts.append(f"User: {message}")
        prompt_parts.append("Assistant:")

        prompt = "\n".join(prompt_parts)

        inputs = tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=2048,
        )

        input_ids = inputs["input_ids"].to(DEVICE)
        attention_mask = inputs["attention_mask"].to(DEVICE)

        with torch.inference_mode():
            output = model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_new_tokens=int(max_new_tokens),
                temperature=float(temperature),
                top_p=float(top_p),
                repetition_penalty=float(repetition_penalty),
                do_sample=True,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )

        generated_tokens = output[0][input_ids.shape[-1]:]

        response = tokenizer.decode(
            generated_tokens,
            skip_special_tokens=True,
        ).strip()

        if not response:
            response = "FastLLM did not generate a response."

        history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": response},
        ]

        # Release temporary tensors
        del inputs
        del input_ids
        del attention_mask
        del output

        if DEVICE == "cuda":
            torch.cuda.empty_cache()

        gc.collect()

        return history, ""

    except Exception as e:
        print("Generation error:")
        print(repr(e))

        error_message = f"❌ Generation error:\n\n`{str(e)}`"

        history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": error_message},
        ]

        return history, ""


# ---------------------------------------------------------
# Clear conversation
# ---------------------------------------------------------

def clear_chat():
    return []


# ---------------------------------------------------------
# UI
# ---------------------------------------------------------

css = """
.gradio-container {
    max-width: 1100px !important;
    margin: auto !important;
}

.title {
    text-align: center;
    margin-bottom: 4px;
}

.subtitle {
    text-align: center;
    opacity: 0.7;
    margin-bottom: 20px;
}

.status {
    text-align: center;
    font-size: 13px;
    opacity: 0.65;
}
"""

with gr.Blocks(
    css=css,
    title="FastLLM",
) as demo:

    gr.Markdown(
        """
        # ⚡ FastLLM
        """,
        elem_classes=["title"],
    )

    gr.Markdown(
        """
        A 150M parameter causal language model built from scratch by **devoppro**.
        """,
        elem_classes=["subtitle"],
    )

    chatbot = gr.Chatbot(
        label="FastLLM",
        height=560,
        type="messages",
        bubble_full_width=False,
    )

    with gr.Row():

        message = gr.Textbox(
            placeholder="Message FastLLM...",
            label="",
            scale=5,
            lines=2,
        )

        send = gr.Button(
            "Send",
            variant="primary",
            scale=1,
        )

    with gr.Row():

        temperature = gr.Slider(
            minimum=0.1,
            maximum=2.0,
            value=0.7,
            step=0.05,
            label="Temperature",
        )

        top_p = gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.9,
            step=0.05,
            label="Top P",
        )

        max_tokens = gr.Slider(
            minimum=16,
            maximum=1024,
            value=256,
            step=16,
            label="Max New Tokens",
        )

        repetition_penalty = gr.Slider(
            minimum=1.0,
            maximum=2.0,
            value=1.05,
            step=0.01,
            label="Repetition Penalty",
        )

    with gr.Row():

        clear = gr.Button(
            "🗑️ Clear conversation"
        )

    gr.Markdown(
        f"""
        <div class="status">
        Model: <b>{MODEL_ID}</b> · Device: <b>{DEVICE.upper()}</b>
        </div>
        """,
        elem_classes=["status"],
    )

    # -----------------------------------------------------
    # Events
    # -----------------------------------------------------

    send.click(
        generate_response,
        inputs=[
            message,
            chatbot,
            temperature,
            top_p,
            max_tokens,
            repetition_penalty,
        ],
        outputs=[
            chatbot,
            message,
        ],
    )

    message.submit(
        generate_response,
        inputs=[
            message,
            chatbot,
            temperature,
            top_p,
            max_tokens,
            repetition_penalty,
        ],
        outputs=[
            chatbot,
            message,
        ],
    )

    clear.click(
        clear_chat,
        inputs=[],
        outputs=[chatbot],
    )


# ---------------------------------------------------------
# Launch
# ---------------------------------------------------------

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
    )