#!/usr/bin/env python3
"""ARC Agent Client V1: strict stdlib client for the synthetic pilot relay.

It intentionally supports only the fixed ARC API, HTTPS endpoints under an
approved pilot suffix, globally-routable literal edge addresses, and a local
mode-0600 session file.  It never prints an enrollment code or session key.
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import hmac
import http.client
import ipaddress
import json
import os
import re
import secrets
import socket
import ssl
import stat
import sys
import time
from pathlib import Path
from typing import Any
from urllib.parse import urlsplit

SCHEMA = "ARC_AGENT_CLIENT_V1"
ENROLLMENT_FILE = Path(".arc-enrollment-v1.json")
SESSION_FILE = Path(".arc-session-v1.json")
MAX_JSON_BYTES = 2 * 1024 * 1024
ID = re.compile(r"[A-Za-z0-9_.:-]{1,64}\Z")
HEX32 = re.compile(r"[0-9a-f]{32}\Z")
HEX64 = re.compile(r"[0-9a-f]{64}\Z")
B64U43 = re.compile(r"[A-Za-z0-9_-]{43}\Z")
HOST_LABEL = r"[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?"
QUICK_HOST = re.compile(rf"{HOST_LABEL}\.trycloudflare\.com\Z")
TAILNET_HOST = re.compile(rf"{HOST_LABEL}\.{HOST_LABEL}\.ts\.net\Z")
ROOMS = {"ENTRY", "BENCH", "COMMONS", "JOB", "ENGINEERING", "REVIEW", "APPROVAL", "INCIDENT"}
ALLOWED = (
    re.compile(r"GET /v1/room-manifests/(?:ENTRY|BENCH|COMMONS|JOB|ENGINEERING|REVIEW|APPROVAL|INCIDENT)\Z"),
    re.compile(r"POST /v1/(?:room-acks|events/poll|messages|workspaces|tests|artifacts|reviews)\Z"),
    re.compile(r"POST /v1/workspaces/[0-9a-f]{32}/objects\Z"),
    re.compile(r"(?:GET|PUT) /v1/workspaces/[0-9a-f]{32}/objects/[0-9a-f]{32}\Z"),
    re.compile(r"GET /v1/artifacts/[0-9a-f]{32}\Z"),
    re.compile(r"GET /v1/changeplans/[A-Za-z0-9_-]{16}\Z"),
)


class ClientError(RuntimeError):
    pass


def fail(message: str) -> None:
    raise ClientError(message)


def b64u(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).decode().rstrip("=")


def canonical(value: Any) -> bytes:
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"), allow_nan=False).encode("utf-8")


def strict_json(data: bytes, label: str) -> dict[str, Any]:
    if len(data) > MAX_JSON_BYTES:
        fail(label + " too large")
    try:
        text = data.decode("utf-8")
        def pairs(items: list[tuple[str, Any]]) -> dict[str, Any]:
            out: dict[str, Any] = {}
            for key, value in items:
                if key in out:
                    fail(label + " duplicate key")
                out[key] = value
            return out
        value = json.loads(text, object_pairs_hook=pairs, parse_constant=lambda _x: fail(label + " invalid number"))
    except (UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise ClientError(label + " invalid JSON") from exc
    if not isinstance(value, dict):
        fail(label + " must be object")
    return value


def regular_private(path: Path, *, maximum: int = 65536) -> bytes:
    try:
        info = path.lstat()
    except FileNotFoundError as exc:
        raise ClientError(f"required private file absent: {path.name}") from exc
    if not stat.S_ISREG(info.st_mode) or info.st_nlink != 1 or info.st_uid != os.geteuid() or stat.S_IMODE(info.st_mode) != 0o600:
        fail(f"unsafe private file metadata: {path.name}")
    if info.st_size > maximum:
        fail(f"private file too large: {path.name}")
    descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
    try:
        chunks: list[bytes] = []
        remaining = maximum + 1
        while remaining:
            part = os.read(descriptor, remaining)
            if not part:
                break
            chunks.append(part)
            remaining -= len(part)
        data = b"".join(chunks)
        if len(data) > maximum:
            fail(f"private file too large: {path.name}")
        return data
    finally:
        os.close(descriptor)


def atomic_private(path: Path, data: bytes) -> None:
    if path.parent != Path("."):
        fail("state path must remain in current workspace")
    temporary = Path(f".{path.name}.{secrets.token_hex(8)}.tmp")
    descriptor = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0), 0o600)
    try:
        with os.fdopen(descriptor, "wb") as handle:
            handle.write(data)
            handle.flush()
            os.fsync(handle.fileno())
        os.chmod(temporary, 0o600)
        if path.exists() or path.is_symlink():
            fail("session file already exists")
        os.replace(temporary, path)
        parent = os.open(".", os.O_RDONLY)
        try:
            os.fsync(parent)
        finally:
            os.close(parent)
    finally:
        try:
            temporary.unlink()
        except FileNotFoundError:
            pass


def endpoint_hostname(endpoint: str) -> str:
    parsed = urlsplit(endpoint)
    if parsed.scheme != "https" or parsed.username is not None or parsed.password is not None:
        fail("endpoint must be credential-free HTTPS")
    try:
        port = parsed.port
    except ValueError as exc:
        raise ClientError("invalid endpoint port") from exc
    host = parsed.hostname
    if not host or host != host.casefold() or port not in (None, 443):
        fail("endpoint host or port rejected")
    if parsed.path not in ("",) or parsed.query or parsed.fragment or endpoint != "https://" + host:
        fail("endpoint must not contain path, query, fragment, or explicit port")
    if QUICK_HOST.fullmatch(host) is None and TAILNET_HOST.fullmatch(host) is None:
        fail("endpoint suffix rejected")
    return host


def global_addresses(host: str) -> list[tuple[int, str]]:
    try:
        rows = socket.getaddrinfo(host, 443, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP)
    except socket.gaierror as exc:
        raise ClientError("endpoint DNS unavailable") from exc
    found: list[tuple[int, str]] = []
    for family, socktype, proto, _canon, sockaddr in rows:
        if family not in (socket.AF_INET, socket.AF_INET6) or socktype != socket.SOCK_STREAM or proto not in (0, socket.IPPROTO_TCP):
            continue
        literal = sockaddr[0]
        try:
            ip = ipaddress.ip_address(literal)
        except ValueError:
            continue
        if not ip.is_global or ip.is_multicast or ip.is_unspecified or ip.is_loopback or ip.is_link_local or ip.is_private or ip.is_reserved:
            continue
        pair = (family, ip.compressed)
        if pair not in found:
            found.append(pair)
    if not found:
        fail("endpoint has no globally-routable address")
    return found


class LiteralHTTPSConnection(http.client.HTTPSConnection):
    def __init__(self, hostname: str, family: int, literal: str) -> None:
        super().__init__(hostname, port=443, timeout=20, context=ssl.create_default_context())
        self._family = family
        self._literal = literal

    def connect(self) -> None:
        raw = socket.socket(self._family, socket.SOCK_STREAM, socket.IPPROTO_TCP)
        raw.settimeout(self.timeout)
        try:
            address: tuple[Any, ...] = (self._literal, 443, 0, 0) if self._family == socket.AF_INET6 else (self._literal, 443)
            raw.connect(address)
            self.sock = self._context.wrap_socket(raw, server_hostname=self.host)
        except BaseException:
            raw.close()
            raise


def request(endpoint: str, method: str, path: str, body: bytes, extra_headers: dict[str, str] | None = None) -> tuple[int, dict[str, Any]]:
    host = endpoint_hostname(endpoint)
    if not path.startswith("/v1/") or any(x in path for x in ("?", "#", "%")) or len(path) > 160:
        fail("request target rejected")
    headers = {"Accept": "application/json", "Content-Type": "application/json", "Host": host, "User-Agent": "ARC-Agent-Client/1"}
    if extra_headers:
        headers.update(extra_headers)
    last: BaseException | None = None
    for family, literal in global_addresses(host):
        connection = LiteralHTTPSConnection(host, family, literal)
        try:
            connection.request(method, path, body=body if method in {"POST", "PUT"} else None, headers=headers)
            response = connection.getresponse()
            raw = response.read(MAX_JSON_BYTES + 1)
            if len(raw) > MAX_JSON_BYTES:
                fail("response too large")
            if response.getheader("Content-Type", "").split(";", 1)[0].strip().casefold() != "application/json":
                fail("response media type rejected")
            return response.status, strict_json(raw, "response")
        except (OSError, ssl.SSLError, http.client.HTTPException, ClientError) as exc:
            last = exc
        finally:
            connection.close()
    raise ClientError("HTTPS request failed") from last


def load_session() -> dict[str, Any]:
    value = strict_json(regular_private(SESSION_FILE), "session")
    if set(value) != {"schema", "endpoint", "session_id", "agent_id", "session_key", "expires_at", "slot"}:
        fail("session fields rejected")
    if value["schema"] != "ARC_LOCAL_SESSION_V1" or not isinstance(value["endpoint"], str):
        fail("session schema rejected")
    endpoint_hostname(value["endpoint"])
    if not all(isinstance(value[k], str) and ID.fullmatch(value[k]) for k in ("session_id", "agent_id")):
        fail("session identity rejected")
    if not isinstance(value["session_key"], str) or B64U43.fullmatch(value["session_key"]) is None:
        fail("session key encoding rejected")
    if not isinstance(value["expires_at"], int) or isinstance(value["expires_at"], bool) or value["expires_at"] < int(time.time()):
        fail("session expired")
    if not isinstance(value["slot"], int) or isinstance(value["slot"], bool) or not 1 <= value["slot"] <= 10:
        fail("session slot rejected")
    return value


def signed_headers(session: dict[str, Any], method: str, path: str, body: bytes) -> dict[str, str]:
    timestamp = int(time.time())
    nonce = secrets.token_hex(16)
    body_hash = hashlib.sha256(body).hexdigest()
    fields = ("ARC-HMAC-V1", method, path, "", "application/json", body_hash, str(timestamp), nonce, session["agent_id"], session["session_id"], "")
    canonical_request = "\n".join(fields).encode()
    key = base64.urlsafe_b64decode(session["session_key"] + "=")
    mac = b64u(hmac.new(key, canonical_request, hashlib.sha256).digest())
    return {
        "Authorization": f"ARC-HMAC-V1 {session['session_id']}:{mac}",
        "X-ARC-Agent-ID": session["agent_id"],
        "X-ARC-Timestamp": str(timestamp),
        "X-ARC-Nonce": nonce,
    }


def command_enroll(_args: argparse.Namespace) -> int:
    source = strict_json(regular_private(ENROLLMENT_FILE), "enrollment")
    required = {"schema", "endpoint", "agent_id", "prompt_id", "enrollment_code"}
    if set(source) != required or source.get("schema") != "ARC_ENROLLMENT_INPUT_V1":
        fail("enrollment fields rejected")
    if not all(isinstance(source[k], str) for k in required):
        fail("enrollment values rejected")
    endpoint_hostname(source["endpoint"])
    if ID.fullmatch(source["agent_id"]) is None or ID.fullmatch(source["prompt_id"]) is None or not 16 <= len(source["enrollment_code"]) <= 64:
        fail("enrollment identity rejected")
    body = canonical({k: source[k] for k in ("enrollment_code", "agent_id", "prompt_id")})
    status, response = request(source["endpoint"], "POST", "/v1/session/enroll", body)
    expected = {"schema", "session_id", "agent_id", "session_key", "expires_at", "slot", "pilot_result"}
    if status != 201 or set(response) != expected or response.get("schema") != "ARC_SESSION_CREDENTIAL_V1" or response.get("agent_id") != source["agent_id"] or response.get("pilot_result") != "PROVISIONAL_PILOT":
        fail("enrollment rejected")
    session = {"schema": "ARC_LOCAL_SESSION_V1", "endpoint": source["endpoint"], **{k: response[k] for k in ("session_id", "agent_id", "session_key", "expires_at", "slot")}}
    atomic_private(SESSION_FILE, canonical(session) + b"\n")
    ENROLLMENT_FILE.unlink()
    print(json.dumps({"schema": SCHEMA, "status": "ENROLLED", "agent_id": session["agent_id"], "slot": session["slot"], "expires_at": session["expires_at"], "credential_printed": False}, sort_keys=True, separators=(",", ":")))
    return 0


def command_health(_args: argparse.Namespace) -> int:
    source = strict_json(regular_private(ENROLLMENT_FILE), "enrollment") if ENROLLMENT_FILE.exists() else load_session()
    endpoint = source.get("endpoint")
    if not isinstance(endpoint, str):
        fail("endpoint unavailable")
    status, response = request(endpoint, "GET", "/v1/health", b"")
    if status != 200 or response.get("schema") != "ARC_RELAY_HEALTH_V1" or response.get("status") != "READY":
        fail("relay health not ready")
    print(json.dumps({"schema": SCHEMA, "status": "READY", "server_time": response.get("server_time")}, sort_keys=True, separators=(",", ":")))
    return 0


def command_request(args: argparse.Namespace) -> int:
    method = args.method
    path = args.path
    if not any(rule.fullmatch(method + " " + path) for rule in ALLOWED):
        fail("method/path not allowlisted")
    if method in {"POST", "PUT"}:
        body = regular_private(Path(args.body), maximum=1024 * 1024)
        if method == "POST":
            strict_json(body, "request body")
    else:
        if args.body is not None:
            fail("GET body rejected")
        body = b""
    session = load_session()
    status, response = request(session["endpoint"], method, path, body, signed_headers(session, method, path, body))
    print(json.dumps({"schema": SCHEMA, "http_status": status, "response": response}, ensure_ascii=False, sort_keys=True, separators=(",", ":")))
    return 0 if 200 <= status < 300 else 1


def command_destroy(_args: argparse.Namespace) -> int:
    data = regular_private(SESSION_FILE)
    SESSION_FILE.unlink()
    del data
    print(json.dumps({"schema": SCHEMA, "status": "LOCAL_SESSION_REMOVED"}, sort_keys=True, separators=(",", ":")))
    return 0


def parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="Strict ARC synthetic pilot HTTPS client")
    sub = p.add_subparsers(dest="command", required=True)
    for name, handler in (("health", command_health), ("enroll", command_enroll), ("destroy", command_destroy)):
        q = sub.add_parser(name); q.set_defaults(handler=handler)
    q = sub.add_parser("request")
    q.add_argument("method", choices=("GET", "POST", "PUT"))
    q.add_argument("path")
    q.add_argument("--body")
    q.set_defaults(handler=command_request)
    return p


def main() -> int:
    os.umask(0o077)
    try:
        args = parser().parse_args()
        return int(args.handler(args))
    except ClientError as exc:
        print(json.dumps({"schema": "ARC_AGENT_CLIENT_ERROR_V1", "status": "FAILED_SAFE", "error": str(exc)}, sort_keys=True, separators=(",", ":")), file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())