"""
Probe script for gpt-image-2 via the Responses API image_generation tool.

gpt-image-2 is not exposed through /v1/images/generations or /v1/images/edits.
Per OpenAI's docs, it's reachable only via the Responses API with
tools=[{"type": "image_generation", ...}], driven by a text model
(e.g. gpt-5). The underlying image model (gpt-image-1, -1.5, -2, -1-mini) is
selected by the system, not by the caller.

This script empirically probes the tool's parameter surface and writes a
markdown report to scripts/gpt_image_2_probe.md. Run it once before
implementing the /gpt-image-2 backend route — its findings shape the
GptImage2Request schema, edit semantics, and SSE protocol.

Usage:
    cd api
    python scripts/probe_gpt_image_2.py

Cost: ~$2-5 in API calls.
"""

import base64
import io
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.md"

# Each scenario: (label, tool_options_dict, prompt_override_or_None)
SCENARIOS: list[tuple[str, dict[str, Any], Optional[str]]] = [
    # 1. Minimal generate
    ("minimal_generate", {}, None),

    # 2. Sizes
    ("size_1024x1024", {"size": "1024x1024"}, None),
    ("size_1024x1536", {"size": "1024x1536"}, None),
    ("size_1536x1024", {"size": "1536x1024"}, None),
    ("size_auto",      {"size": "auto"}, None),

    # 3. Qualities
    ("quality_low",    {"quality": "low"}, None),
    ("quality_medium", {"quality": "medium"}, None),
    ("quality_high",   {"quality": "high"}, None),
    ("quality_auto",   {"quality": "auto"}, None),

    # 4. Formats + compression
    ("format_png",         {"format": "png"}, None),
    ("format_jpeg",        {"format": "jpeg"}, None),
    ("format_webp",        {"format": "webp"}, None),
    ("jpeg_compression_0",   {"format": "jpeg", "compression": 0}, None),
    ("jpeg_compression_50",  {"format": "jpeg", "compression": 50}, None),
    ("jpeg_compression_100", {"format": "jpeg", "compression": 100}, None),
    ("webp_compression_50",  {"format": "webp", "compression": 50}, None),

    # 5. Background
    ("background_transparent",     {"background": "transparent", "format": "png"}, None),
    ("background_opaque",          {"background": "opaque"}, None),
    ("background_auto",            {"background": "auto"}, None),
    ("background_transparent_jpeg", {"background": "transparent", "format": "jpeg"}, None),

    # 6. input_fidelity (Images-API param; tool docs don't list it — confirm)
    ("input_fidelity_low",  {"input_fidelity": "low"}, None),
    ("input_fidelity_high", {"input_fidelity": "high"}, None),

    # 7. partial_images (no stream — confirm tool accepts the option even
    #    when the request itself isn't streamed)
    ("partial_images_2_no_stream", {"partial_images": 2}, None),
    ("partial_images_3_no_stream", {"partial_images": 3}, None),
]


def _serialize(obj: Any) -> Any:
    """Best-effort JSON-friendly serialization of SDK response objects."""
    if hasattr(obj, "model_dump"):
        try:
            return obj.model_dump()
        except Exception:
            pass
    if hasattr(obj, "to_dict"):
        try:
            return obj.to_dict()
        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(response: Any) -> dict[str, Any]:
    """Pull the interesting bits out of a Responses API result."""
    summary: dict[str, Any] = {}
    try:
        summary["id"] = getattr(response, "id", None)
        summary["status"] = getattr(response, "status", None)
        usage = getattr(response, "usage", None)
        if usage is not None:
            summary["usage"] = _serialize(usage)
        output = getattr(response, "output", None) or []
        item_summaries = []
        for item in output:
            item_type = getattr(item, "type", None)
            entry: dict[str, Any] = {"type": item_type}
            for attr in ("id", "status", "size", "quality", "format",
                         "background", "compression", "model", "name"):
                if hasattr(item, attr):
                    entry[attr] = getattr(item, attr)
            result = getattr(item, "result", None)
            if isinstance(result, str):
                entry["result_b64_len"] = len(result)
            elif result is not None:
                entry["result_repr"] = repr(result)[:200]
            item_summaries.append(entry)
        summary["output_items"] = item_summaries
    except Exception as e:
        summary["summarize_error"] = f"{type(e).__name__}: {e}"
    return summary


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


def run_generate(label: str, options: dict[str, Any], prompt: Optional[str]) -> dict[str, Any]:
    tool: dict[str, Any] = {"type": "image_generation"}
    tool.update(options)
    used_prompt = prompt or "A small red circle on a clean white background"
    started = time.time()
    try:
        response = client.responses.create(
            model=DRIVER_MODEL,
            input=used_prompt,
            tools=[tool],
        )
        elapsed = time.time() - started
        return {
            "label": label,
            "ok": True,
            "elapsed_s": round(elapsed, 2),
            "tool_options": options,
            "prompt": used_prompt,
            "summary": _summarize_response(response),
        }
    except Exception as e:
        elapsed = time.time() - started
        return {
            "label": label,
            "ok": False,
            "elapsed_s": round(elapsed, 2),
            "tool_options": options,
            "prompt": used_prompt,
            "error_type": type(e).__name__,
            "error": str(e)[:600],
        }


def run_edit_with_input_image(b64_png: str) -> dict[str, Any]:
    """Test edit / inpainting flow by passing an input_image."""
    tool: dict[str, Any] = {
        "type": "image_generation",
        "input_image": b64_png,
        # action variants are also probed — leave default first
    }
    started = time.time()
    try:
        response = client.responses.create(
            model=DRIVER_MODEL,
            input="Change the circle's color to bright blue. Keep the background.",
            tools=[tool],
        )
        return {
            "label": "edit_with_input_image",
            "ok": True,
            "elapsed_s": round(time.time() - started, 2),
            "summary": _summarize_response(response),
        }
    except Exception as e:
        return {
            "label": "edit_with_input_image",
            "ok": False,
            "elapsed_s": round(time.time() - started, 2),
            "error_type": type(e).__name__,
            "error": str(e)[:600],
        }


def run_edit_with_action_edit(b64_png: str) -> dict[str, Any]:
    tool: dict[str, Any] = {
        "type": "image_generation",
        "action": "edit",
        "input_image": b64_png,
    }
    started = time.time()
    try:
        response = client.responses.create(
            model=DRIVER_MODEL,
            input="Recolor to green.",
            tools=[tool],
        )
        return {
            "label": "edit_action_edit",
            "ok": True,
            "elapsed_s": round(time.time() - started, 2),
            "summary": _summarize_response(response),
        }
    except Exception as e:
        return {
            "label": "edit_action_edit",
            "ok": False,
            "elapsed_s": round(time.time() - started, 2),
            "error_type": type(e).__name__,
            "error": str(e)[:600],
        }


def run_previous_response_chain(initial_response_id: str) -> dict[str, Any]:
    started = time.time()
    try:
        response = client.responses.create(
            model=DRIVER_MODEL,
            previous_response_id=initial_response_id,
            input="Now make the shape a square instead of a circle.",
            tools=[{"type": "image_generation"}],
        )
        return {
            "label": "previous_response_chain",
            "ok": True,
            "elapsed_s": round(time.time() - started, 2),
            "summary": _summarize_response(response),
        }
    except Exception as e:
        return {
            "label": "previous_response_chain",
            "ok": False,
            "elapsed_s": round(time.time() - started, 2),
            "error_type": type(e).__name__,
            "error": str(e)[:600],
        }


def run_streaming_partials() -> dict[str, Any]:
    """Probe streaming with partial_images=2."""
    started = time.time()
    events_log: list[dict[str, Any]] = []
    try:
        stream = client.responses.create(
            model=DRIVER_MODEL,
            input="A serene mountain landscape at sunrise.",
            tools=[{"type": "image_generation", "partial_images": 2}],
            stream=True,
        )
        for event in stream:
            evt_type = getattr(event, "type", None)
            entry: dict[str, Any] = {"type": evt_type, "t": round(time.time() - started, 2)}
            for attr in ("partial_image_index", "item_id", "output_index", "sequence_number"):
                if hasattr(event, attr):
                    entry[attr] = getattr(event, attr)
            b64 = getattr(event, "b64_json", None) or getattr(event, "partial_image_b64", None)
            if isinstance(b64, str):
                entry["b64_len"] = len(b64)
            events_log.append(entry)
        return {
            "label": "streaming_partials",
            "ok": True,
            "elapsed_s": round(time.time() - started, 2),
            "event_count": len(events_log),
            "events": events_log,
        }
    except Exception as e:
        return {
            "label": "streaming_partials",
            "ok": False,
            "elapsed_s": round(time.time() - started, 2),
            "error_type": type(e).__name__,
            "error": str(e)[:600],
            "events_so_far": events_log,
        }


def render_report(results: list[dict[str, Any]]) -> str:
    lines: list[str] = []
    lines.append("# gpt-image-2 Probe Report")
    lines.append("")
    lines.append(f"_Driver model:_ `{DRIVER_MODEL}`  ")
    lines.append(f"_Generated:_ {time.strftime('%Y-%m-%d %H:%M:%S')}")
    lines.append("")
    lines.append("## Summary")
    lines.append("")
    lines.append("| Scenario | Result | Elapsed (s) | Notes |")
    lines.append("| --- | --- | --- | --- |")
    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:
            summary = r.get("summary") or {}
            items = summary.get("output_items") or []
            for it in items:
                if it.get("type") == "image_generation_call":
                    bits: list[str] = []
                    for k in ("size", "quality", "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")
        # Drop any heavy b64 payloads before dumping
        cleaned = {k: v for k, v in r.items() if k != "events"}
        if r.get("label") == "streaming_partials":
            cleaned["events"] = r.get("events", [])
        lines.append(json.dumps(cleaned, indent=2, default=str))
        lines.append("```")
        lines.append("")
    return "\n".join(lines)


def main() -> None:
    print("=" * 60)
    print("  gpt-image-2 Probe (Responses API + image_generation tool)")
    print("=" * 60)
    print()
    print(f"Driver model: {DRIVER_MODEL}")
    print(f"Scenarios: {len(SCENARIOS)} parameter probes + edit + chain + streaming")
    print(f"Estimated cost: ~$2-5")
    print(f"Report: {REPORT_PATH}")
    print()

    confirm = input("Proceed? [y/N] ").strip().lower()
    if confirm != "y":
        print("Aborted.")
        sys.exit(0)

    results: list[dict[str, Any]] = []

    # Parameter sweeps
    for label, options, prompt in SCENARIOS:
        sys.stdout.write(f"  {label:<32s} ... ")
        sys.stdout.flush()
        result = run_generate(label, options, prompt)
        print("ok" if result["ok"] else f"FAIL ({result.get('error_type')})")
        results.append(result)

    # Edit-flow probes — need a source image
    print()
    print("Generating source image for edit probes...")
    source_b64: Optional[str] = None
    seed = run_generate("edit_seed", {"size": "1024x1024", "quality": "low", "format": "png"},
                        "A simple red circle on a white background")
    if seed["ok"]:
        # Extract b64 from the original (un-summarized) response — re-run extraction
        # by calling once more since _summarize discards the bytes. Cheaper to
        # just call again with a fresh request.
        try:
            resp = client.responses.create(
                model=DRIVER_MODEL,
                input="A simple red circle on a white background",
                tools=[{"type": "image_generation", "size": "1024x1024",
                        "quality": "low", "format": "png"}],
            )
            source_b64 = _extract_image_b64(resp)
            initial_response_id = getattr(resp, "id", None)
        except Exception as e:
            print(f"  Could not regenerate source image: {e}")
            initial_response_id = None
    else:
        print("  Seed generation failed; skipping edit probes.")
        initial_response_id = None

    if source_b64:
        sys.stdout.write("  edit_with_input_image            ... ")
        sys.stdout.flush()
        r = run_edit_with_input_image(source_b64)
        print("ok" if r["ok"] else f"FAIL ({r.get('error_type')})")
        results.append(r)

        sys.stdout.write("  edit_action_edit                 ... ")
        sys.stdout.flush()
        r = run_edit_with_action_edit(source_b64)
        print("ok" if r["ok"] else f"FAIL ({r.get('error_type')})")
        results.append(r)

    if initial_response_id:
        sys.stdout.write("  previous_response_chain          ... ")
        sys.stdout.flush()
        r = run_previous_response_chain(initial_response_id)
        print("ok" if r["ok"] else f"FAIL ({r.get('error_type')})")
        results.append(r)

    # Streaming
    sys.stdout.write("  streaming_partials               ... ")
    sys.stdout.flush()
    r = run_streaming_partials()
    print("ok" if r["ok"] else f"FAIL ({r.get('error_type')})")
    results.append(r)

    # Write report
    REPORT_PATH.write_text(render_report(results))
    print()
    print(f"Report written: {REPORT_PATH}")

    pass_count = sum(1 for r in results if r.get("ok"))
    fail_count = len(results) - pass_count
    print(f"Passed: {pass_count}  Failed: {fail_count}")


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