import time

from e2b_desktop import Sandbox
from qwen_vl_utils import process_vision_info
from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "Qwen/Qwen2.5-VL-3B-Instruct", torch_dtype="auto", device_map="auto"
)
processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct")


def move_center(desktop, width, height):
    x = width / 2
    y = height / 2
    desktop.move_mouse(x, y)
    print(" - Moved mouse to", x, y)
    desktop.right_click()
    print(" - Right clicked")
    print(" - Waiting 2 seconds...\n")
    time.sleep(2)


print("Starting desktop sandbox...")
desktop = Sandbox(api_key="e2b_1da1c9786e4bd971e6e105cf95a423f588403b73")
width, height = desktop.get_screen_size()
print("Screen size:", width, height)
desktop.stream.start()
stream_url = desktop.stream.get_url()

print("Stream URL:", stream_url)
move_center(desktop, width, height)

# Take a screenshot and save it as "screenshot.png" locally
image = desktop.screenshot()
# Save the image to a file
with open("screenshot.png", "wb") as f:
    f.write(image)


messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": "screenshot.png",
            },
            {"type": "text", "text": "Describe this image."},
        ],
    }
]

# Preparation for inference
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image_inputs, video_inputs = process_vision_info(messages)
inputs = processor(
    text=[text],
    images=image_inputs,
    videos=video_inputs,
    padding=True,
    return_tensors="pt",
)
inputs = inputs.to("cuda")

# Inference: Generation of the output
generated_ids = model.generate(**inputs, max_new_tokens=128)
generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
output_text = processor.batch_decode(
    generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
)
print(output_text)

print("\n> Stopping desktop stream...")
desktop.stream.stop()
print(" - Desktop stream stopped")
print("\n> Killing desktop sandbox...")
# Kill  sandbox.
desktop.kill()
print(" - Desktop sandbox killed")
