import torch
from transformers import pipeline
import gradio as gr
from datetime import datetime

class BrightLlamaChatbot:
    def __init__(self):
        self.model_id = "ezcz/Llama-3.2-3B-BrightLlamaChat-LoRA"
        self.pipe = pipeline(
            "text-generation",
            model=self.model_id,
            torch_dtype=torch.bfloat16,
            device_map="auto",
        )
        self.system_prompt = """You are a helpful AI assistant focused on coding and reasoning tasks. 
        You provide clear, accurate responses while maintaining a friendly tone."""
        
    def format_message(self, role, content):
        timestamp = datetime.now().strftime("%H:%M:%S")
        return f"<div class='message {role}'><span class='time'>[{timestamp}]</span>{content}</div>"
    
    def generate_response(self, user_input, chat_history):
        if not user_input:
            return chat_history
            
        # Format conversation history for the model
        messages = [{"role": "system", "content": self.system_prompt}]
        for msg in chat_history.split("\n"):
            if msg.strip():
                if "User:" in msg:
                    messages.append({"role": "user", "content": msg.replace("User:", "").strip()})
                elif "Assistant:" in msg:
                    messages.append({"role": "assistant", "content": msg.replace("Assistant:", "").strip()})
        
        messages.append({"role": "user", "content": user_input})
        
        # Generate response
        response = self.pipe(messages, max_new_tokens=512, return_full_text=False, temperature=0.7, top_p=0.9)
        assistant_response = response[0]["generated_text"]
        
        # Format and update chat history
        updated_history = chat_history + "\n" + f"User: {user_input}" + "\n" + f"Assistant: {assistant_response}"
        
        return updated_history

    def create_interface(self):
        with gr.Blocks(css=self.get_custom_css()) as interface:
            gr.HTML("<h1>🦙 Bright Llama Chatbot</h1>")
            gr.HTML("<p>An AI assistant using ezcz/Llama-3.2-3B-BrightLlamaChat-LoRA</p>")
            
            with gr.Row():
                with gr.Column(scale=4):
                    chatbot = gr.Textbox(
                        show_label=False,
                        placeholder="Conversation history...",
                        lines=15,
                        max_lines=15,
                        interactive=False
                    )
                    
                    with gr.Row():
                        with gr.Column(scale=8):
                            user_input = gr.Textbox(
                                show_label=False,
                                placeholder="Type your message here...",
                                lines=2
                            )
                        with gr.Column(scale=1):
                            submit_btn = gr.Button("Send", variant="primary")
                            clear_btn = gr.Button("Clear")
                
                with gr.Column(scale=1):
                    gr.HTML("<h3>Bright Llama LoRA Chatbot:</h3>")
                    gr.HTML("""
                        <ul>
                            <li>💬 Natural conversation</li>
                            <li>🐢Running on CPU - expect slow responses</li>
                        </ul>
                    """)
                    
                    with gr.Accordion("Settings", open=False):
                        temperature = gr.Slider(
                            minimum=0.1,
                            maximum=1.0,
                            value=0.7,
                            step=0.1,
                            label="Temperature"
                        )
                        
                        max_length = gr.Slider(
                            minimum=64,
                            maximum=1024,
                            value=512,
                            step=64,
                            label="Max Length"
                        )
            
            submit_btn.click(
                fn=self.generate_response,
                inputs=[user_input, chatbot],
                outputs=chatbot
            ).then(
                fn=lambda: "",
                outputs=user_input
            )
            
            clear_btn.click(
                fn=lambda: ("", ""),
                outputs=[user_input, chatbot]
            )
            
        return interface
    
    def get_custom_css(self):
        return """
            .message {
                padding: 10px;
                margin: 5px;
                border-radius: 8px;
            }
            .user {
                background-color: #f0f0f0;
                margin-left: 20px;
            }
            .assistant {
                background-color: #e3f2fd;
                margin-right: 20px;
            }
            .time {
                color: #666;
                font-size: 0.8em;
                margin-right: 10px;
            }
            pre {
                background-color: #2b2b2b;
                color: #ffffff;
                padding: 10px;
                border-radius: 5px;
                overflow-x: auto;
            }
            code {
                font-family: 'Courier New', monospace;
            }
        """

if __name__ == "__main__":
    chatbot = BrightLlamaChatbot()
    interface = chatbot.create_interface()
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )