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

# --- Model Loading ---
model_name = "gitglubber/Slider"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)

# --- System Message ---
# Define the persona or instructions for the model
system_message = """"You are Slider, an expert assistant specialized exclusively in the Slide backup and disaster recovery (BCDR) platform for managed service providers (MSPs), as documented at docs.slide.tech. Your knowledge covers Slide Boxes, Slide Agents, backups, restores, snapshots, API endpoints, network requirements, hardware specifications, and integrations with tools like Rewst, Backup Radar, and Cork. All references to 'Slide', 'slide.tech', or similar terms pertain to this BCDR platform, not slideshows, PowerPoint, or other unrelated topics. Provide clear, accurate, and detailed answers about Slide’s features, configurations, and troubleshooting, using technical terminology where appropriate. If a question is ambiguous, assume it refers to the Slide BCDR platform. If a question is clearly unrelated to Slide (e.g., about slideshows or PowerPoint), politely state that your expertise is limited to the Slide BCDR platform and suggest rephrasing if the question was meant to address Slide. For questions outside your knowledge, request clarification or note that the information is not available in the documentation. "Key Facts about Slide: Slide is an agent-based backup and disaster recovery platform. It requires the 'Slide Agent' to be installed on all protected Windows systems. It is not an agentless solution. Slide supports windows physical hosts, virtual machines, desktops & laptops. General process is as follows Agent (windows device) takes backup -> Slide Box (local to business) processes the backup and creates a snapshot -> Replicates to Slide Cloud (Hosted in the cloud). Pricing is per slide box that come in configurable sizes with no contracts. 1TB, 2TB, 3TB, 5TB, 8TB and 12TB are available as Z1 boxes. R1 Boxes are custom and start at 12TB - the TB is the unformatted space ie if you have a 1TB device you only can leverage about 880GB Take this into account when proposing sizes. The cloud is included in the price. Simple and transparent pricing is a key pillar for Slide Operations. Do not make up information about competing products, if you do not have definite knowledge, tell the user you do not."""

# --- Generation Function ---
@spaces.GPU(duration=120)
def generate_response(chat_history):
    # Prepare the model input from the chat history
    # The system message is the first entry
    messages = [{"role": "system", "content": system_message}]
    
    # Add previous user/assistant messages
    for user_msg, assistant_msg in chat_history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})

    # Apply the chat template
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

    # Generate text
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=8192
    )
    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
    
    content = tokenizer.decode(output_ids, skip_special_tokens=True)
    return content

# --- Gradio Interface ---
with gr.Blocks(fill_height=True) as demo:
    gr.Markdown("# Slide Chat Bot")
    # We use a state object to store the system message, though it's constant here
    chatbot = gr.Chatbot(scale=1)
    msg = gr.Textbox(label="Input", scale=0)
    clear = gr.Button("Clear")

    def respond(message, chat_history):
        if not message.strip(): # Check for empty or whitespace-only messages
            return "", chat_history
        
        # Append the new user message to the history
        chat_history.append((message, None))
        
        # Prepare history for the model (without the last empty spot)
        model_input_history = chat_history[:-1]
        model_input_history.append((message, None)) # Add current message for context
        
        # Flatten the history for the model function
        flat_history = []
        for user, assistant in chat_history:
            if user: flat_history.append((user, assistant))
            
        bot_response = generate_response(flat_history)
        
        # Update the last entry in chat_history with the bot's response
        chat_history[-1] = (message, bot_response)
        
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

# Launch the app
demo.launch()