import gradio as gr
import subprocess

# Function to load a model using Hugging Face Spaces and enable GPU
def load_model_with_gpu(model_name):
    print(f"Attempting to load {model_name} with GPU enabled...")
    try:
        # Use subprocess to run hf.space_info and get GPU setting
        result = subprocess.run(
            ["python", "-c", f"from huggingface_hub import space_info; print(space_info('{model_name}').hardware)"],
            capture_output=True,
            text=True,
            check=True
        )
        hardware = result.stdout.strip()
        print(f"Hardware for {model_name}: {hardware}")

        demo = gr.load(name=model_name, src="spaces")

        # Return the loaded model demo
        print(f"Successfully loaded {model_name}")
        return demo

    except Exception as e:
        print(f"Error loading model {model_name}: {e}")
        return None

# Load the models with GPU enabled (if available)
deepseek_r1_distill = load_model_with_gpu("deepseek-ai/DeepSeek-R1-Distill-Qwen-32B")
deepseek_r1 = load_model_with_gpu("deepseek-ai/DeepSeek-R1")
deepseek_r1_zero = load_model_with_gpu("deepseek-ai/DeepSeek-R1-Zero")

# --- Chatbot function ---
def chatbot(input_text, history, model_choice, system_message, max_new_tokens, temperature, top_p):
    history = history or []
    print(f"Input: {input_text}, History: {history}, Model: {model_choice}")

    # Choose the model based on user selection
    if model_choice == "DeepSeek-R1-Distill-Qwen-32B" and deepseek_r1_distill:
        model_demo = deepseek_r1_distill
    elif model_choice == "DeepSeek-R1" and deepseek_r1:
        model_demo = deepseek_r1
    elif model_choice == "DeepSeek-R1-Zero" and deepseek_r1_zero:
        model_demo = deepseek_r1_zero
    else:
        default_response = "Model not selected or could not be loaded."
        history.append((input_text, default_response))
        return history, history, "", model_choice, system_message, max_new_tokens, temperature, top_p
        
    # Adjust the call to the model, remove default_value if not applicable
    model_output = model_demo(input_text, history, max_new_tokens, temperature, top_p, system_message)

    # Check if model_output is iterable and has expected number of elements
    if not isinstance(model_output, (list, tuple)) or len(model_output) < 2:
        error_message = "Model output does not have the expected format."
        history.append((input_text, error_message))
        return history, history, "", model_choice, system_message, max_new_tokens, temperature, top_p
    
    response = model_output[-1][1] if model_output[-1][1] else "Model did not return a response."
    history.append((input_text, response))
    return history, history, "", model_choice, system_message, max_new_tokens, temperature, top_p

# --- Gradio Interface ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # DeepSeek Chatbot
        Created by [ruslanmv.com](https://ruslanmv.com/)

        This is a demo of different DeepSeek models. Select a model, type your message, and click "Submit".
        You can also adjust optional parameters like system message, max new tokens, temperature, and top-p.
        """
    )

    with gr.Row():
        with gr.Column():
            chatbot_output = gr.Chatbot(label="DeepSeek Chatbot", height=500)
            msg = gr.Textbox(label="Your Message", placeholder="Type your message here...")
            with gr.Row():
                submit_btn = gr.Button("Submit", variant="primary")
                clear_btn = gr.ClearButton([msg, chatbot_output])

    # Options moved below the chat interface
    with gr.Row():
        with gr.Accordion("Options", open=True):
            model_choice = gr.Radio(
                choices=["DeepSeek-R1-Distill-Qwen-32B", "DeepSeek-R1", "DeepSeek-R1-Zero"],
                label="Choose a Model",
                value="DeepSeek-R1"
            )
            with gr.Accordion("Optional Parameters", open=False):
                system_message = gr.Textbox(
                    label="System Message",
                    value="You are a friendly Chatbot created by ruslanmv.com",
                    lines=2,
                )
                max_new_tokens = gr.Slider(
                    minimum=1, maximum=4000, value=200, label="Max New Tokens"
                )
                temperature = gr.Slider(
                    minimum=0.10, maximum=4.00, value=0.70, label="Temperature"
                )
                top_p = gr.Slider(
                    minimum=0.10, maximum=1.00, value=0.90, label="Top-p (nucleus sampling)"
                )

    # Maintain chat history
    chat_history = gr.State([])

    # Event handling
    submit_btn.click(
        chatbot,
        [msg, chat_history, model_choice, system_message, max_new_tokens, temperature, top_p],
        [chatbot_output, chat_history, msg, model_choice, system_message, max_new_tokens, temperature, top_p],
    )
    msg.submit(
        chatbot,
        [msg, chat_history, model_choice, system_message, max_new_tokens, temperature, top_p],
        [chatbot_output, chat_history, msg, model_choice, system_message, max_new_tokens, temperature, top_p],
    )

# Launch the demo
if __name__ == "__main__":
    demo.launch()