import gradio as gr
import torch
import re

from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TextIteratorStreamer
)

from threading import Thread

# ============================================
# CONFIG
# ============================================

MODEL_ID = "huntwter/Obsidian-1"

MAX_NEW_TOKENS = 512
TEMPERATURE = 0.7
TOP_P = 0.95

SYSTEM_PROMPT = (
    "You are Obsidian, a reasoning-focused "
    "cybersecurity engineering AI created by Jatin Sharma."
)

# ============================================
# LOAD MODEL
# ============================================

print("Loading tokenizer...")

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

print("Loading model...")

model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)

model.eval()

print("Model loaded successfully.")

# ============================================
# CSS
# ============================================

custom_css = """
body,
.gradio-container {
    background: #f0f4f9 !important;
    font-family: Arial, sans-serif !important;
}

.chat-container {
    display: flex;
    flex-direction: column;
    gap: 12px;
    padding: 18px;
    height: calc(100vh - 220px);
    overflow-y: auto;
    background: white;
    border-radius: 16px;
}

.chat-row {
    display: flex;
    width: 100%;
}

.chat-row.user {
    justify-content: flex-end;
}

.chat-row.assistant {
    justify-content: flex-start;
}

.bubble {
    max-width: 72%;
    padding: 12px 16px;
    border-radius: 16px;
    line-height: 1.5;
    font-size: 15px;
}

.user-bubble {
    background: #2563eb;
    color: white;
}

.assistant-bubble {
    background: white;
    color: black;
    border: 1px solid #ddd;
}

.think-block {
    background: #f8fafc;
    border-left: 3px solid #2563eb;
    padding: 12px;
    margin: 10px 0;
}
"""

# ============================================
# THINK FORMAT
# ============================================

def format_think_blocks(text):

    pattern = r'<think>(.*?)</think>'

    def replacer(match):

        content = match.group(1)

        return f"""
        <div class="think-block">
            <b>Thinking...</b><br>
            {content}
        </div>
        """

    return re.sub(
        pattern,
        replacer,
        text,
        flags=re.DOTALL
    )

# ============================================
# CHAT HTML
# ============================================

def render_chat_html(history):

    rows = []

    for msg in history:

        role = msg["role"]
        content = msg["content"]

        if role == "user":

            rows.append(
                f"""
                <div class="chat-row user">
                    <div class="bubble user-bubble">
                        {content}
                    </div>
                </div>
                """
            )

        else:

            rows.append(
                f"""
                <div class="chat-row assistant">
                    <div class="bubble assistant-bubble">
                        {content}
                    </div>
                </div>
                """
            )

    return f"""
    <div class="chat-container">
        {"".join(rows)}
    </div>
    """

# ============================================
# PROMPT BUILDER
# ============================================

def build_prompt(message, history):

    prompt = f"System: {SYSTEM_PROMPT}\n\n"

    for msg in history:

        role = msg["role"]
        content = msg["content"]

        if role == "user":

            prompt += f"User: {content}\n"

        else:

            prompt += f"Assistant: {content}\n"

    prompt += f"User: {message}\n"
    prompt += "Assistant: "

    return prompt

# ============================================
# GENERATE
# ============================================

def predict(message, history):

    try:

        prompt = build_prompt(message, history)

        inputs = tokenizer(
            prompt,
            return_tensors="pt"
        ).to(model.device)

        streamer = TextIteratorStreamer(
            tokenizer,
            skip_prompt=True,
            skip_special_tokens=True
        )

        generation_kwargs = dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=True,
            temperature=TEMPERATURE,
            top_p=TOP_P,
            repetition_penalty=1.1
        )

        thread = Thread(
            target=model.generate,
            kwargs=generation_kwargs
        )

        thread.start()

        partial_text = ""

        for token in streamer:

            partial_text += token

            formatted = format_think_blocks(partial_text)

            yield formatted

    except Exception as e:

        yield f"""
        <div style="color:red;padding:20px;">
            Error:<br><br>
            {str(e)}
        </div>
        """

# ============================================
# UI
# ============================================

with gr.Blocks(
    title="Obsidian-v1",
    css=custom_css
) as app:

    gr.HTML("""
    <h1 style="text-align:center;">
        Obsidian-v1
    </h1>
    """)

    chat_html = gr.HTML(
        value=render_chat_html([]),
        sanitize_html=False
    )

    with gr.Row():

        msg = gr.Textbox(
            placeholder="Type a message...",
            show_label=False,
            scale=8
        )

        send_btn = gr.Button(
            "Send",
            scale=1
        )

    history_state = gr.State([])

    # ============================================
    # USER ACTION
    # ============================================

    def user_action(user_message, history):

        if not user_message.strip():

            return "", render_chat_html(history), history

        history.append({
            "role": "user",
            "content": user_message
        })

        history.append({
            "role": "assistant",
            "content": ""
        })

        return (
            "",
            render_chat_html(history),
            history
        )

    # ============================================
    # BOT ACTION
    # ============================================

    def bot_action(history):

        user_message = history[-2]["content"]

        conversation = history[:-1]

        for response in predict(
            user_message,
            conversation
        ):

            history[-1]["content"] = response

            yield (
                render_chat_html(history),
                history
            )

    # ============================================
    # EVENTS
    # ============================================

    msg.submit(
        user_action,
        [msg, history_state],
        [msg, chat_html, history_state],
        queue=False
    ).then(
        bot_action,
        [history_state],
        [chat_html, history_state]
    )

    send_btn.click(
        user_action,
        [msg, history_state],
        [msg, chat_html, history_state],
        queue=False
    ).then(
        bot_action,
        [history_state],
        [chat_html, history_state]
    )

# ============================================
# LAUNCH
# ============================================

if __name__ == "__main__":

    app.launch()