#!/usr/bin/env python3
"""Stable stdio entrypoint for Codex image tasks.

Codex owns the outer MCP connection. This process must therefore stay alive
even when the implementation returns a tool error or is updated on disk.
The implementation module is reloaded between requests when its mtime changes.
"""

from __future__ import annotations

import importlib.util
import json
import sys
import traceback
from pathlib import Path
from typing import Any


MCP_ROOT = Path(__file__).resolve().parent
if str(MCP_ROOT) not in sys.path:
    sys.path.insert(0, str(MCP_ROOT))
AGENT_PATH = MCP_ROOT / 'yoho_image_agent_mcp.py'
MODULE_NAME = 'yoho_image_agent_runtime'
WIRE_MODE = 'line'
_agent: Any = None
_agent_mtime_ns: int | None = None


def send(payload: dict[str, Any]) -> None:
    body = json.dumps(payload, ensure_ascii=False, separators=(',', ':')).encode('utf-8')
    if WIRE_MODE == 'framed':
        sys.stdout.buffer.write(f'Content-Length: {len(body)}\r\n\r\n'.encode('ascii'))
        sys.stdout.buffer.write(body)
    else:
        sys.stdout.buffer.write(body + b'\n')
    sys.stdout.buffer.flush()


def receive() -> dict[str, Any] | None:
    global WIRE_MODE
    while True:
        line = sys.stdin.buffer.readline()
        if not line:
            return None
        raw = line.strip()
        if not raw:
            continue
        if raw.lower().startswith(b'content-length:'):
            try:
                length = int(raw.split(b':', 1)[1].strip())
            except (IndexError, ValueError) as exc:
                raise ValueError('Invalid Content-Length header') from exc
            while True:
                header = sys.stdin.buffer.readline()
                if not header:
                    return None
                if header in {b'\r\n', b'\n'}:
                    break
            body = sys.stdin.buffer.read(length)
            if len(body) != length:
                raise ValueError('MCP frame ended before Content-Length bytes were read')
            WIRE_MODE = 'framed'
            value = json.loads(body.decode('utf-8-sig'))
        else:
            value = json.loads(raw.decode('utf-8-sig'))
        if not isinstance(value, dict):
            raise ValueError('MCP message must be a JSON object')
        return value


def load_agent() -> Any:
    global _agent, _agent_mtime_ns
    current_mtime_ns = AGENT_PATH.stat().st_mtime_ns
    if _agent is not None and _agent_mtime_ns == current_mtime_ns:
        return _agent
    spec = importlib.util.spec_from_file_location(MODULE_NAME, AGENT_PATH)
    if spec is None or spec.loader is None:
        raise RuntimeError(f'Unable to load image implementation: {AGENT_PATH}')
    candidate = importlib.util.module_from_spec(spec)
    try:
        spec.loader.exec_module(candidate)
    except Exception:
        # Keep serving with the last known-good module during an incomplete or
        # invalid update. The next request retries loading the new file.
        if _agent is not None:
            print(traceback.format_exc(), file=sys.stderr, flush=True)
            return _agent
        raise
    if not callable(getattr(candidate, 'handle_message', None)):
        raise RuntimeError('Image implementation does not expose handle_message')
    _agent = candidate
    _agent_mtime_ns = current_mtime_ns
    return _agent


def main() -> None:
    while True:
        request_id: Any = None
        try:
            message = receive()
            if message is None:
                return
            request_id = message.get('id')
            response = load_agent().handle_message(message)
            if response is not None:
                send(response)
        except ValueError as exc:
            print(f'MCP input error: {exc}', file=sys.stderr, flush=True)
            if request_id is not None:
                send({'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32600, 'message': str(exc)}})
        except Exception as exc:  # noqa: BLE001 - the bridge must never drop Codex's transport
            print(traceback.format_exc(), file=sys.stderr, flush=True)
            if request_id is not None:
                send({'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32603, 'message': f'Image bridge error: {type(exc).__name__}: {exc}'}})


if __name__ == '__main__':
    main()
