# ============================================================
# Pi Agent 长期只读凭证（moduleId: auth-agent-tokens, 可重生 ✅）
# 方向 A · P1：/api/agent/* 机器凭证。
#   - HMAC(base64url body + secret) 与 licenses.py/providers.py 同风格；
#   - token 本身只携带 tokenId/kind/scope/时间；principal/tenant 从登记文件解析；
#   - APS_AGENT_TOKEN_SECRET 缺失时签发与鉴权显式失败，绝不回退匿名/本地管理员。
# ============================================================
from __future__ import annotations

import hashlib
import hmac
import json
import os
import threading
import time
import uuid
from base64 import urlsafe_b64decode as _urlsafe_b64decode
from base64 import urlsafe_b64encode as _urlsafe_b64encode
from pathlib import Path
from typing import Any

from server.auth.context import IdentityContext
from server.auth.providers import AuthError

AGENT_TOKEN_VERSION = 1
AGENT_TOKEN_KIND = "agent.read"
AGENT_TOKEN_SCOPE_READ = "read"
AGENT_TOKEN_FILE_NAME = "agent-tokens.json"
AGENT_TOKEN_MAX_TTL_DAYS = 3650


def _b64encode(raw: bytes) -> str:
    return _urlsafe_b64encode(raw).decode("ascii").rstrip("=")


def _b64decode_text(raw: str) -> bytes:
    return _urlsafe_b64decode(raw + "=" * (-len(raw) % 4))


def _configured_secret() -> bytes:
    value = (os.environ.get("APS_AGENT_TOKEN_SECRET") or "").strip()
    if not value:
        raise AuthError(
            "AGENT_TOKEN_NOT_CONFIGURED",
            "Pi Agent 凭证服务未配置：请设置 APS_AGENT_TOKEN_SECRET",
            503,
        )
    return value.encode("utf-8")


def _default_path() -> Path:
    try:
        from server.aps_home import data_dir

        return data_dir() / AGENT_TOKEN_FILE_NAME
    except (OSError, RuntimeError, ImportError):
        return (
            Path(os.environ.get("APS_DATA_DIR", "server/data")) / AGENT_TOKEN_FILE_NAME
        )


def _configured_path() -> Path:
    raw = (os.environ.get("APS_AGENT_TOKEN_PATH") or "").strip()
    path = Path(raw).expanduser() if raw else _default_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    return path.resolve()


def _issuer_roles(issuer: IdentityContext | None) -> set[str]:
    if issuer is None:
        return {"system", "admin"}
    roles = {str(role).lower() for role in (issuer.roles or ())}
    if issuer.auth_kind in {"disabled", "local-admin"}:
        roles.update({"system", "admin"})
    return roles


def _can_manage(issuer: IdentityContext | None, record: dict[str, Any]) -> bool:
    if issuer is None:
        return True
    if str(record.get("tenantUuid") or "") != issuer.tenant_uuid:
        return False
    roles = _issuer_roles(issuer)
    return bool(roles & {"system", "admin"})


class AgentTokenService:
    """File-backed agent read-token registry with HMAC-signed tokens."""

    def __init__(self, secret: bytes | None = None, path: Path | None = None) -> None:
        self.secret = secret
        self.path = (path or _configured_path()).resolve()
        self._lock = threading.RLock()
        self.path.parent.mkdir(parents=True, exist_ok=True)

    def _require_secret(self) -> bytes:
        if not self.secret:
            raise AuthError(
                "AGENT_TOKEN_NOT_CONFIGURED",
                "Pi Agent 凭证服务未配置：请设置 APS_AGENT_TOKEN_SECRET",
                503,
            )
        return self.secret

    def _load(self) -> dict[str, Any]:
        try:
            if not self.path.is_file():
                return {"version": AGENT_TOKEN_VERSION, "tokens": []}
            with self.path.open("r", encoding="utf-8") as handle:
                document = json.load(handle)
            if not isinstance(document, dict) or not isinstance(
                document.get("tokens"), list
            ):
                raise TypeError("registry schema")
            return document
        except AuthError:
            raise
        except Exception as exc:
            raise AuthError(
                "AGENT_TOKEN_STORE_CORRUPT",
                f"Pi Agent 凭证登记文件不可读：{self.path}",
                503,
            ) from exc

    def _save(self, document: dict[str, Any]) -> None:
        import tempfile

        with self._lock:
            self.path.parent.mkdir(parents=True, exist_ok=True)
            fd = None
            tmp = None
            try:
                fd, tmp_name = tempfile.mkstemp(
                    dir=str(self.path.parent), prefix=".agent-tokens-", suffix=".tmp"
                )
                tmp = Path(tmp_name)
                with os.fdopen(fd, "w", encoding="utf-8") as handle:
                    json.dump(document, handle, ensure_ascii=False, indent=2)
                os.replace(tmp, self.path)
            except BaseException:
                if fd is not None:
                    try:
                        os.close(fd)
                    except OSError:
                        pass
                if tmp is not None and tmp.exists():
                    try:
                        tmp.unlink()
                    except OSError:
                        pass
                raise

    def _sign(self, payload: dict[str, Any]) -> str:
        raw = json.dumps(payload, separators=(",", ":"), ensure_ascii=True).encode(
            "utf-8"
        )
        body = _b64encode(raw)
        signature = _b64encode(
            hmac.new(
                self._require_secret(), body.encode("ascii"), hashlib.sha256
            ).digest()
        )
        return f"{body}.{signature}"

    def _verify_token(self, token: str) -> dict[str, Any]:
        secret = self._require_secret()
        try:
            body, signature = token.split(".", 1)
            expected = _b64encode(
                hmac.new(secret, body.encode("ascii"), hashlib.sha256).digest()
            )
            if not hmac.compare_digest(signature, expected):
                raise ValueError("signature")
            payload = json.loads(_b64decode_text(body).decode("utf-8"))
            if (
                int(payload.get("v") or 0) != AGENT_TOKEN_VERSION
                or str(payload.get("kind") or "") != AGENT_TOKEN_KIND
            ):
                raise ValueError("token shape")
            token_id = str(payload.get("tokenId") or "")
            if not token_id or len(token_id) > 128:
                raise ValueError("token id")
            return {
                "tokenId": token_id,
                "iat": int(payload.get("iat") or 0),
                "exp": int(payload.get("exp") or 0),
            }
        except AuthError:
            raise
        except Exception as exc:
            raise AuthError("AGENT_TOKEN_INVALID", "Pi Agent 凭证无效", 401) from exc

    def _find(self, token_id: str) -> dict[str, Any] | None:
        document = self._load()
        for record in document.get("tokens") or []:
            if str(record.get("tokenId") or "") == token_id:
                return record
        return None

    def issue(
        self,
        principal: IdentityContext,
        label: str,
        *,
        ttl_days: int = 365,
    ) -> tuple[IdentityContext, str]:
        self._require_secret()
        label = str(label or "").strip()
        if not label or len(label) > 80:
            raise AuthError(
                "AGENT_TOKEN_LABEL_INVALID", "凭证标签必须为 1-80 字符", 400
            )
        try:
            user_id = int(principal.user_id)
            ttl_days = int(ttl_days)
        except (TypeError, ValueError) as exc:
            raise AuthError(
                "AGENT_TOKEN_PRINCIPAL_INVALID", "签发主体标识无效", 400
            ) from exc
        if user_id < 0 or user_id > 2**63 - 1:
            raise AuthError(
                "AGENT_TOKEN_PRINCIPAL_INVALID", "签发主体 user_id 超出合法范围", 400
            )
        tenant_uuid = str(principal.tenant_uuid or "").strip()
        if not tenant_uuid or len(tenant_uuid) > 64:
            raise AuthError(
                "AGENT_TOKEN_PRINCIPAL_INVALID", "签发主体租户标识无效", 400
            )
        if not 1 <= ttl_days <= AGENT_TOKEN_MAX_TTL_DAYS:
            raise AuthError("AGENT_TOKEN_TTL_INVALID", "凭证有效期须为 1-3650 天", 400)

        now = int(time.time())
        token_id = uuid.uuid4().hex
        payload = {
            "v": AGENT_TOKEN_VERSION,
            "kind": AGENT_TOKEN_KIND,
            "scope": [AGENT_TOKEN_SCOPE_READ],
            "tokenId": token_id,
            "iat": now,
            "exp": now + ttl_days * 24 * 60 * 60,
        }
        token = self._sign(payload)
        created_by = str(principal.username or principal.fullname or principal.user_id)
        record = {
            "tokenId": token_id,
            "label": label,
            "principalUserId": user_id,
            "principalUsername": str(principal.username or f"user-{user_id}"),
            "principalFullname": str(principal.fullname or ""),
            "tenantUuid": tenant_uuid,
            "scope": [AGENT_TOKEN_SCOPE_READ],
            "iat": now,
            "exp": now + ttl_days * 24 * 60 * 60,
            "revoked": False,
            "createdBy": created_by,
            "createdAt": now,
        }
        with self._lock:
            document = self._load()
            document["version"] = AGENT_TOKEN_VERSION
            document.setdefault("tokens", []).append(record)
            self._save(document)
        return self._identity_from_record(record), token

    @staticmethod
    def _identity_from_record(record: dict[str, Any]) -> IdentityContext:
        return IdentityContext(
            user_id=int(record.get("principalUserId") or 0),
            username=str(
                record.get("principalUsername") or f"agent-{record.get('tokenId')}"
            ),
            fullname=str(
                record.get("principalFullname") or record.get("label") or "Pi Agent"
            ),
            tenant_uuid=str(record.get("tenantUuid") or "platform"),
            roles=("pi-agent",),
            expires_at=int(record.get("exp") or 0),
            auth_kind="agent",
        )

    def authenticate(self, token: str) -> IdentityContext:
        payload = self._verify_token(token)
        now = int(time.time())
        if payload["exp"] <= now:
            raise AuthError("AGENT_TOKEN_EXPIRED", "Pi Agent 凭证已到期", 401)
        record = self._find(str(payload["tokenId"]))
        if record is None:
            raise AuthError("AGENT_TOKEN_INVALID", "Pi Agent 凭证未登记或已删除", 401)
        if bool(record.get("revoked")):
            raise AuthError("AGENT_TOKEN_REVOKED", "Pi Agent 凭证已吊销", 401)
        if (
            int(record.get("exp") or 0) != payload["exp"]
            or int(record.get("iat") or 0) != payload["iat"]
        ):
            raise AuthError("AGENT_TOKEN_INVALID", "Pi Agent 凭证登记不一致", 401)
        if AGENT_TOKEN_SCOPE_READ not in (record.get("scope") or []):
            raise AuthError(
                "AGENT_TOKEN_SCOPE_INVALID", "Pi Agent 凭证缺少只读作用域", 403
            )
        return self._identity_from_record(record)

    def list_tokens(self, tenant_uuid: str | None = None) -> list[dict[str, Any]]:
        document = self._load()
        rows: list[dict[str, Any]] = []
        for record in document.get("tokens") or []:
            if tenant_uuid and str(record.get("tenantUuid") or "") != tenant_uuid:
                continue
            rows.append(
                {
                    "tokenId": str(record.get("tokenId") or ""),
                    "label": str(record.get("label") or ""),
                    "principalUserId": int(record.get("principalUserId") or 0),
                    "principalUsername": str(record.get("principalUsername") or ""),
                    "tenantUuid": str(record.get("tenantUuid") or ""),
                    "scope": list(record.get("scope") or []),
                    "iat": int(record.get("iat") or 0),
                    "exp": int(record.get("exp") or 0),
                    "revoked": bool(record.get("revoked")),
                    "createdBy": str(record.get("createdBy") or ""),
                    "createdAt": int(record.get("createdAt") or 0),
                }
            )
        return rows

    def revoke(
        self,
        token_id: str,
        *,
        tenant_uuid: str | None = None,
        issuer: IdentityContext | None = None,
    ) -> dict[str, Any]:
        token_id = str(token_id or "").strip()
        if not token_id:
            raise AuthError("AGENT_TOKEN_ID_REQUIRED", "缺少凭证 ID", 400)
        self._require_secret()
        with self._lock:
            document = self._load()
            for record in document.get("tokens") or []:
                if str(record.get("tokenId") or "") != token_id:
                    continue
                if tenant_uuid and str(record.get("tenantUuid") or "") != tenant_uuid:
                    raise AuthError("AGENT_TOKEN_NOT_FOUND", "凭证不存在", 404)
                if not _can_manage(issuer, record):
                    raise AuthError(
                        "AGENT_TOKEN_FORBIDDEN", "当前身份无权吊销该凭证", 403
                    )
                record["revoked"] = True
                self._save(document)
                return {
                    "tokenId": token_id,
                    "revoked": True,
                    "tenantUuid": str(record.get("tenantUuid") or ""),
                }
        raise AuthError("AGENT_TOKEN_NOT_FOUND", "凭证不存在", 404)


_service: AgentTokenService | None = None
_service_signature: tuple[bytes | None, str] | None = None


def reset_agent_token_service() -> None:
    """丢弃进程级缓存，供环境变更后的测试与运维脚本调用。"""
    global _service, _service_signature
    _service = None
    _service_signature = None


def _get_service() -> AgentTokenService:
    global _service, _service_signature
    secret = (os.environ.get("APS_AGENT_TOKEN_SECRET") or "").strip().encode(
        "utf-8"
    ) or None
    signature = (secret, str(_configured_path()))
    if _service is None or _service_signature != signature:
        _service = AgentTokenService(secret=secret, path=Path(signature[1]))
        _service_signature = signature
    return _service


def issue(
    principal: IdentityContext, label: str, *, ttl_days: int = 365
) -> tuple[IdentityContext, str]:
    return _get_service().issue(principal, label, ttl_days=ttl_days)


def authenticate(token: str) -> IdentityContext:
    return _get_service().authenticate(token)


def list_tokens(tenant_uuid: str | None = None) -> list[dict[str, Any]]:
    return _get_service().list_tokens(tenant_uuid=tenant_uuid)


def revoke(
    token_id: str,
    *,
    tenant_uuid: str | None = None,
    issuer: IdentityContext | None = None,
) -> dict[str, Any]:
    return _get_service().revoke(token_id, tenant_uuid=tenant_uuid, issuer=issuer)