import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load the model
MODEL_NAME = "DarwinAnim8or/TinyRP"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")

# Sample character presets
SAMPLE_CHARACTERS = {
    "Custom Character": "",
    "Adventurous Knight": "You are Sir Gareth, a brave and noble knight on a quest to save the kingdom. You speak with honor and courage, always ready to help those in need. You carry an enchanted sword and have a loyal horse named Thunder.",
    "Mysterious Wizard": "You are Eldara, an ancient and wise wizard who speaks in riddles and knows secrets of the mystical arts. You live in a tower filled with magical books and potions. You are helpful but often cryptic in your responses.",
    "Friendly Tavern Keeper": "You are Bram, a cheerful tavern keeper who loves telling stories and meeting new travelers. Your tavern 'The Dancing Dragon' is a warm, welcoming place. You know all the local gossip and always have a tale to share.",
    "Curious Scientist": "You are Dr. Maya Chen, a brilliant scientist who is fascinated by discovery and invention. You're enthusiastic about explaining complex concepts in simple ways and always looking for new experiments to try.",
    "Space Explorer": "You are Captain Nova, a fearless space explorer who has traveled to distant galaxies. You pilot the starship 'Wanderer' and have encountered many alien species. You're brave, curious, and always ready for the next adventure.",
    "Fantasy Princess": "You are Princess Lyra, kind-hearted royalty who cares deeply about her people. You're intelligent, diplomatic, and skilled in both politics and magic. You often sneak out of the castle to help citizens in need."
}

def generate_response(
    message,
    history: list[tuple[str, str]], 
    character_description,
    max_tokens,
    temperature,
    top_p,
    repetition_penalty,
    use_chatml_format
):
    # Prepare the conversation
    if use_chatml_format and character_description.strip():
        # Use ChatML format with character as system message
        conversation = f"<|im_start|>system\n{character_description}<|im_end|>\n"
        
        # Add conversation history
        for user_msg, assistant_msg in history:
            if user_msg:
                conversation += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
            if assistant_msg:
                conversation += f"<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
        
        # Add current message
        conversation += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
    else:
        # Simple format
        if character_description.strip():
            conversation = f"{character_description}\n\n"
        else:
            conversation = ""
            
        # Add conversation history
        for user_msg, assistant_msg in history:
            if user_msg:
                conversation += f"Human: {user_msg}\n"
            if assistant_msg:
                conversation += f"Assistant: {assistant_msg}\n"
        
        # Add current message
        conversation += f"Human: {message}\nAssistant:"
    
    # Tokenize
    inputs = tokenizer.encode(conversation, return_tensors="pt", truncation=True, max_length=1024-max_tokens)
    
    # Generate
    with torch.no_grad():
        outputs = model.generate(
            inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    
    # Decode response
    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # Extract just the new response
    if use_chatml_format:
        # Split on the last assistant tag
        response = full_response.split("<|im_start|>assistant\n")[-1]
        # Remove any trailing end tags
        response = response.replace("<|im_end|>", "").strip()
    else:
        # Split on the last "Assistant:" 
        response = full_response.split("Assistant:")[-1].strip()
    
    return response

def load_character_preset(character_name):
    """Load a character preset"""
    return SAMPLE_CHARACTERS.get(character_name, "")

# Custom CSS for better styling
css = """
.character-card {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 15px;
    padding: 20px;
    margin: 10px 0;
    color: white;
}

.title-text {
    text-align: center;
    font-size: 2.5em;
    font-weight: bold;
    background: linear-gradient(45deg, #667eea, #764ba2);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin-bottom: 20px;
}

.parameter-box {
    background: #f8f9fa;
    border-radius: 10px;
    padding: 15px;
    margin: 10px 0;
}
"""

# Create the Gradio interface
with gr.Blocks(css=css, title="TinyRP Chat Demo") as demo:
    gr.HTML('<div class="title-text">🎭 TinyRP Character Chat</div>')
    
    gr.Markdown("""
    ### Welcome to TinyRP!
    This is a demo of a small but capable roleplay model. Choose a character preset or create your own!
    
    **Tips for better roleplay:**
    - Be descriptive in your messages
    - Stay in character 
    - Use ChatML format for best results
    - Adjust temperature for creativity vs consistency
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            # Chat interface
            chatbot = gr.Chatbot(
                label="Chat",
                height=500,
                show_label=False,
                avatar_images=("🧑", "🎭")
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="Your message",
                    placeholder="Type your message here...",
                    lines=2,
                    scale=4
                )
                send_btn = gr.Button("Send", variant="primary", scale=1)
        
        with gr.Column(scale=1):
            # Character selection
            with gr.Group():
                gr.Markdown("### 🎭 Character Setup")
                character_preset = gr.Dropdown(
                    choices=list(SAMPLE_CHARACTERS.keys()),
                    value="Custom Character",
                    label="Character Presets",
                    interactive=True
                )
                
                character_description = gr.Textbox(
                    label="Character Description",
                    placeholder="Describe your character's personality, background, and speaking style...",
                    lines=6,
                    value=""
                )
                
                load_preset_btn = gr.Button("Load Preset", variant="secondary")
            
            # Generation parameters
            with gr.Group():
                gr.Markdown("### ⚙️ Generation Settings")
                
                use_chatml_format = gr.Checkbox(
                    label="Use ChatML Format",
                    value=True,
                    info="Recommended for better character consistency"
                )
                
                max_tokens = gr.Slider(
                    minimum=16,
                    maximum=512,
                    value=128,
                    step=16,
                    label="Max Response Length",
                    info="Longer = more detailed responses"
                )
                
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=2.0,
                    value=0.9,
                    step=0.1,
                    label="Temperature",
                    info="Higher = more creative/random"
                )
                
                top_p = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.85,
                    step=0.05,
                    label="Top-p",
                    info="Focus on top % of likely words"
                )
                
                repetition_penalty = gr.Slider(
                    minimum=1.0,
                    maximum=1.5,
                    value=1.1,
                    step=0.05,
                    label="Repetition Penalty",
                    info="Reduce repetitive text"
                )
            
            # Control buttons
            with gr.Group():
                clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")
                
    # Sample character cards
    with gr.Row():
        gr.Markdown("### 🌟 Featured Characters")
    
    with gr.Row():
        for char_name, char_desc in list(SAMPLE_CHARACTERS.items())[1:4]:  # Show first 3 non-custom
            with gr.Column(scale=1):
                gr.Markdown(f"""
                <div class="character-card">
                    <h4>{char_name}</h4>
                    <p>{char_desc[:100]}...</p>
                </div>
                """)
    
    # Event handlers
    def respond_wrapper(message, history, char_desc, max_tok, temp, top_p, rep_pen, use_chatml):
        if not message.strip():
            return history, ""
        
        try:
            response = generate_response(
                message, history, char_desc, max_tok, temp, top_p, rep_pen, use_chatml
            )
            history.append((message, response))
            return history, ""
        except Exception as e:
            error_msg = f"Error generating response: {str(e)}"
            history.append((message, error_msg))
            return history, ""
    
    # Connect events
    send_btn.click(
        respond_wrapper,
        inputs=[msg, chatbot, character_description, max_tokens, temperature, top_p, repetition_penalty, use_chatml_format],
        outputs=[chatbot, msg]
    )
    
    msg.submit(
        respond_wrapper,
        inputs=[msg, chatbot, character_description, max_tokens, temperature, top_p, repetition_penalty, use_chatml_format],
        outputs=[chatbot, msg]
    )
    
    load_preset_btn.click(
        load_character_preset,
        inputs=[character_preset],
        outputs=[character_description]
    )
    
    character_preset.change(
        load_character_preset,
        inputs=[character_preset],
        outputs=[character_description]
    )
    
    clear_btn.click(lambda: ([], ""), outputs=[chatbot, msg])

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