import base64, boto3, json, sys, time
from pathlib import Path

REGION = "us-east-1"
BUCKET = "medgemma-async-2026"
ENDPOINT = "medgemma-vllm-endpoint"
IMAGE_PATH = Path(sys.argv[1] if len(sys.argv) > 1 else "demo_images/ROCOv2_2023_valid_000001.jpg")
QUESTION = "What abnormalities do you see in this image?"

sagemaker = boto3.client("sagemaker-runtime", region_name=REGION)
s3 = boto3.client("s3", region_name=REGION)

image_b64 = base64.b64encode(IMAGE_PATH.read_bytes()).decode()

payload = {
    "messages": [
        {
            "role": "user",
            "content": [
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"}},
                {"type": "text", "text": QUESTION},
            ],
        }
    ],
    "max_tokens": 300,
}

input_key = f"input/test_vision_{int(time.time())}.json"
s3.put_object(Bucket=BUCKET, Key=input_key, Body=json.dumps(payload).encode())

response = sagemaker.invoke_endpoint_async(
    EndpointName=ENDPOINT,
    InputLocation=f"s3://{BUCKET}/{input_key}",
    ContentType="application/json",
)
output_uri = response["OutputLocation"]
output_key = output_uri.replace(f"s3://{BUCKET}/", "")

for i in range(60):
    time.sleep(10)
    try:
        obj = s3.get_object(Bucket=BUCKET, Key=output_key)
        print(obj["Body"].read().decode())
        break
    except s3.exceptions.NoSuchKey:
        print(f"  waiting... {(i+1)*10}s")