import json
import os

from PIL import Image
from tqdm import tqdm

IMAGE_DIR = "/opt/datasets/bd/gid/png"
OUTPUT_DIR = "batchapi/data/bboxes"
BATCH_INPUT = "batchapi/batch_input.jsonl"
MAP_FILE = "batchapi/image_file_ids.json"

os.makedirs(OUTPUT_DIR, exist_ok=True)

# Load uploaded OpenAI file IDs
with open(MAP_FILE, "r") as f:
    image_map = json.load(f)


def already_processed(img_name):
    json_name = img_name.rsplit(".", 1)[0] + ".json"
    return os.path.exists(os.path.join(OUTPUT_DIR, json_name))


with open(BATCH_INPUT, "w") as f:
    for img in tqdm(os.listdir(IMAGE_DIR)):
        if not img.lower().endswith((".png", ".jpg", ".jpeg")):
            continue

        if already_processed(img):
            continue

        if img not in image_map:
            print(f"Skipping {img} (not uploaded)")
            continue

        img_path = os.path.join(IMAGE_DIR, img)

        # 🔹 Read image dimensions
        try:
            with Image.open(img_path) as im:
                width, height = im.size
        except Exception as e:
            print(f"Failed to read {img}: {e}")
            continue

        # 🔹 Inject dimensions into prompt
        PROMPT = f"""You are analyzing a satellite/aerial image.

Image dimensions:
- Width: {width} pixels
- Height: {height} pixels

You are analyzing a satellite/aerial image. Examine the image carefully and identify all regions that match these land cover types: farmland, built-up, meadow, forest, and water.

For each land cover type, provide bounding box coordinates for EVERY distinct region you can identify in the image.

CRITICAL: You must respond with ONLY a JSON object. Do not include any explanatory text, markdown formatting, or code blocks. Your entire response should be valid JSON that can be parsed directly.

Return this exact JSON structure:
{{
    "built-up": [[x1, y1, x2, y2], [x1, y1, x2, y2], ...],
    "farmland": [[x1, y1, x2, y2], [x1, y1, x2, y2], ...],
    "forest": [[x1, y1, x2, y2], [x1, y1, x2, y2], ...],
    "meadow": [[x1, y1, x2, y2], [x1, y1, x2, y2], ...],
    "water": [[x1, y1, x2, y2], [x1, y1, x2, y2], ...]
}}

Coordinate format:
- x1, y1: top-left corner (percentages 0-100)
- x2, y2: bottom-right corner (percentages 0-100)
- Use empty array [] if a class is not visible
- Include multiple boxes for the same class if there are multiple regions

Example valid response:
{{"built-up": [[60, 10, 90, 40]], "farmland": [[10, 20, 50, 60]], "forest": [[5, 70, 30, 95]], "meadow": [], "water": [[70, 50, 95, 80]]}}

Now analyze the provided image and return only the JSON object with bounding boxes.
"""
        # """Examine the image carefully and identify all regions that match these land cover types:
        # farmland, built-up, meadow, forest, and water.

        # For each land cover type, provide bounding box coordinates for EVERY distinct region.

        # CRITICAL:
        # - Respond with ONLY a JSON object.
        # - No markdown.
        # - No explanation text.
        # - Valid JSON only.

        # Return this exact JSON structure:
        # {{
        #     "farmland": [[x1, y1, x2, y2]],
        #     "built-up": [[x1, y1, x2, y2]],
        #     "meadow": [[x1, y1, x2, y2]],
        #     "forest": [[x1, y1, x2, y2]],
        #     "water": [[x1, y1, x2, y2]]
        # }}

        # Coordinate format:
        # - x1, y1: top-left corner
        # - x2, y2: bottom-right corner
        # - Coordinates must be in PERCENTAGES (0–100)
        # - Use [] if class is not visible
        # - Include multiple boxes if needed

        # Now analyze the provided image and return only the JSON object."""

        file_id = image_map[img]

        request = {
            "custom_id": img,
            "method": "POST",
            "url": "/v1/responses",
            "body": {
                "model": "gpt-5.2",  # safer for vision batch
                "input": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "input_text", "text": PROMPT},
                            {
                                "type": "input_image",
                                # "image_url": f"openai://{file_id}"
                                "file_id": file_id,
                            },
                        ],
                    }
                ],
                # "response_format": {"type": "json"}
            },
            # }
        }

        f.write(json.dumps(request) + "\n")

print("Batch input file created:", BATCH_INPUT)