"""Keep image-generation pricing on the shared cost-map key.

Run at BUILD time (the runtime rootfs is read-only under Cloudron). Non-fatal:
if the target file or anchor strings are missing (e.g. after an upstream
version bump), it prints a warning and exits 0 so the build — and the
gateway — still succeeds.

Why: LiteLLM's image cost calculator (default_image_cost_calculator) looks
pricing up ONLY in the shared ``{provider}/{model}`` key of
``litellm.model_cost`` — it never consults the per-deployment entry. But the
router deliberately strips per-deployment pricing from that shared key
("so one deployment's overrides don't pollute sibling deployments"), which
makes every successful image generation on a custom backend invisible:
the cost lookup raises, the spend-log write is skipped, and SpendLogs /
the dashboards show nothing. Upstream: BerriAI/litellm#10023 and #12471,
closed not-planned.

Patch: for image_generation deployments, re-add the custom pricing fields to
the shared-key entry. Safe for us: the co-op gateway has a single deployment
per image backend, so there are no sibling deployments to pollute.
"""
import sys
import glob

ANCHOR = "        # Always register the (possibly mode-preserved) shared backend info."
INSERT = '''        # CO-OP PATCH (2026-10-19): image-generation deployments must keep
        # their pricing on the shared cost-map key — the image cost calculator
        # (default_image_cost_calculator) reads ONLY this key and never the
        # per-deployment model_id entry, so stripping pricing here makes every
        # successful image generation invisible in SpendLogs (upstream:
        # BerriAI/litellm#10023, #12471, closed not-planned). Safe: the co-op
        # gateway has a single deployment per image backend — no siblings.
        if deployment_mode == "image_generation":
            for _pk, _pv in model_info.items():
                if _pk in CustomPricingLiteLLMParams.model_fields and _pv is not None:
                    shared_model_info[_pk] = _pv

'''

candidates = glob.glob(
    "/app/.venv/lib/python*/site-packages/litellm/router.py"
)

patched_any = False
for p in candidates:
    with open(p) as f:
        src = f.read()

    if "CO-OP PATCH (2026-10-19)" in src:
        print(f"INFO: image-cost patch already present in {p}")
        continue

    if ANCHOR not in src:
        print(f"WARN: anchor not found in {p}; patch not applied (non-fatal)",
              file=sys.stderr)
        continue

    new = src.replace(ANCHOR, INSERT + ANCHOR, 1)

    if new == src:
        print(f"WARN: insertion failed in {p} (non-fatal)", file=sys.stderr)
        continue

    with open(p, "w") as f:
        f.write(new)
    print(f"INFO: patched image-cost registration in {p}")
    patched_any = True

if not candidates:
    print("WARN: router.py not found; patch not applied (non-fatal)",
          file=sys.stderr)
elif not patched_any:
    print("WARN: no files patched (non-fatal)", file=sys.stderr)

sys.exit(0)