import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import re

# 1. Load weights explicitly
MODEL_ID = "Axiom-AI/Quasar-1-0.8B"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True
)

# Helper function to parse both <think> and <answer> tag structures
def format_reasoning(text):
    """
    Parses <think>...</think> and <answer>...</answer> tags 
    to apply highly distinct, clean visual wrappers.
    """
    # Extract Thinking Content
    thinking_html = ""
    if "<think>" in text:
        if "</think>" in text:
            thinking_content = text.split("<think>")[1].split("</think>")[0].strip()
            thinking_html = f"""<details open><summary style="color: #ff7a00; font-weight: bold; cursor: pointer; margin-bottom: 4px;">🤔 Thinking Process (Click to collapse)</summary>
<div style="color: #555555; font-style: italic; background-color: #fdfdfd; padding: 12px; border-left: 4px solid #ff7a00; border-radius: 4px; margin: 4px 0 16px 0; font-size: 0.95em; line-height: 1.5;">
{thinking_content}
</div></details>"""
        else:
            # Currently streaming thought process
            thinking_content = text.split("<think>")[1].strip()
            thinking_html = f"""<details open><summary style="color: #ff7a00; font-weight: bold;">🤔 Thinking...</summary>
<div style="color: #555555; font-style: italic; background-color: #fdfdfd; padding: 12px; border-left: 4px solid #ff7a00; border-radius: 4px; margin: 4px 0 16px 0; font-size: 0.95em; line-height: 1.5;">
{thinking_content}
</div></details>"""

    # Extract Answer Content
    answer_html = ""
    if "<answer>" in text:
        if "</answer>" in text:
            answer_content = text.split("<answer>")[1].split("</answer>")[0].strip()
            answer_html = f"""<div style="background-color: #f4f6f8; border-left: 4px solid #2e4053; padding: 16px; border-radius: 4px; color: #1c2833; box-shadow: inset 0 1px 3px rgba(0,0,0,0.05); line-height: 1.6;">
{answer_content}
</div>"""
        else:
            # Currently streaming final answer output
            answer_content = text.split("<answer>")[1].strip()
            answer_html = f"""<div style="background-color: #f4f6f8; border-left: 4px solid #2e4053; padding: 16px; border-radius: 4px; color: #1c2833; line-height: 1.6;">
<span style="color: #2e4053; font-weight: bold; display: block; margin-bottom: 8px;">📝 Writing Response...</span>
{answer_content}
</div>"""
            
    # Combine them for the streaming UI wrapper display
    if thinking_html or answer_html:
        return f"{thinking_html}\n{answer_html}"
        
    return text

# 2. Main Generation Logic Handler
def generate_response(message, history, system_prompt, temperature, top_p, top_k, max_tokens):
    # Setup standard base template structures
    prompt = f"System: {system_prompt}\n" if system_prompt else ""
    
    # Pack up chat context from the dictionary format
    for turn in history:
        role = turn.get("role")
        content = turn.get("content")
        if role == "user":
            prompt += f"User: {content}\n"
        elif role == "assistant":
            # Strip out our custom HTML formatting back to raw tags for context history
            clean_content = content
            if "🤔" in content or "background-color" in content:
                # RegEx conversion logic to revert beautiful layout code back to raw tags for model's context awareness
                import re
                # Check for completed or uncompleted think elements
                think_match = re.search(r'<div style="color: #555555;.*?>(.*?)</div>', clean_content, re.DOTALL)
                ans_match = re.search(r'</div></details>\s*<div style="background-color: #f4f6f8;.*?>.*?</span>(.*?)</div>', clean_content, re.DOTALL) or re.search(r'</div></details>\s*<div style="background-color: #f4f6f8;.*?>(.*?)</div>', clean_content, re.DOTALL)
                
                raw_think = f"<think>{think_match.group(1).strip()}</think>" if think_match else ""
                raw_ans = f"<answer>{ans_match.group(1).strip()}</answer>" if ans_match else ""
                clean_content = f"{raw_think}\n{raw_ans}".strip()
                
            prompt += f"Assistant: {clean_content}\n"
    
    # Inject active raw message with Quasar's default target sequence patterns
    prompt += f"User: {message}\nAssistant: <think>"
    
    inputs = tokenizer(prompt, return_tensors="pt")
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
    
    generation_kwargs = dict(
        inputs,
        streamer=streamer,
        max_new_tokens=int(max_tokens),
        do_sample=True if temperature > 0.0 else False,
        temperature=float(temperature) if temperature > 0.0 else 1.0,
        top_p=float(top_p),
        top_k=int(top_k)
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    partial_text = "<think>"
    for new_text in streamer:
        partial_text += new_text
        yield format_reasoning(partial_text)

# 3. Clean Interface Assembly 
with gr.Blocks() as demo:
    gr.Markdown("# Quasar-1-0.8B Control Engine")
    
    with gr.Row():
        # Left Workspace Column: Chat Layout
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(height=500)
            msg_input = gr.Textbox(placeholder="Enter reasoning prompt...", label="Your Message")
            with gr.Row():
                submit_btn = gr.Button("Send", variant="primary")
                clear_btn = gr.Button("Clear Chat")

        # Right Workspace Column: Customizable Model Generation Properties
        with gr.Column(scale=1):
            gr.Markdown("### Inference Parameters")
            system_prompt = gr.Textbox(
                value="You are a helpful assistant that reasons step-by-step.", 
                label="System Prompt", 
                lines=3
            )
            temperature = gr.Slider(minimum=0.0, maximum=1.5, value=0.4, step=0.05, label="Temperature")
            top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, step=0.05, label="Top-P")
            top_k = gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top-K")
            max_tokens = gr.Slider(minimum=64, maximum=2048, value=1024, step=64, label="Max New Tokens")

    # Gradio dictionary append functions
    def user_side_submit(user_message, history):
        history.append({"role": "user", "content": user_message})
        return "", history

    def bot_side_inference(history, system_p, temp, t_p, t_k, max_t):
        user_message = history[-1]["content"]
        
        # Initialize assistant message element with a loading state style
        history.append({"role": "assistant", "content": "🤔 *Thinking...*"})
        
        for partial_reply in generate_response(user_message, history[:-2], system_p, temp, t_p, t_k, max_t):
            history[-1]["content"] = partial_reply
            yield history

    # Event Wiring Flow
    msg_input.submit(
        user_side_submit, [msg_input, chatbot], [msg_input, chatbot], queue=False
    ).then(
        bot_side_inference, 
        [chatbot, system_prompt, temperature, top_p, top_k, max_tokens], 
        chatbot
    )
    
    submit_btn.click(
        user_side_submit, [msg_input, chatbot], [msg_input, chatbot], queue=False
    ).then(
        bot_side_inference, 
        [chatbot, system_prompt, temperature, top_p, top_k, max_tokens], 
        chatbot
    )
    
    clear_btn.click(lambda: [], None, chatbot, queue=False)

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