import base64
import gzip
import json
import logging
import time
from threading import Lock

import httpx

logger = logging.getLogger(__name__)

_jwt_token: str | None = None
_jwt_lock = Lock()


def _is_token_expired(token: str | None) -> bool:
    if not token:
        return True
    try:
        parts = token.split(".")
        if len(parts) != 3:
            return True
        padded = parts[1] + "=="
        payload = json.loads(base64.b64decode(padded))
        exp = payload.get("exp", 0)
        return bool(time.time() >= exp - 60)  # 60s buffer
    except Exception:
        return True


async def authorize(api_base_url: str, secret_token: str) -> str:
    global _jwt_token
    body = json.dumps({"secret": secret_token}).encode("utf-8")
    compressed = gzip.compress(body)
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(
            f"{api_base_url}/auth/v2/sign-in-with-secret-token",
            content=compressed,
            headers={
                "Content-Type": "application/json",
                "Content-Encoding": "gzip",
                "x-client-name": "py-agent",
            },
        )
        response.raise_for_status()
        data = response.json()
        jwt: str = data["idToken"]
    with _jwt_lock:
        _jwt_token = jwt
    return jwt


async def get_jwt_token(api_base_url: str, secret_token: str) -> str:
    with _jwt_lock:
        if _jwt_token and not _is_token_expired(_jwt_token):
            return _jwt_token
    return await authorize(api_base_url, secret_token)
