"""
Follow-up probe for gpt-image-2. Initial probe revealed:
  - 'format' is rejected (Unknown parameter tools[0].format)
  - 'partial_images' requires stream=True
  - 'background', sizes, qualities, input_fidelity all OK

This script tests the likely-correct parameter names ('output_format',
plus 'compression' standalone) and finally exercises the edit + chained
flows that previously skipped because the seed call used 'format'.

Usage:
    cd api
    python scripts/probe_gpt_image_2_followup.py

Cost: ~$1-2.
"""

import json
import sys
import time
import traceback
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_followup.md"


SCENARIOS: list[tuple[str, dict[str, Any]]] = [
    # Hypothesis: param name is 'output_format', not 'format'
    ("output_format_png",  {"output_format": "png"}),
    ("output_format_jpeg", {"output_format": "jpeg"}),
    ("output_format_webp", {"output_format": "webp"}),

    # Compression alone (no format param)
    ("compression_50_alone",         {"compression": 50}),
    ("output_format_jpeg_compression_50",  {"output_format": "jpeg", "compression": 50}),
    ("output_format_webp_compression_75",  {"output_format": "webp", "compression": 75}),

    # Transparent background paired with png via output_format
    ("transparent_with_output_format_png", {"output_format": "png", "background": "transparent"}),

    # Combination similar to today's gpt-image-1.5 defaults
    ("combo_1024x1024_medium_png", {
        "size": "1024x1024",
        "quality": "medium",
        "output_format": "png",
        "background": "auto",
    }),
]


def _serialize(obj: Any) -> Any:
    if hasattr(obj, "model_dump"):
        try:
            return obj.model_dump()
        except Exception:
            pass
    if isinstance(obj, (list, tuple)):
        return [_serialize(x) for x in obj]
    if isinstance(obj, dict):
        return {k: _serialize(v) for k, v in obj.items()}
    if isinstance(obj, (str, int, float, bool)) or obj is None:
        return obj
    return repr(obj)


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", "format", "background",
                     "compression", "model", "input_fidelity"):
            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
    usage = getattr(response, "usage", None)
    if usage is not None:
        out["usage"] = _serialize(usage)
    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 run_generate(label: str, options: dict[str, Any]) -> dict[str, Any]:
    tool: dict[str, Any] = {"type": "image_generation"}
    tool.update(options)
    started = time.time()
    try:
        resp = client.responses.create(
            model=DRIVER_MODEL,
            input="A small red circle on a clean white background",
            tools=[tool],
        )
        return {
            "label": label, "ok": True,
            "elapsed_s": round(time.time() - started, 2),
            "tool_options": options,
            "summary": _summarize(resp),
        }
    except Exception as e:
        return {
            "label": label, "ok": False,
            "elapsed_s": round(time.time() - started, 2),
            "tool_options": options,
            "error_type": type(e).__name__,
            "error": str(e)[:600],
        }


def run_seed_and_edits() -> list[dict[str, Any]]:
    """Generate a seed image without 'format', then try edit flows."""
    results: list[dict[str, Any]] = []
    print("  Generating seed image (no output_format)...")
    started = time.time()
    try:
        seed_resp = client.responses.create(
            model=DRIVER_MODEL,
            input="A simple red circle on a white background",
            tools=[{"type": "image_generation", "size": "1024x1024", "quality": "low"}],
        )
        seed_b64 = _extract_b64(seed_resp)
        seed_id = getattr(seed_resp, "id", None)
        results.append({
            "label": "seed_no_format",
            "ok": seed_b64 is not None,
            "elapsed_s": round(time.time() - started, 2),
            "summary": _summarize(seed_resp),
        })
    except Exception as e:
        results.append({
            "label": "seed_no_format", "ok": False,
            "elapsed_s": round(time.time() - started, 2),
            "error_type": type(e).__name__, "error": str(e)[:600],
        })
        return results

    if not seed_b64:
        print("  Seed produced no b64; skipping edit probes.")
        return results

    # Edit via input_image alone
    print("  edit_input_image_only ...")
    started = time.time()
    try:
        resp = client.responses.create(
            model=DRIVER_MODEL,
            input="Change the circle's color to bright blue.",
            tools=[{"type": "image_generation", "input_image": seed_b64}],
        )
        results.append({
            "label": "edit_input_image_only", "ok": True,
            "elapsed_s": round(time.time() - started, 2),
            "summary": _summarize(resp),
        })
    except Exception as e:
        results.append({
            "label": "edit_input_image_only", "ok": False,
            "elapsed_s": round(time.time() - started, 2),
            "error_type": type(e).__name__, "error": str(e)[:600],
        })

    # Edit via action: edit
    print("  edit_action_edit ...")
    started = time.time()
    try:
        resp = client.responses.create(
            model=DRIVER_MODEL,
            input="Recolor to green.",
            tools=[{"type": "image_generation", "action": "edit", "input_image": seed_b64}],
        )
        results.append({
            "label": "edit_action_edit", "ok": True,
            "elapsed_s": round(time.time() - started, 2),
            "summary": _summarize(resp),
        })
    except Exception as e:
        results.append({
            "label": "edit_action_edit", "ok": False,
            "elapsed_s": round(time.time() - started, 2),
            "error_type": type(e).__name__, "error": str(e)[:600],
        })

    # previous_response_id chain
    if seed_id:
        print("  previous_response_chain ...")
        started = time.time()
        try:
            resp = client.responses.create(
                model=DRIVER_MODEL,
                previous_response_id=seed_id,
                input="Now make the shape a square instead of a circle.",
                tools=[{"type": "image_generation"}],
            )
            results.append({
                "label": "previous_response_chain", "ok": True,
                "elapsed_s": round(time.time() - started, 2),
                "summary": _summarize(resp),
            })
        except Exception as e:
            results.append({
                "label": "previous_response_chain", "ok": False,
                "elapsed_s": round(time.time() - started, 2),
                "error_type": type(e).__name__, "error": str(e)[:600],
            })

    return results


def render(results: list[dict[str, Any]]) -> str:
    lines = [
        "# gpt-image-2 Probe — Follow-up",
        "",
        f"_Driver model:_ `{DRIVER_MODEL}`  ",
        f"_Generated:_ {time.strftime('%Y-%m-%d %H:%M:%S')}",
        "",
        "## Summary",
        "",
        "| Scenario | Result | Elapsed (s) | Notes |",
        "| --- | --- | --- | --- |",
    ]
    for r in results:
        status = "✅ ok" if r.get("ok") else "❌ fail"
        notes = ""
        if not r.get("ok"):
            notes = f"`{r.get('error_type', '')}`: {(r.get('error') or '')[:80]}"
        else:
            for it in (r.get("summary", {}).get("output_items") or []):
                if it.get("type") == "image_generation_call":
                    bits = []
                    for k in ("size", "quality", "output_format", "format",
                              "background", "compression", "model"):
                        v = it.get(k)
                        if v is not None:
                            bits.append(f"{k}={v}")
                    if "result_b64_len" in it:
                        bits.append(f"b64_len={it['result_b64_len']}")
                    notes = ", ".join(bits)
                    break
        lines.append(f"| {r['label']} | {status} | {r.get('elapsed_s', '')} | {notes} |")
    lines.append("")
    lines.append("## Full results")
    lines.append("")
    for r in results:
        lines.append(f"### {r['label']}")
        lines.append("")
        lines.append("```json")
        lines.append(json.dumps(r, indent=2, default=str))
        lines.append("```")
        lines.append("")
    return "\n".join(lines)


def main() -> None:
    print("gpt-image-2 follow-up probe")
    print(f"Scenarios: {len(SCENARIOS)} param probes + seed/edit/chain")
    confirm = input("Proceed? [y/N] ").strip().lower()
    if confirm != "y":
        sys.exit(0)

    results: list[dict[str, Any]] = []
    for label, options in SCENARIOS:
        sys.stdout.write(f"  {label:<40s} ... ")
        sys.stdout.flush()
        r = run_generate(label, options)
        print("ok" if r["ok"] else f"FAIL ({r.get('error_type')})")
        results.append(r)

    print()
    results.extend(run_seed_and_edits())

    REPORT_PATH.write_text(render(results))
    print()
    print(f"Report: {REPORT_PATH}")
    pc = sum(1 for r in results if r.get("ok"))
    print(f"Passed: {pc}  Failed: {len(results) - pc}")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nInterrupted.")
        sys.exit(1)
    except Exception:
        traceback.print_exc()
        sys.exit(1)