#!/usr/bin/env python3
"""Enable CLIProxyAPI Codex auths for upstream websocket transport.

CLIProxyAPI decides whether it may forward Codex Responses websocket traffic to
OpenAI from the selected auth file metadata.  Older auth files can be missing
the top-level ``websockets`` flag, which leaves the client-to-CPA leg on
websocket but still uses the non-websocket upstream path.
"""

from __future__ import annotations

import json
import os
import sys
from pathlib import Path
from typing import Any


def _is_codex_auth(path: Path, payload: Any) -> bool:
    if not isinstance(payload, dict):
        return False

    auth_type = str(payload.get("type") or "").strip().lower()
    if auth_type == "codex":
        return True

    # Most CLIProxyAPI Codex auths are named ``codex-...json``.  Keep this as a
    # fallback for older records that may not have a normalized type field yet.
    return path.name.lower().startswith("codex-")


def _write_json(path: Path, payload: dict[str, Any]) -> None:
    tmp_path = path.with_name(f".{path.name}.tmp")
    with tmp_path.open("w", encoding="utf-8") as handle:
        json.dump(payload, handle, ensure_ascii=False, indent=2)
        handle.write("\n")
    os.replace(tmp_path, path)


def enable_codex_auth_websockets(auth_dir: str | os.PathLike[str]) -> int:
    root = Path(auth_dir)
    if not root.is_dir():
        return 0

    changed = 0
    for path in sorted(root.rglob("*.json")):
        try:
            with path.open("r", encoding="utf-8") as handle:
                payload = json.load(handle)
        except (OSError, json.JSONDecodeError):
            continue

        if not _is_codex_auth(path, payload):
            continue
        if payload.get("websockets") is True:
            continue

        payload["websockets"] = True
        _write_json(path, payload)
        changed += 1

    return changed


def main(argv: list[str]) -> int:
    if len(argv) < 2:
        print("usage: codex_ws_auths.py AUTH_DIR [AUTH_DIR...]", file=sys.stderr)
        return 2

    total = 0
    for auth_dir in argv[1:]:
        changed = enable_codex_auth_websockets(auth_dir)
        total += changed
        if changed:
            print(f"enabled codex websockets for {changed} auth file(s) under {auth_dir}", file=sys.stderr)

    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))