import requests
import time

# Define input and output file paths
INPUT_FILE = "blah.png"
OUTPUT_FILE = "blaho.png"

# Your API key for Claid.ai
API_KEY = "YOUR_API_KEY"

# Define the API endpoint and headers
TASK_URL = "https://api.claid.ai/v1/task/"
HEADERS = {"Authorization": f"Bearer {API_KEY}"}

# Step 1: Create a task
task_data = {
    "type": "photo",
    "input": [{"source": "UPLOAD"}],
    "params": {"scale": 2},  # Scale 2x
}

response = requests.post(TASK_URL, json=task_data, headers=HEADERS)
response.raise_for_status()
task_id = response.json()["id"]

print(f"Task created with ID: {task_id}")

# Step 2: Upload the image
upload_url = response.json()["input"][0]["upload_url"]
with open(INPUT_FILE, "rb") as f:
    upload_response = requests.put(upload_url, data=f)
upload_response.raise_for_status()
print(f"Image {INPUT_FILE} uploaded successfully.")

# Step 3: Wait for processing to complete
print("Waiting for processing to complete...")
while True:
    task_status = requests.get(f"{TASK_URL}{task_id}/", headers=HEADERS).json()
    if task_status["status"] == "finished":
        print("Processing complete.")
        break
    elif task_status["status"] == "failed":
        print("Task failed:", task_status)
        exit(1)
    time.sleep(5)  # Wait 5 seconds before checking again

# Step 4: Download the upscaled image
output_url = task_status["output"][0]["file"]
output_response = requests.get(output_url)
output_response.raise_for_status()

with open(OUTPUT_FILE, "wb") as f:
    f.write(output_response.content)

print(f"Upscaled image saved to {OUTPUT_FILE}")