"""ChatGPT / Codex API Gateway — thin wrapper around gateway_base.

Accepts plaintext HTTP from sandboxes on the internal Docker network,
validates container identity, injects OAuth credentials, and forwards to
https://chatgpt.com over HTTPS.  Streams responses back chunk-by-chunk
without buffering.

Port allocation:
  :9848  Anthropic gateway (gateway.py)
  :9849  OpenAI gateway    (openai_gateway.py)
  :9850  GitHub gateway     (github_gateway.py)
  :9851  Gemini gateway     (gemini_gateway.py)
  :9852  ChatGPT gateway    (this module, HTTP)
  :443   ChatGPT gateway    (this module, TLS — transparent interception)

Credential injection:
  Uses OAuth tokens from Codex CLI auth.json (mounted at CODEX_AUTH_FILE).
  Tokens are refreshed automatically via OAuthTokenManager when expired.

Transparent TLS interception:
  The Codex CLI in subscription mode hardcodes requests to chatgpt.com and
  ignores base URL env vars.  To intercept this traffic without MITM:

  1. /etc/hosts in the sandbox resolves chatgpt.com → unified-proxy IP
  2. NO_PROXY ensures the request bypasses Squid (direct to gateway)
  3. The gateway terminates TLS on port 443 using a cert signed by the
     mitmproxy CA (trusted by the sandbox via the combined CA bundle)
  4. The gateway injects OAuth credentials and forwards to real chatgpt.com

  The TLS cert is generated by entrypoint.sh (generate_chatgpt_gateway_cert)
  only when CODEX_AUTH_FILE is set and MITM is enabled (mitmproxy CA exists).
  When the cert is absent, only the HTTP listener on port 9852 starts.
"""

import asyncio
import os
import signal
import ssl
import sys
from typing import Optional

from aiohttp import web

# ---------------------------------------------------------------------------
# Ensure /opt/proxy is on sys.path so imports work inside the container
# ---------------------------------------------------------------------------
_PROXY_DIR = "/opt/proxy"
if _PROXY_DIR not in sys.path:
    sys.path.insert(0, _PROXY_DIR)

from gateway_base import create_gateway_app, gateway_error  # noqa: E402
from logging_config import get_logger, setup_logging  # noqa: E402
from registry import ContainerRegistry  # noqa: E402

logger = get_logger(__name__)

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

CHATGPT_GATEWAY_PORT = int(os.environ.get("CHATGPT_GATEWAY_PORT", "9852"))
CHATGPT_GATEWAY_BIND = os.environ.get("CHATGPT_GATEWAY_BIND", "0.0.0.0")

# TLS listener for transparent interception (chatgpt.com via /etc/hosts)
CHATGPT_TLS_PORT = int(os.environ.get("CHATGPT_TLS_PORT", "443"))
CHATGPT_TLS_CERT = os.environ.get("CHATGPT_TLS_CERT", "/etc/proxy/certs/chatgpt-gw.pem")
CHATGPT_TLS_KEY = os.environ.get("CHATGPT_TLS_KEY", "/etc/proxy/certs/chatgpt-gw.key")


# ---------------------------------------------------------------------------
# OAuth token manager (module-level so request_hook can access it)
# ---------------------------------------------------------------------------

_token_manager = None


def _init_token_manager() -> None:
    """Initialize the OAuth token manager from CODEX_AUTH_FILE."""
    global _token_manager

    auth_file = os.environ.get("CODEX_AUTH_FILE", "").strip()
    if not auth_file:
        logger.warning("chatgpt-gateway: CODEX_AUTH_FILE not set")
        return

    if not os.path.exists(auth_file):
        logger.warning(f"chatgpt-gateway: auth file not found: {auth_file}")
        return

    try:
        # Import from the addons oauth_managers package (available in container)
        from addons.oauth_managers.codex import OAuthTokenManager
        _token_manager = OAuthTokenManager(auth_file)
        logger.info(f"chatgpt-gateway: OAuth token manager initialized from {auth_file}")
    except ImportError:
        # Fallback: try the top-level module (codex_token_manager.py)
        try:
            from codex_token_manager import OAuthTokenManager as FallbackManager
            _token_manager = FallbackManager(auth_file)
            logger.info(f"chatgpt-gateway: OAuth token manager initialized (fallback) from {auth_file}")
        except ImportError:
            logger.error("chatgpt-gateway: OAuthTokenManager not available")
    except (FileNotFoundError, ValueError) as e:
        logger.warning(f"chatgpt-gateway: failed to load auth file: {e}")
    except Exception as e:
        logger.error(f"chatgpt-gateway: failed to initialize token manager: {e}")


# ---------------------------------------------------------------------------
# Credential loading
# ---------------------------------------------------------------------------

def _load_chatgpt_credential() -> Optional[dict]:
    """Load the ChatGPT/Codex OAuth credential.

    Reads the current access token from the OAuthTokenManager.
    Returns dict with 'header' and 'value' keys, or None.
    """
    if _token_manager is None:
        return None

    try:
        token = _token_manager.get_valid_token()
        return {"header": "Authorization", "value": f"Bearer {token}"}
    except Exception as e:
        logger.error(f"chatgpt-gateway: failed to get initial token: {e}")
        return None


# ---------------------------------------------------------------------------
# Request hook — refresh token before each request
# ---------------------------------------------------------------------------

async def _chatgpt_request_hook(
    request: web.Request,
    method: str,
    body: bytes,
    container_id: str,
) -> Optional[web.Response]:
    """Refresh OAuth token before forwarding to upstream.

    The gateway_base credential is loaded once at startup, but OAuth tokens
    expire.  This hook mutates the credential dict in-place so the shared
    proxy handler sees a fresh token on every request.
    """
    if _token_manager is None:
        return gateway_error(502, "ChatGPT OAuth credential not configured")

    try:
        # Run synchronous token refresh in a thread to avoid blocking the
        # event loop — the Codex OAuthTokenManager may do synchronous HTTP
        # via httpx when refreshing expired tokens.
        token = await asyncio.to_thread(_token_manager.get_valid_token)
        # Store refreshed credential on the app dict. This is safe in the
        # single-process asyncio model because the dict assignment is atomic
        # within the event loop (no await between read and write).
        new_cred = {
            "header": "Authorization",
            "value": f"Bearer {token}",
        }
        request.app["credential"] = new_cred
    except Exception as e:
        logger.error(f"chatgpt-gateway: OAuth token refresh failed: {e}")
        return gateway_error(502, "ChatGPT OAuth token refresh failed")

    return None  # Continue to upstream forwarding


# ---------------------------------------------------------------------------
# Application factory
# ---------------------------------------------------------------------------

def create_chatgpt_gateway_app(
    registry: Optional[ContainerRegistry] = None,
) -> web.Application:
    """Create the aiohttp application for the ChatGPT/Codex gateway."""
    _init_token_manager()

    return create_gateway_app(
        upstream_base_url="https://chatgpt.com",
        upstream_host="chatgpt.com",
        service_name="chatgpt-gateway",
        credential_loader=_load_chatgpt_credential,
        routes=[
            ("*", "/{path_info:.*}"),
        ],
        port=CHATGPT_GATEWAY_PORT,
        bind=CHATGPT_GATEWAY_BIND,
        max_request_body=32 * 1024 * 1024,
        connect_timeout=30,
        total_timeout=600,
        credential_required=True,
        request_hook=_chatgpt_request_hook,
        registry=registry,
    )


def run_chatgpt_gateway_server(
    host: str = CHATGPT_GATEWAY_BIND,
    port: int = CHATGPT_GATEWAY_PORT,
    app: Optional[web.Application] = None,
    tls_port: int = CHATGPT_TLS_PORT,
    tls_cert: str = CHATGPT_TLS_CERT,
    tls_key: str = CHATGPT_TLS_KEY,
) -> None:
    """Run the ChatGPT/Codex gateway server (blocking).

    Starts an HTTP listener on *port* (default 9852) and, when TLS cert
    files exist, an HTTPS listener on *tls_port* (default 443) for
    transparent interception of chatgpt.com traffic.
    """
    if app is None:
        app = create_chatgpt_gateway_app()

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    runner = web.AppRunner(app)
    loop.run_until_complete(runner.setup())

    # HTTP listener (port 9852)
    http_site = web.TCPSite(runner, host, port)
    loop.run_until_complete(http_site.start())
    logger.info(f"chatgpt-gateway: HTTP on {host}:{port}")

    # TLS listener (port 443) — only if cert exists
    if os.path.isfile(tls_cert) and os.path.isfile(tls_key):
        try:
            ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
            ctx.minimum_version = ssl.TLSVersion.TLSv1_2
            ctx.set_ciphers("ECDHE+AESGCM:ECDHE+CHACHA20:DHE+AESGCM")
            ctx.options |= ssl.OP_NO_COMPRESSION
            ctx.load_cert_chain(tls_cert, tls_key)
            tls_site = web.TCPSite(runner, host, tls_port, ssl_context=ctx)
            loop.run_until_complete(tls_site.start())
            logger.info(f"chatgpt-gateway: TLS on {host}:{tls_port}")
        except (ssl.SSLError, OSError) as exc:
            logger.error(f"chatgpt-gateway: failed to load TLS cert/key: {exc}")
            logger.warning("chatgpt-gateway: HTTPS listener disabled due to cert error")
    else:
        logger.warning("chatgpt-gateway: TLS cert not found, HTTPS listener disabled")

    # Register signal handlers for graceful shutdown (matching web.run_app behavior)
    shutdown_event = asyncio.Event()

    def _shutdown_handler() -> None:
        logger.info("chatgpt-gateway: received shutdown signal")
        shutdown_event.set()

    for sig in (signal.SIGINT, signal.SIGTERM):
        loop.add_signal_handler(sig, _shutdown_handler)

    try:
        loop.run_until_complete(shutdown_event.wait())
    finally:
        loop.run_until_complete(runner.cleanup())
        loop.close()


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    setup_logging()
    run_chatgpt_gateway_server()
