"""
Final probe: how does gpt-image-2 accept an arbitrary input image for edit?
- 'input_image' as a tool option is rejected.
- 'previous_response_id' works but only chains within Responses API state.

Hypothesis: pass image as a content part in the request `input`
(multimodal user message) and let the image_generation tool consume it.

Usage:
    cd api
    python scripts/probe_gpt_image_2_edit.py
"""

import base64
import json
import sys
import time
from pathlib import Path
from typing import Any, Optional

from dotenv import load_dotenv

env_path = Path(__file__).resolve().parent.parent / '.env'
load_dotenv(dotenv_path=env_path)

from openai import OpenAI

client = OpenAI()
DRIVER_MODEL = "gpt-5"
REPORT_PATH = Path(__file__).resolve().parent / "gpt_image_2_probe_edit.md"


def _summarize(response: Any) -> dict[str, Any]:
    out: dict[str, Any] = {
        "id": getattr(response, "id", None),
        "status": getattr(response, "status", None),
    }
    items = []
    for item in (getattr(response, "output", None) or []):
        entry: dict[str, Any] = {"type": getattr(item, "type", None)}
        for attr in ("size", "quality", "output_format", "background", "model"):
            if hasattr(item, attr):
                entry[attr] = getattr(item, attr)
        result = getattr(item, "result", None)
        if isinstance(result, str):
            entry["result_b64_len"] = len(result)
        items.append(entry)
    out["output_items"] = items
    return out


def _extract_b64(response: Any) -> Optional[str]:
    for item in (getattr(response, "output", None) or []):
        if getattr(item, "type", None) == "image_generation_call":
            r = getattr(item, "result", None)
            if isinstance(r, str):
                return r
    return None


def make_seed() -> tuple[Optional[str], Optional[str]]:
    print("Generating seed image...")
    resp = client.responses.create(
        model=DRIVER_MODEL,
        input="A simple red circle on a white background",
        tools=[{"type": "image_generation", "size": "1024x1024", "quality": "low"}],
    )
    return _extract_b64(resp), getattr(resp, "id", None)


def try_edit_via_content_part(seed_b64: str) -> dict[str, Any]:
    """Pass image as multimodal content part in input."""
    started = time.time()
    label = "edit_via_content_part_input_image"
    try:
        resp = client.responses.create(
            model=DRIVER_MODEL,
            input=[
                {
                    "role": "user",
                    "content": [
                        {"type": "input_text",
                         "text": "Recolor the circle to bright blue. Keep the white background."},
                        {"type": "input_image",
                         "image_url": f"data:image/png;base64,{seed_b64}"},
                    ],
                }
            ],
            tools=[{"type": "image_generation"}],
        )
        return {
            "label": label, "ok": True,
            "elapsed_s": round(time.time() - started, 2),
            "summary": _summarize(resp),
        }
    except Exception as e:
        return {
            "label": label, "ok": False,
            "elapsed_s": round(time.time() - started, 2),
            "error_type": type(e).__name__,
            "error": str(e)[:600],
        }


def try_edit_via_file_id(seed_bytes: bytes) -> dict[str, Any]:
    """Upload via Files API and reference by file_id in content part."""
    started = time.time()
    label = "edit_via_file_id"
    try:
        # Upload bytes to Files API
        from io import BytesIO
        buf = BytesIO(seed_bytes)
        buf.name = "seed.png"
        upload = client.files.create(file=buf, purpose="vision")
        file_id = upload.id
        resp = client.responses.create(
            model=DRIVER_MODEL,
            input=[
                {
                    "role": "user",
                    "content": [
                        {"type": "input_text",
                         "text": "Recolor the circle to green. Keep the background."},
                        {"type": "input_image", "file_id": file_id},
                    ],
                }
            ],
            tools=[{"type": "image_generation"}],
        )
        return {
            "label": label, "ok": True,
            "elapsed_s": round(time.time() - started, 2),
            "file_id": file_id,
            "summary": _summarize(resp),
        }
    except Exception as e:
        return {
            "label": label, "ok": False,
            "elapsed_s": round(time.time() - started, 2),
            "error_type": type(e).__name__,
            "error": str(e)[:600],
        }


def main() -> None:
    confirm = input("Proceed with edit-flow probe (~$1)? [y/N] ").strip().lower()
    if confirm != "y":
        sys.exit(0)

    results: list[dict[str, Any]] = []
    seed_b64, seed_id = make_seed()
    if not seed_b64:
        print("Seed failed.")
        sys.exit(1)

    print("  edit_via_content_part_input_image ...")
    r = try_edit_via_content_part(seed_b64)
    print(f"    {'ok' if r['ok'] else 'FAIL: ' + str(r.get('error'))[:100]}")
    results.append(r)

    print("  edit_via_file_id ...")
    seed_bytes = base64.b64decode(seed_b64)
    r = try_edit_via_file_id(seed_bytes)
    print(f"    {'ok' if r['ok'] else 'FAIL: ' + str(r.get('error'))[:100]}")
    results.append(r)

    REPORT_PATH.write_text(json.dumps(results, indent=2, default=str))
    print(f"\nReport: {REPORT_PATH}")


if __name__ == "__main__":
    main()