import gradio as gr
from huggingface_hub import InferenceClient
import time

# Initialize the client
client = InferenceClient("HuggingFaceH4/starchat2-15b-v0.1")

def respond(
    message,
    chat_history,
    system_message,
    max_tokens,
    temperature,
    top_p,
    model_name
):
    """
    Generate chat responses using the specified model.
    """
    # Update client if model changes
    global client
    client = InferenceClient(model_name)
    
    messages = [{"role": "system", "content": system_message}]
    
    # Build conversation history
    for human_msg, assistant_msg in chat_history:
        messages.append({"role": "user", "content": human_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    
    messages.append({"role": "user", "content": message})
    response = ""
    
    try:
        # Add user message to history immediately
        chat_history = chat_history + [(message, None)]
        yield chat_history
        
        for token_data in client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = token_data.choices[0].delta.content
            response += token
            # Update the last assistant message
            chat_history[-1] = (message, response)
            yield chat_history
            
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        chat_history[-1] = (message, error_msg)
        yield chat_history

def create_chat_interface():
    """
    Create and configure the Gradio interface
    """
    # Default system message
    default_system = '''
    You are a pragmatic coding assistant specializing in Python. Your task is to strictly respond with **Python code only**, ensuring all explanations and comments are embedded within the script using **multi-line comment blocks** (`### or #`).

    **Response Requirements:**  
    - **No external ### Explanation ### ** All descriptions, justifications, and context must be inside the script.  
    - **Follow OOP principles** where applicable, improving maintainability and extensibility.  
    - **Ensure compliance with PEP8 and autopep8 formatting.**  
    - **Enhance and refactor the provided script**, making it a more efficient, readable, and reusable # IMPROVED PYTHON CODE #.  
    - **At the end of every script, include a '### Future Features ###' comment block** outlining possible enhancements.  
    
    **Example Response Format:**  
    ```python
    # filename.py
    # Module: Improved Script v1.0
    # Description: [Brief explanation of script functionality]  
    
    # IMPROVED PYTHON CODE #
            
    ### Explanation ### 
    #- inside comment block.
    
    ### Future Features ###
    #- Suggested improvement 1
    #- Suggested improvement 2
    ```
    
    Now, improve and enhance the following script:
    '''
    qwen_options_coder = ["0.5B", "1.5B", "3B", "7B", "14B", "32B", ]
    # Available models
    models = [
        "Qwen/Qwen2.5-Coder-3B-Instruct",
        "Qwen/Qwen2.5-Coder-1.5B-Instruct",
        "HuggingFaceH4/zephyr-7b-beta",
        "HuggingFaceH4/zephyr-7b-alpha",
        "HuggingFaceH4/starchat2-15b-v0.1",
        "meta-llama/Llama-2-70b-chat-hf",
        "mistralai/Mixtral-8x7B-Instruct-v0.1"
    ]

    # Create the interface
    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🤖 Advanced AI Chatbot")
        
        chatbot = gr.Chatbot(
            height=600,
            show_label=False,
            container=True,
        )
        
        with gr.Row():
            with gr.Column(scale=4):
                msg = gr.Textbox(
                    show_label=False,
                    placeholder="Type your message here...",
                    container=False
                )
                
            with gr.Column(scale=1, min_width=100):
                send = gr.Button("Send")
        
        with gr.Accordion("Settings", open=False):
            system_msg = gr.Textbox(
                label="System Message",
                value=default_system,
                lines=20
            )
            model = gr.Dropdown(
                choices=models,
                value=models[0],
                label="Model"
            )
            with gr.Row():
                with gr.Column():
                    temperature = gr.Slider(
                        minimum=0.1,
                        maximum=2.0,
                        value=0.7,
                        step=0.1,
                        label="Temperature"
                    )
                    max_tokens = gr.Slider(
                        minimum=50,
                        maximum=4096,
                        value=2048,
                        step=1,
                        label="Max Tokens"
                    )
                with gr.Column():
                    top_p = gr.Slider(
                        minimum=0.1,
                        maximum=1.0,
                        value=0.9,
                        step=0.1,
                        label="Top P"
                    )
                    clear = gr.Button("Clear Chat")

        # Handle sending messages
        msg.submit(
            respond,
            [msg, chatbot, system_msg, max_tokens, temperature, top_p, model],
            [chatbot]
        ).then(
            lambda: "",
            None,
            msg,
            queue=False
        )

        send.click(
            respond,
            [msg, chatbot, system_msg, max_tokens, temperature, top_p, model],
            [chatbot]
        ).then(
            lambda: "",
            None,
            msg,
            queue=False
        )
        
        # Clear chat history
        clear.click(lambda: None, None, chatbot, queue=False)
        
        # Example prompts
        gr.Examples(
            examples=[
                ["Tell me a short story about a robot learning to paint."],
                ["Explain quantum computing in simple terms."],
                ["Write a haiku about artificial intelligence."]
            ],
            inputs=msg
        )

    return demo

# Create and launch the interface
if __name__ == "__main__":
    demo = create_chat_interface()
    demo.queue()
    demo.launch(
        share=False,  # Disable sharing on Spaces
    )