import os
import re
import json
import asyncio
import itertools
from typing import Any, Dict, List, Tuple, Optional

import httpx
from cryptography.fernet import Fernet
from fastapi import FastAPI, Request, Depends, HTTPException
from fastapi.responses import StreamingResponse, Response
from fastapi.middleware.cors import CORSMiddleware

# ----------------------------------------------------------------------------
# Config from secrets
# ----------------------------------------------------------------------------
KEY = os.environ.get("KEY", "").strip()
AES_KEY = os.environ.get("AES_KEY", "").strip()

if not KEY:
    raise RuntimeError("KEY secret is not set")
if not AES_KEY:
    raise RuntimeError("AES_KEY secret is not set")

# ----------------------------------------------------------------------------
# Decrypt models.json at runtime. The plaintext is never on disk.
# ----------------------------------------------------------------------------
_fernet = Fernet(AES_KEY.encode() if isinstance(AES_KEY, str) else AES_KEY)
with open("models.json.enc", "rb") as _f:
    MODELS: Dict[str, Dict[str, Any]] = json.loads(_fernet.decrypt(_f.read()))

# Per-model round-robin state
_RR: Dict[str, itertools.cycle] = {
    name: itertools.cycle(range(len(cfg["targets"])))
    for name, cfg in MODELS.items()
}
_RR_LOCKS: Dict[str, asyncio.Lock] = {name: asyncio.Lock() for name in MODELS}

# Unique (url, key) endpoints across all models, for /health probing
def _unique_endpoints() -> List[Tuple[str, str]]:
    seen = set(); out = []
    for cfg in MODELS.values():
        for t in cfg["targets"]:
            k = (t["url"], t["key"])
            if k not in seen:
                seen.add(k); out.append(k)
    return out

ENDPOINTS = _unique_endpoints()

# ----------------------------------------------------------------------------
# Auth -- accept the key from any header a harness might send:
# Authorization: Bearer KEY (OpenAI SDK, LiteLLM, ...)
# x-api-key: KEY                  (Anthropic SDK, Claude Code)
# api-key: KEY                    (Azure-style)
# x-goog-api-key: KEY             (Google AI SDK)
# ----------------------------------------------------------------------------
_AUTH_HEADERS = ("authorization", "x-api-key", "api-key", "x-goog-api-key")


def _provided_key(request: Request) -> Optional[str]:
    for h in _AUTH_HEADERS:
        v = request.headers.get(h)
        if not v:
            continue
        v = v.strip()
        if v.lower().startswith("bearer "):
            v = v[7:].strip()
        if v:
            return v
    # allow ?api_key= / ?key= query too
    return request.query_params.get("api_key") or request.query_params.get("key")


async def require_auth(request: Request):
    if _provided_key(request) == KEY:
        return True
    raise HTTPException(status_code=401, detail="invalid api key")

# ----------------------------------------------------------------------------
# App -- docs/redoc/openapi fully disabled and explicitly blocked
# ----------------------------------------------------------------------------
app = FastAPI(
    title="router",
    docs_url=None,
    redoc_url=None,
    openapi_url=None,
)

# CORS -- browser harnesses (Continue, web UIs, etc.) send preflight OPTIONS.
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["*"],
)

@app.get("/docs", include_in_schema=False)
@app.get("/redoc", include_in_schema=False)
@app.get("/openapi.json", include_in_schema=False)
@app.get("/docs.json", include_in_schema=False)
async def _block_docs():
    raise HTTPException(status_code=404)

# ----------------------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------------------
# ----------------------------------------------------------------------------
# Cooldown / circuit-breaker: when a target fails with a retryable error,
# it is banned for BAN_TTL seconds. _pick_target skips banned targets and
# only falls back to them if every target is currently banned (so a cold
# start or all-banned state still serves the request).
# ----------------------------------------------------------------------------
import time as _time
_BAN: Dict[Tuple[str, str], float] = {}
# Tiered bans: 1113 = dead account (rarely recovers); 429 = transient;
# connect/5xx = probably transient. Long bans stop the prober from
# re-testing confirmed-dead keys every cycle.
_BAN_TTL_LONG = 3600.0    # 1113 insufficient balance / 401 / 403
_BAN_TTL_MED = 120.0      # connect error / 5xx
_BAN_TTL_SHORT = 20.0     # 429 rate limit (recovers fast)
_PROBE_INTERVAL = 600.0  # re-probe every 10 min (cheaper than 4)
_PROBE_CONCURRENCY = 40   # fast first-cycle classification (~15s for 683)


# Shared connection pool -- ONE persistent client reused by every request
# and by the prober. This eliminates the TLS handshake cost (~200-400ms)
# that a per-request httpx.AsyncClient pays on every retry. Keep-alive keeps
# the TCP+TLS connection to api.z.ai / bigmodel.cn warm.
_CLIENT: Optional[httpx.AsyncClient] = None


async def _client() -> httpx.AsyncClient:
    global _CLIENT
    if _CLIENT is None or _CLIENT.is_closed:
        _CLIENT = httpx.AsyncClient(
            timeout=httpx.Timeout(600.0, connect=10.0, read=600.0),
            limits=httpx.Limits(
                max_connections=300,
                max_keepalive_connections=80,
                keepalive_expiry=300.0,
            ),
        )
    return _CLIENT


def _ban(target, ttl=_BAN_TTL_MED):
    _BAN[(target["url"], target["key"])] = _time.monotonic() + ttl


def _is_banned(target) -> bool:
    exp = _BAN.get((target["url"], target["key"]))
    return exp is not None and exp > _time.monotonic()


# ----------------------------------------------------------------------------
# Background balance prober. /models lies -- it returns 200 even for drained
# keys. Only a real 1-token chat reveals the 1113 (insufficient balance)
# error. This task continuously classifies every key, banning drained ones
# and keeping the good pool hot so requests hit a good key on the first try.
# Cost: ~1 token per good key per cycle (drained keys error out at 0 cost).
# ----------------------------------------------------------------------------
async def _probe_balance(client: httpx.AsyncClient, url: str, key: str) -> bool:
    try:
        r = await client.post(
            url.rstrip("/") + "/chat/completions",
            headers={"Authorization": "Bearer " + key, "Content-Type": "application/json"},
            json={"model": "glm-4.6", "messages": [{"role": "user", "content": "."}], "max_tokens": 1},
            timeout=15.0,
        )
    except Exception:
        return False
    if r.status_code != 200:
        return False
    return not _is_retryable_error(r.content)


async def _classify_all():
    """One-time bootstrap: probe EVERY key once to populate the ban table fast
    (~15s at concurrency 40) so the first real request doesn't walk 640 dead
    keys. Good keys may briefly 429 but recover within a minute."""
    sem = asyncio.Semaphore(_PROBE_CONCURRENCY)
    c = await _client()
    async def one(u, k):
        async with sem:
            ok = await _probe_balance(c, u, k)
        if ok:
            _BAN.pop((u, k), None)
        else:
            _BAN[(u, k)] = _time.monotonic() + _BAN_TTL_LONG
    await asyncio.gather(*[one(u, k) for u, k in ENDPOINTS])


async def _probe_all():
    """Re-probe BANNED keys whose ban is about to expire, to catch recoveries.
    Working/un-banned keys are never probed -- live traffic validates them and
    the prober would just 429 them. This is a recovery sweep, not a full scan."""
    sem = asyncio.Semaphore(_PROBE_CONCURRENCY)
    c = await _client()
    now = _time.monotonic()
    # Keys banned within the next PROBE_INTERVAL + buffer are candidates.
    horizon = now + _PROBE_INTERVAL + 60
    todo = [(u, k) for (u, k) in ENDPOINTS
            if _BAN.get((u, k), 0) <= horizon and _BAN.get((u, k), 0) > now]

    async def one(u, k):
        async with sem:
            ok = await _probe_balance(c, u, k)
        if ok:
            _BAN.pop((u, k), None)   # recovered -> available again
        else:
            _BAN[(u, k)] = _time.monotonic() + _BAN_TTL_LONG
    if todo:
        await asyncio.gather(*[one(u, k) for u, k in todo])


async def _prober_loop():
    # Bootstrap: classify all keys once so the ban table is populated.
    try:
        await _classify_all()
    except Exception:
        pass
    # Steady state: only sweep recovering banned keys, never touch working
    # ones (live traffic validates those).
    while True:
        await asyncio.sleep(_PROBE_INTERVAL)
        try:
            await _probe_all()
        except Exception:
            pass


@app.on_event("startup")
async def _start_prober():
    await _client()   # warm the connection pool
    asyncio.create_task(_prober_loop())


@app.on_event("shutdown")
async def _close_client():
    global _CLIENT
    if _CLIENT is not None:
        await _CLIENT.aclose()
        _CLIENT = None


async def _pick_target(model_name: str):
    """Round-robin next target, skipping banned ones unless all are banned."""
    cfg = MODELS[model_name]
    targets = cfg["targets"]
    n = len(targets)
    async with _RR_LOCKS[model_name]:
        # First pass: find the next non-banned target.
        for _ in range(n):
            idx = next(_RR[model_name])
            t = targets[idx]
            if not _is_banned(t):
                return t
        # Everything is banned -- return the round-robin pick anyway so the
        # request still goes out. Ban timestamps will have started expiring.
        idx = next(_RR[model_name])
        return targets[idx]

# hop-by-hop / router-only headers we never copy upstream or downstream
_HOP = {
    "host", "content-length", "transfer-encoding", "connection",
    "keep-alive", "proxy-authenticate", "proxy-authorization",
    "te", "trailers", "upgrade",
}


def _resolve_model(payload: Optional[dict], request: Request) -> Optional[str]:
    """Find the requested model from JSON body, query, or header."""
    if isinstance(payload, dict):
        m = payload.get("model")
        if isinstance(m, str) and m:
            return m
    m = request.query_params.get("model")
    if m:
        return m
    m = request.headers.get("x-router-model")
    if m:
        return m
    return None

# GLM/Zhipu per-key error codes that mean "this key is unusable, try another":
# 1111/1112 invalid or expired key, 1113 insufficient balance,
# 1114/1115/1116/1117 quota / resource-pack exhausted.
_RETRYABLE_CODES = {"1111", "1112", "1113", "1114", "1115", "1116", "1117", "1261"}


# ----------------------------------------------------------------------------
# Protocol / harness detection
# ----------------------------------------------------------------------------
# GLM exposes two native protocol surfaces:
#   OpenAI-compatible:   <host>/api/paas/v4/<path>   (and /api/coding/paas/v4)
#   Anthropic-compatible: <host>/api/anthropic/<path>
# We detect which harness is calling us from the request path and route to the
# matching upstream base, so OpenAI SDK, Anthropic SDK / Claude Code, LiteLLM,
# LangChain, Cline, Continue, aider, etc. all work with zero translation.

def _detect_format(path: str) -> str:
    p = path.lower().lstrip("/")
    if (
        p.startswith("v1/messages")
        or p == "messages"
        or p.startswith("messages/")
        or "count_tokens" in p
    ):
        return "anthropic"
    if p.startswith("v1beta/") or ":generatecontent" in p or ":streamgeneratecontent" in p:
        return "gemini"
    return "openai"


_ANTR_RE = re.compile(r"^(https?://[^/]+/api)(/.*)?$")


def _target_base(target_url: str, fmt: str) -> str:
    """Map an OpenAI-format target base to the right base for the protocol."""
    base = target_url.rstrip("/")
    if fmt == "anthropic":
        m = _ANTR_RE.match(base)
        if m:
            return m.group(1) + "/anthropic"
        return base
    # openai + gemini (gemini falls back to openai base, best-effort)
    return base


def _forward_path(path: str, fmt: str) -> str:
    """Compute the path segment to append to the upstream base."""
    p = path.lstrip("/")
    if fmt == "openai":
        # base already ends in /v4; drop a leading v1/ from the harness path
        if p.startswith("v1/"):
            p = p[3:]
        return p
    if fmt == "anthropic":
        # upstream anthropic base has no version; the harness's /v1/ stays
        if not p.startswith("v1/"):
            p = "v1/" + p
        return p
    return p  # gemini passthrough


_GEM_MODEL_RE = re.compile(r"/models/([^/:]+)", re.I)


def _extract_model_gemini(path: str) -> Optional[str]:
    m = _GEM_MODEL_RE.search("/" + path)
    return m.group(1) if m else None



def _is_retryable_error(buf: bytes) -> bool:
    """Return True if buf is a GLM error JSON we should retry on another key.

    Handles both protocol shapes:
      OpenAI:    {"error":{"code":"1113","message":"..."}}
      Anthropic: {"type":"error","error":{"type":"...","code":"1113","message":"..."}}
    """
    if not buf or buf[:1] != b"{":
        return False
    try:
        j = json.loads(buf)
    except Exception:
        return False
    if not isinstance(j, dict):
        return False
    # Both shapes carry the GLM code under error.code or top-level code.
    err = j.get("error") if isinstance(j.get("error"), dict) else {}
    code = str(err.get("code") or j.get("code") or "")
    if code in _RETRYABLE_CODES:
        return True
    msg = str(err.get("message") or j.get("message") or "").lower()
    if any(s in msg for s in ("balance", "quota", "余额", "配额", "insufficient")):
        return True
    # Anthropic rate_limit_error / overloaded -> try another key.
    etype = str(err.get("type") or "").lower()
    if etype in ("rate_limit_error", "overloaded_error"):
        return True
    return False

# ----------------------------------------------------------------------------
# Health -- instant. Reads the ban table maintained by the background
# prober. No outbound calls (the old version fired 683 probes per request
# and starved the event loop).
# ----------------------------------------------------------------------------
@app.get("/health")
@app.get("/healthz")
async def health():
    now_mono = _time.monotonic()
    banned = sum(1 for exp in _BAN.values() if exp > now_mono)
    return {
        "total": len(ENDPOINTS),
        "available": len(ENDPOINTS) - banned,
        "banned": banned,
        "models": len(MODELS),
    }

# ----------------------------------------------------------------------------
# Model list
# ----------------------------------------------------------------------------
@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [
            {"id": name, "object": "model", "owned_by": "router"}
            for name in MODELS
        ],
    }

@app.get("/")
async def root():
    now_mono = _time.monotonic()
    banned = sum(1 for exp in _BAN.values() if exp > now_mono)
    return {
        "ok": True,
        "models": len(MODELS),
        "endpoints": len(ENDPOINTS),
        "available": len(ENDPOINTS) - banned,
        "banned": banned,
    }

# ----------------------------------------------------------------------------
# Transparent catch-all proxy. Mounted under /{path:path} so harnesses that
# include /v1/ (OpenAI/Anthropic) AND those that drop it both work. Public
# routes above (/health, /v1/models, /) are matched before this catch-all.
# ----------------------------------------------------------------------------
@app.api_route(
    "/{path:path}",
    methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"],
    dependencies=[Depends(require_auth)],
)
async def proxy(path: str, request: Request):
    fmt = _detect_format(path)

    raw = await request.body()
    ctype = request.headers.get("content-type", "")

    # Parse body only if JSON -- otherwise forward bytes verbatim (multipart,
    # binary, etc.). Tools / messages / stream / thinking all live in JSON
    # and are passed through untouched.
    payload: Optional[dict] = None
    body_to_send: bytes = raw
    is_json = "application/json" in ctype.lower() or ctype == ""
    if is_json and raw:
        try:
            payload = json.loads(raw)
            if not isinstance(payload, dict):
                payload = None
        except Exception:
            payload = None

    # Model resolution works across all harnesses:
    #   OpenAI/Anthropic: "model" in JSON body
    #   Gemini: model in URL path (v1beta/models/<model>:generateContent)
    model_name = _resolve_model(payload, request)
    if model_name is None and fmt == "gemini":
        model_name = _extract_model_gemini(path)
    if not model_name or model_name not in MODELS:
        raise HTTPException(
            status_code=400,
            detail=f"unknown model: {model_name!r}; available: {list(MODELS)}",
        )

    cfg = MODELS[model_name]
    n_targets = len(cfg["targets"])
    last_err = "no targets"
    fwd_rel = _forward_path(path, fmt)

    for _ in range(n_targets):
        target = await _pick_target(model_name)
        base = _target_base(target["url"], fmt)

        # Build outbound body: swap ONLY the model field when JSON. For Gemini
        # the model lives in the URL path; since upstream_model == model here
        # we forward the path verbatim.
        if isinstance(payload, dict):
            out = dict(payload)
            out["model"] = target["upstream_model"]
            send_bytes = json.dumps(out, ensure_ascii=False).encode()
        else:
            send_bytes = body_to_send

        url = base.rstrip("/") + "/" + fwd_rel.lstrip("/")

        # Forward every client header transparently except hop-by-hop + the
        # auth headers we must override for the upstream.
        fwd = {
            k: v for k, v in request.headers.items()
            if k.lower() not in _HOP
            and k.lower() not in _AUTH_HEADERS
        }
        # Anthropic upstream wants x-api-key, OpenAI/Gemini want Bearer.
        if fmt == "anthropic":
            fwd["x-api-key"] = target["key"]
        else:
            fwd["Authorization"] = "Bearer " + target["key"]
        if send_bytes:
            fwd["content-length"] = str(len(send_bytes))
        if "content-type" not in {k.lower() for k in fwd} and ctype:
            fwd["Content-Type"] = ctype

        query = str(request.url.query)

        client = await _client()
        try:
            req = client.build_request(
                request.method, url,
                headers=fwd,
                content=send_bytes if send_bytes else None,
                params=query or None,
            )
            resp = await client.send(req, stream=True)
        except Exception as e:
            _ban(target, _BAN_TTL_MED)
            last_err = f"connect: {e}"
            continue

        # Classify the upstream response. The key insight: GLM serves
        # balance-drained errors (code 1113) with HTTP 429 status, NOT a real
        # rate-limit. So we must inspect the body before choosing a ban TTL:
        #   1113/quota in body -> dead account, LONG ban (don't re-walk it)
        #   genuine 429         -> transient, SHORT ban
        #   5xx / connect err   -> transient, MED ban
        #   401/403             -> dead auth, LONG ban
        ctype_resp = resp.headers.get("content-type", "")
        is_sse = ctype_resp.lower().startswith("text/event-stream")

        if resp.status_code >= 400 or ("application/json" in ctype_resp.lower() and not is_sse):
            body_buf = await resp.aread()
            await resp.aclose()
            retryable = _is_retryable_error(body_buf)
            if retryable:
                # 1113 balance / quota / drained -> dead key, don't revisit
                _ban(target, _BAN_TTL_LONG)
                last_err = f"drained: {body_buf[:120]!r}"
                continue
            if resp.status_code == 429:
                # genuine rate limit (no 1113) -> recovers fast
                _ban(target, _BAN_TTL_SHORT)
                last_err = f"rate-limited: {body_buf[:120]!r}"
                continue
            if resp.status_code in (401, 403):
                _ban(target, _BAN_TTL_LONG)
                last_err = f"auth {resp.status_code}: {body_buf[:120]!r}"
                continue
            if resp.status_code >= 500:
                _ban(target, _BAN_TTL_MED)
                last_err = f"upstream {resp.status_code}: {body_buf[:120]!r}"
                continue
            # Non-retryable 4xx (bad request, unknown model, ...) -> return
            resp_headers = {
                k: v for k, v in resp.headers.items()
                if k.lower() not in _HOP and k.lower() != "content-encoding"
            }
            return Response(
                content=body_buf,
                status_code=resp.status_code,
                headers=resp_headers,
                media_type=resp.headers.get("content-type"),
            )

        # Pass through response headers (content-type, x-request-id,
        # anthropic-*, caching headers, etc.) untouched.
        resp_headers = {
            k: v for k, v in resp.headers.items()
            if k.lower() not in _HOP
        }

        async def gen():
            try:
                async for chunk in resp.aiter_raw():
                    if chunk:
                        yield chunk
            finally:
                await resp.aclose()

        return StreamingResponse(
            gen(),
            status_code=resp.status_code,
            headers=resp_headers,
            media_type=resp.headers.get("content-type"),
        )

    raise HTTPException(status_code=502, detail=f"all targets failed: {last_err}")