#!/usr/bin/env python3
"""
Extended live smoke test against the production Manager. Exercises:

  A. Routing surface
     - /api/status              → 200 healthy
     - /api/devices/            → 401 (auth-protected, route mounted)
     - /api/devices/pair        → 401 without auth
     - /api/devices/12345       PATCH/DELETE → 401 without auth
     - Web SPA index            → 200 (assets loadable)
     - /v1/models with no auth  → 401 (legacy v1 path still gated)

  B. V2 diagnostic header contract
     Exercises FIVE distinct V2 failure modes and asserts that each
     returns:
       - HTTP 401
       - X-Heicode-Server-Time present, parseable, within 60s of local
       - X-Heicode-Auth-Error = expected machine-readable code
     This pins the whole error-code mapping we shipped in fix(devices).

  C. Legacy compat
     - V1 sk- request to /v1/models without bearer  → 401
     - V1 sk- request with bogus bearer "sk-xxx"    → 401
     Confirms the 30-day legacy compat path still runs (V2 dispatcher
     isn't accidentally swallowing non-V2 traffic).

Exit code: 0 == all green, 1 == any failure.
"""

import sys
import time
import urllib.request
import urllib.error

BASE = "https://code.xinghanlab.com"
UA = "heicode-smoke-extended/1.0"

FAILS = []

# Honour HTTPS_PROXY env (clash/wireguard/etc) — urllib doesn't pick
# it up automatically on Windows. Needed when the local machine has a
# DNS hijack returning sinkhole IPs for code.xinghanlab.com.
import os
_proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy")
if _proxy:
    proxy_handler = urllib.request.ProxyHandler({"https": _proxy, "http": _proxy})
    _opener = urllib.request.build_opener(proxy_handler)
else:
    _opener = urllib.request.build_opener()


def http(method, path, headers=None, body=None):
    req = urllib.request.Request(BASE + path, method=method, data=body)
    req.add_header("User-Agent", UA)
    for k, v in (headers or {}).items():
        req.add_header(k, v)
    try:
        r = _opener.open(req, timeout=15)
        return r.status, {k.lower(): v for k, v in r.headers.items()}, r.read().decode("utf-8", "replace")
    except urllib.error.HTTPError as e:
        return e.code, {k.lower(): v for k, v in e.headers.items()}, e.read().decode("utf-8", "replace")


def check(label, ok, detail=""):
    print(f"[{'OK  ' if ok else 'FAIL'}] {label}" + (f"  -- {detail}" if detail else ""))
    if not ok:
        FAILS.append(label)


import base64

# Exactly 32 bytes (0x41 * 32) — valid X25519 pubkey LENGTH (every
# 32-byte sequence is a valid X25519 point per spec). Using "wrong
# byte count" pubkey strings here would short-circuit the dispatcher
# at eph_pubkey_malformed before we can test the downstream code
# paths, so always start from a 32-byte default.
VALID_LEN_EPH_PUBKEY = base64.standard_b64encode(b"\x41" * 32).decode()


def v2_headers(eph_pubkey_b64=VALID_LEN_EPH_PUBKEY,
               timestamp_ms=None, nonce="dead" * 8):
    return {
        "Content-Encoding": "heicode-aead-v1",
        "X-Heicode-Device-Id": "smoke-bogus",
        "X-Heicode-Timestamp": str(timestamp_ms or int(time.time() * 1000)),
        "X-Heicode-Nonce": nonce,
        "X-Heicode-Fingerprint": "0" * 64,
        "X-Heicode-Eph-Pubkey": eph_pubkey_b64,
        "X-Heicode-Signature": "BBBBBBBB",
        "Content-Type": "application/json",
    }


def assert_v2_failure(label, headers, body, expected_code):
    code, hdrs, _ = http("POST", "/v1/messages", headers=headers, body=body)
    diag = hdrs.get("x-heicode-auth-error", "")
    ts = hdrs.get("x-heicode-server-time", "")
    ts_ok = False
    try:
        ts_ok = ts and abs(int(ts) - int(time.time() * 1000)) < 60_000
    except ValueError:
        pass
    ok = code == 401 and diag == expected_code and ts_ok
    check(label, ok, f"HTTP {code} code={diag!r} time={ts!r}")


# ---------------------------------------------------------------- A. Routing
print(f"\n=== A. Routing surface ({BASE}) ===\n")

code, _, body = http("GET", "/api/status")
check("/api/status returns 200 + success:true",
      code == 200 and '"success":true' in body, f"HTTP {code}")

code, _, _ = http("GET", "/api/devices/")
check("/api/devices/ unauthenticated → 401", code == 401, f"HTTP {code}")

code, _, _ = http("POST", "/api/devices/pair",
                  headers={"Content-Type": "application/json"},
                  body=b'{"device_id":"x","public_key":"","fingerprint":""}')
check("/api/devices/pair unauthenticated → 401", code == 401, f"HTTP {code}")

code, _, _ = http("PATCH", "/api/devices/99999",
                  headers={"Content-Type": "application/json"},
                  body=b'{"device_name":"hax"}')
check("/api/devices/:id PATCH unauthenticated → 401", code == 401, f"HTTP {code}")

code, _, _ = http("DELETE", "/api/devices/99999")
check("/api/devices/:id DELETE unauthenticated → 401", code == 401, f"HTTP {code}")

code, _, body = http("GET", "/")
check("Web SPA index → 200", code == 200 and "<html" in body.lower(),
      f"HTTP {code} len={len(body)}")

code, _, _ = http("GET", "/v1/models")
check("/v1/models without auth → 401", code == 401, f"HTTP {code}")

# ---------------------------------------------------------- B. V2 diag codes
print("\n=== B. V2 diagnostic header contract ===\n")

# B1: malformed ephemeral pubkey (not base64 → eph_pubkey_malformed)
assert_v2_failure(
    "V2 malformed eph_pubkey → eph_pubkey_malformed",
    {**v2_headers(eph_pubkey_b64="not_base64_!!!"), },
    b"\x00" * 12,
    "eph_pubkey_malformed",
)

# B2: pubkey base64-decodes but wrong length (not 32 bytes) → eph_pubkey_malformed
assert_v2_failure(
    "V2 wrong-length eph_pubkey → eph_pubkey_malformed",
    {**v2_headers(eph_pubkey_b64="QUE=")},  # decodes to 2 bytes
    b"\x00" * 12,
    "eph_pubkey_malformed",
)

# B3: body too short (<28 bytes for nonce+tag) → body_too_short
assert_v2_failure(
    "V2 body too short → body_too_short",
    v2_headers(),
    b"\x00" * 5,  # too short
    "body_too_short",
)

# B4: missing X-Heicode-Eph-Pubkey header entirely → eph_pubkey_missing
hdrs_no_eph = v2_headers()
hdrs_no_eph.pop("X-Heicode-Eph-Pubkey")
assert_v2_failure(
    "V2 missing eph_pubkey header → eph_pubkey_missing",
    hdrs_no_eph,
    b"\x00" * 12,
    "eph_pubkey_missing",
)

# B5: valid-shape but unknown device + decoy ciphertext → after decrypt
#     gets `body_decrypt_failed` (AEAD tag mismatch). This is the most
#     common real-world failure mode for a tampered request.
assert_v2_failure(
    "V2 valid-shape but undecryptable → body_decrypt_failed",
    v2_headers(),  # 32-byte AAAAAA... pubkey is valid X25519
    b"\xAA" * 28,  # 12 nonce + 16 tag, but ciphertext won't authenticate
    "body_decrypt_failed",
)

# ---------------------------------------------------------- C. Legacy compat
print("\n=== C. Legacy sk- bearer compat path ===\n")

# No Content-Encoding header → legacy dispatcher → ValidateUserToken
code, hdrs, _ = http("POST", "/v1/messages",
                     headers={"Content-Type": "application/json",
                              "Authorization": "Bearer sk-bogus-token-not-real"},
                     body=b'{"model":"claude-haiku","messages":[]}')
# V2 diagnostic headers MUST NOT be present on legacy path responses.
diag = hdrs.get("x-heicode-auth-error", "")
check("Legacy sk- request → 401 (no Content-Encoding)", code == 401, f"HTTP {code}")
check("Legacy path does NOT leak V2 diagnostic headers", diag == "",
      f"x-heicode-auth-error={diag!r}")

# ---------------------------------------------------------- summary
if FAILS:
    print(f"\n{len(FAILS)} smoke check(s) FAILED:")
    for f in FAILS:
        print(f"  - {f}")
    sys.exit(1)

print("\nALL EXTENDED SMOKE CHECKS PASSED")