"""Bounded Codex app-server query and child-process cleanup."""

from __future__ import annotations

import json
import os
import selectors
import subprocess
import time
from typing import Any


CODEX_BIN = os.environ.get("ATLAS_AI_CODEX_BIN", "/opt/data/tools/bin/codex")
CODEX_HOME = os.environ.get("CODEX_HOME", "/runtime-access/codex")
CODEX_CLEANUP_TIMEOUT_SECONDS = 5
CODEX_QUERY_TIMEOUT_SECONDS = min(
    60.0,
    max(20.0, float(os.environ.get("ATLAS_AI_CODEX_QUERY_TIMEOUT_SECONDS", "45"))),
)


def _close_stream(stream: Any) -> bool:
    """Close a subprocess pipe without allowing cleanup errors to escape."""
    if stream is None:
        return True
    try:
        stream.close()
    except Exception:
        return False
    return True


def _cleanup_codex_process(
    process: subprocess.Popen[str], selector: selectors.BaseSelector | None
) -> bool:
    """Stop and reap a Codex child within two fixed waits, then close its pipes."""
    cleanup_ok = True
    if selector is not None:
        try:
            selector.close()
        except Exception:
            cleanup_ok = False

    # Closing the request pipe first also gives a responsive app-server an EOF.
    cleanup_ok = _close_stream(getattr(process, "stdin", None)) and cleanup_ok
    try:
        process.terminate()
    except ProcessLookupError:
        pass
    except Exception:
        cleanup_ok = False

    needs_kill = False
    try:
        process.wait(timeout=CODEX_CLEANUP_TIMEOUT_SECONDS)
    except subprocess.TimeoutExpired:
        needs_kill = True
    except Exception:
        cleanup_ok = False
        needs_kill = True

    if needs_kill:
        try:
            process.kill()
        except ProcessLookupError:
            pass
        except Exception:
            cleanup_ok = False
        try:
            process.wait(timeout=CODEX_CLEANUP_TIMEOUT_SECONDS)
        except Exception:
            cleanup_ok = False

    cleanup_ok = _close_stream(getattr(process, "stdout", None)) and cleanup_ok
    cleanup_ok = _close_stream(getattr(process, "stderr", None)) and cleanup_ok
    return cleanup_ok


def query_codex(timeout: float | None = None) -> tuple[dict[str, Any], dict[str, Any]]:
    """Read Codex account quota and usage through its structured app-server protocol."""
    timeout = CODEX_QUERY_TIMEOUT_SECONDS if timeout is None else timeout
    process = subprocess.Popen(
        [CODEX_BIN, "app-server", "--stdio"],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.DEVNULL,
        text=True,
        bufsize=1,
        env={**os.environ, "CODEX_HOME": CODEX_HOME},
    )
    requests = (
        {
            "id": 1,
            "method": "initialize",
            "params": {
                "clientInfo": {
                    "name": "atlas-ai-usage-exporter",
                    "title": "Atlas AI Usage Exporter",
                    "version": "1.0.0",
                },
                "capabilities": {"experimentalApi": True},
            },
        },
        {"method": "initialized", "params": {}},
        {"id": 2, "method": "account/rateLimits/read", "params": None},
        {"id": 3, "method": "account/usage/read", "params": None},
    )
    selector: selectors.BaseSelector | None = None
    rate_result: dict[str, Any]
    usage_result: dict[str, Any]
    try:
        if process.stdin is None or process.stdout is None:
            raise RuntimeError("Codex app-server pipes are unavailable")
        for message in requests:
            process.stdin.write(json.dumps(message, separators=(",", ":")) + "\n")
        process.stdin.flush()
        selector = selectors.DefaultSelector()
        selector.register(process.stdout, selectors.EVENT_READ)
        responses: dict[int, dict[str, Any]] = {}
        deadline = time.monotonic() + timeout
        while len(responses) < 3 and time.monotonic() < deadline:
            for key, _ in selector.select(min(1, max(0, deadline - time.monotonic()))):
                line = key.fileobj.readline()
                if not line:
                    continue
                message = json.loads(line)
                if message.get("id") in (1, 2, 3):
                    responses[int(message["id"])] = message
        for response_id in (1, 2, 3):
            response = responses.get(response_id)
            if (
                not response
                or "error" in response
                or not isinstance(response.get("result"), dict)
            ):
                raise RuntimeError(f"Codex app-server response {response_id} failed")
        rate_result = responses[2]["result"]
        usage_result = responses[3]["result"]
    finally:
        cleanup_ok = _cleanup_codex_process(process, selector)
    if not cleanup_ok:
        raise RuntimeError("Codex app-server cleanup failed")
    return rate_result, usage_result