import torch
import gradio as gr
import spaces
from transformers import (
    Mistral3ForConditionalGeneration,
    MistralCommonBackend,
)

# Initialize model and tokenizer with ZeroGPU configuration
model_id = "mistralai/Devstral-Small-2-24B-Instruct-2512"

# Load tokenizer
tokenizer = MistralCommonBackend.from_pretrained(model_id)

# Load model with ZeroGPU compatibility
model = Mistral3ForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.float16,  # Use float16 for better GPU efficiency
    low_cpu_mem_usage=True
)

# Move model to GPU when available
if torch.cuda.is_available():
    model = model.to('cuda')

# System prompt
SP = """You are operating as and within Mistral Vibe, a CLI coding-agent built by Mistral AI and powered by default by the Devstral family of models. It wraps Mistral's Devstral models to enable natural language interaction with a local codebase. Use the available tools when helpful.

You can:

- Receive user prompts, project context, and files.
- Send responses and emit function calls (e.g., shell commands, code edits).
- Apply patches, run commands, based on user approvals.

Answer the user's request using the relevant tool(s), if they are available. Check that all the required parameters for each tool call are provided or can reasonably be inferred from context. IF there are no relevant tools or there are missing values for required parameters, ask the user to supply these values; otherwise proceed with the tool calls. If the user provides a specific value for a parameter (for example provided in quotes), make sure to use that value EXACTLY. DO NOT make up values for or ask about optional parameters. Carefully analyze descriptive terms in the request as they may indicate required parameter values that should be included even if not explicitly quoted.

Always try your hardest to use the tools to answer the user's request. If you can't use the tools, explain why and ask the user for more information.

Act as an agentic assistant, if a user asks for a long task, break it down and do it step by step.

When you want to commit changes, you will always use the 'git commit' bash command. It will always
be suffixed with a line telling it was generated by Mistral Vibe with the appropriate co-authoring information.
The format you will always uses is the following heredoc.

```bash
git commit -m "<Commit message here>

Generated by Mistral Vibe.
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>"
```"""

# Tools configuration
tools = [
    {
        "type": "function",
        "function": {
            "name": "add_number",
            "description": "Add two numbers.",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "string", "description": "The first number."},
                    "b": {"type": "string", "description": "The second number."},
                },
                "required": ["a", "b"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "multiply_number",
            "description": "Multiply two numbers.",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "string", "description": "The first number."},
                    "b": {"type": "string", "description": "The second number."},
                },
                "required": ["a", "b"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "substract_number",
            "description": "Substract two numbers.",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "string", "description": "The first number."},
                    "b": {"type": "string", "description": "The second number."},
                },
                "required": ["a", "b"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "write_a_story",
            "description": "Write a story about science fiction and people with badass laser sabers.",
            "parameters": {},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "terminal",
            "description": "Perform operations from the terminal.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The command you wish to launch, e.g `ls`, `rm`, ...",
                    },
                    "args": {
                        "type": "string",
                        "description": "The arguments to pass to the command.",
                    },
                },
                "required": ["command"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "python",
            "description": "Call a Python interpreter with some Python code that will be ran.",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {
                        "type": "string",
                        "description": "The Python code to run",
                    },
                    "result_variable": {
                        "type": "string",
                        "description": "Variable containing the result you'd like to retrieve from the execution.",
                    },
                },
                "required": ["code", "result_variable"],
            },
        },
    },
]

@spaces.GPU(duration=60)  # Use ZeroGPU with 60 second duration
def chat_function_gpu(message, history):
    try:
        # Prepare input messages
        messages = [
            {
                "role": "system",
                "content": SP,
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": message,
                    }
                ],
            },
        ]

        # Tokenize input
        tokenized = tokenizer.apply_chat_template(
            conversation=messages,
            tools=tools,
            return_tensors="pt",
            return_dict=True,
        )

        input_ids = tokenized["input_ids"].to(device="cuda")

        # Generate output with GPU acceleration
        output = model.generate(
            input_ids,
            max_new_tokens=200,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            num_return_sequences=1
        )[0]

        # Decode and return response
        decoded_output = tokenizer.decode(output[len(tokenized["input_ids"][0]) :])
        return decoded_output

    except Exception as e:
        return f"Error processing your request: {str(e)}"

# Fallback CPU function for when GPU is not available
def chat_function_cpu(message, history):
    try:
        # Prepare input messages
        messages = [
            {
                "role": "system",
                "content": SP,
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": message,
                    }
                ],
            },
        ]

        # Tokenize input with CPU configuration
        tokenized = tokenizer.apply_chat_template(
            conversation=messages,
            tools=tools,
            return_tensors="pt",
            return_dict=True,
        )

        input_ids = tokenized["input_ids"].to(device="cpu")

        # Generate output with CPU-optimized settings
        output = model.generate(
            input_ids,
            max_new_tokens=100,  # Reduced for CPU performance
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            num_return_sequences=1
        )[0]

        # Decode and return response
        decoded_output = tokenizer.decode(output[len(tokenized["input_ids"][0]) :])
        return decoded_output

    except Exception as e:
        return f"Error processing your request: {str(e)}"

# Create custom theme optimized for ZeroGPU
custom_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
    text_size="lg",
    spacing_size="lg",
    radius_size="md"
).set(
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_700",
    block_title_text_weight="600",
)

# Create Gradio interface with ZeroGPU support
with gr.Blocks() as demo:
    chatbot = gr.Chatbot(height=600)
    msg = gr.Textbox(
        label="Your Message",
        placeholder="Type your message here...",
        lines=3
    )

    # Clear button
    clear_btn = gr.ClearButton([msg, chatbot])

    # Submit button with loading indicator
    submit_btn = gr.Button("Send", variant="primary")

    # Status indicator
    status_text = gr.Markdown("Ready for your input...")

    def update_status(text):
        return text

    # Event handlers with status updates
    def handle_submit(message, history):
        if torch.cuda.is_available():
            status_text.value = "Processing with ZeroGPU acceleration..."
            response = chat_function_gpu(message, history)
        else:
            status_text.value = "Processing with CPU (ZeroGPU quota may be exhausted)..."
            response = chat_function_cpu(message, history)
        status_text.value = "Ready for your input..."
        return response

    msg.submit(
        fn=handle_submit,
        inputs=[msg, chatbot],
        outputs=[chatbot],
        api_visibility="public"
    )

    submit_btn.click(
        fn=handle_submit,
        inputs=[msg, chatbot],
        outputs=[chatbot],
        api_visibility="public"
    )

    # Examples with ZeroGPU information
    gr.Examples(
        examples=[
            "Can you implement in Python a method to compute the fibonnaci sequence at the nth element with n a parameter passed to the function?",
            "What are the available tools I can use?",
            "Can you write a story about science fiction with laser sabers?"
        ],
        inputs=msg,
        label="Example Prompts (Powered by ZeroGPU when available)"
    )

# Launch with custom theme and ZeroGPU settings
demo.launch(
    theme=custom_theme,
    footer_links=[
        {
            "label": "Built with anycoder",
            "url": "https://huggingface.co/spaces/akhaliq/anycoder"
        },
        {
            "label": "Mistral AI",
            "url": "https://mistral.ai"
        },
        {
            "label": "Hugging Face ZeroGPU",
            "url": "https://huggingface.co/docs/hub/spaces-zerogpu"
        },
        {
            "label": "Hugging Face Spaces",
            "url": "https://huggingface.co/spaces"
        }
    ],
    share=False,  # Disable share for Spaces deployment
    max_threads=4,  # Allow more threads for GPU processing
    show_error=True,
    enable_queue=True
)