"""Grok OAuth helpers for systray-ai.

HARD RULES:
- Never call the token endpoint unless access JWT is expired (or GET returns 401).
- On refresh, write tokens only into the single account auth.json under flock.
- Dual refresh of the same refresh token across homes revokes the session.
"""

from __future__ import annotations

import base64
import fcntl
import json
import os
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Final


BILLING_URL: Final[str] = "https://cli-chat-proxy.grok.com/v1/billing"
USER_URL: Final[str] = "https://cli-chat-proxy.grok.com/v1/user"
TOKEN_URL: Final[str] = "https://auth.x.ai/oauth2/token"
DEVICE_CODE_URL: Final[str] = "https://auth.x.ai/oauth2/device/code"
CLIENT_ID: Final[str] = "b1a00492-073a-47ea-816f-4c329264a828"
ISSUER: Final[str] = "https://auth.x.ai"
SCOPES: Final[str] = "openid profile email offline_access api:access"
AUTH_SLOT: Final[str] = f"{ISSUER}::{CLIENT_ID}"
EXPIRY_SKEW_SECS: Final[int] = 30


class GrokOAuthError(RuntimeError):
    pass


class GrokOAuthAuthError(GrokOAuthError):
    pass


class GrokOAuthTransportError(GrokOAuthError):
    pass


@dataclass(frozen=True)
class GrokTokens:
    access: str
    refresh: str
    expires_at_iso: str | None
    entry: dict[str, Any]
    slot: str


class GrokOAuthSession:
    def __init__(self, auth_path: Path | str, timeout_secs: float = 10.0) -> None:
        self.auth_path = Path(auth_path)
        self._timeout_secs = timeout_secs

    def get_billing(self, *, allow_refresh: bool = True) -> dict[str, Any]:
        return self._authorized_json(BILLING_URL, allow_refresh=allow_refresh)

    def get_user(self, *, allow_refresh: bool = True) -> dict[str, Any]:
        return self._authorized_json(USER_URL, allow_refresh=allow_refresh)

    def _authorized_json(self, url: str, *, allow_refresh: bool = True) -> dict[str, Any]:
        tokens = self._load_tokens()
        try:
            return self._get_json(url, tokens.access)
        except GrokOAuthAuthError:
            if not allow_refresh:
                raise
            tokens = self._refresh_if_needed(tokens, force=True)
            return self._get_json(url, tokens.access)

    def _get_json(self, url: str, access: str) -> dict[str, Any]:
        request = urllib.request.Request(
            url,
            headers={
                "Authorization": f"Bearer {access}",
                "Accept": "application/json",
                "User-Agent": "systray-ai-grok",
            },
            method="GET",
        )
        try:
            with urllib.request.urlopen(request, timeout=self._timeout_secs) as response:
                payload = json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            body = exc.read().decode("utf-8", errors="replace")
            if exc.code in (401, 403):
                raise GrokOAuthAuthError(f"Grok auth rejected for {url}: {exc.code}") from exc
            raise GrokOAuthTransportError(f"Grok GET {url} failed: {exc.code} {body[:200]}") from exc
        except (OSError, TimeoutError, json.JSONDecodeError) as exc:
            raise GrokOAuthTransportError(f"Grok GET {url} failed: {exc}") from exc
        if not isinstance(payload, dict):
            raise GrokOAuthTransportError(f"Grok GET {url} returned non-object")
        return payload

    def _load_tokens(self) -> GrokTokens:
        try:
            payload = json.loads(self.auth_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError) as exc:
            raise GrokOAuthError(f"cannot read {self.auth_path}: {exc}") from exc
        if not isinstance(payload, dict):
            raise GrokOAuthError(f"invalid auth payload in {self.auth_path}")

        for slot, entry in payload.items():
            if not isinstance(entry, dict):
                continue
            access = str(entry.get("key") or entry.get("access_token") or "").strip()
            refresh = str(entry.get("refresh_token") or "").strip()
            if not access or not refresh:
                continue
            if "auth.x.ai" not in str(slot) and not entry.get("oidc_client_id"):
                continue
            return GrokTokens(
                access=access,
                refresh=refresh,
                expires_at_iso=str(entry["expires_at"]) if isinstance(entry.get("expires_at"), str) else None,
                entry=entry,
                slot=str(slot),
            )
        raise GrokOAuthError(f"no Grok OAuth tokens in {self.auth_path}")

    def _access_expired(self, access: str) -> bool:
        try:
            claims = _decode_jwt_payload(access)
            exp = claims.get("exp")
            if not isinstance(exp, (int, float)):
                return False
            return float(exp) <= time.time() + EXPIRY_SKEW_SECS
        except Exception:
            return False

    def _refresh_if_needed(self, tokens: GrokTokens, *, force: bool) -> GrokTokens:
        if not force and not self._access_expired(tokens.access):
            return tokens
        with self._adjacent_lock():
            latest = self._load_tokens()
            if not force and not self._access_expired(latest.access):
                return latest
            if latest.refresh != tokens.refresh and not self._access_expired(latest.access):
                return latest
            refreshed = self._refresh_tokens(latest.refresh)
            new_payload = self._merge_payload(latest, refreshed)
            self._atomic_write_payload(new_payload)
            return self._load_tokens()

    def _refresh_tokens(self, refresh_token: str) -> dict[str, Any]:
        body = urllib.parse.urlencode(
            {
                "grant_type": "refresh_token",
                "refresh_token": refresh_token,
                "client_id": CLIENT_ID,
            }
        ).encode("utf-8")
        request = urllib.request.Request(
            TOKEN_URL,
            data=body,
            headers={"Content-Type": "application/x-www-form-urlencoded"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=self._timeout_secs) as response:
                payload = json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            text = exc.read().decode("utf-8", errors="replace")
            raise GrokOAuthAuthError(f"refresh failed: {exc.code} {text[:200]}") from exc
        except (OSError, TimeoutError, json.JSONDecodeError) as exc:
            raise GrokOAuthTransportError(f"refresh transport failed: {exc}") from exc
        if not isinstance(payload, dict) or not payload.get("access_token"):
            raise GrokOAuthAuthError("refresh response missing access_token")
        return payload

    def _merge_payload(self, latest: GrokTokens, refreshed: dict[str, Any]) -> dict[str, Any]:
        try:
            root = json.loads(self.auth_path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError):
            root = {}
        if not isinstance(root, dict):
            root = {}
        entry = dict(latest.entry)
        entry["key"] = refreshed["access_token"]
        if refreshed.get("refresh_token"):
            entry["refresh_token"] = refreshed["refresh_token"]
        expires_in = refreshed.get("expires_in")
        if isinstance(expires_in, (int, float)):
            entry["expires_at"] = time.strftime(
                "%Y-%m-%dT%H:%M:%SZ",
                time.gmtime(time.time() + float(expires_in)),
            )
        entry.setdefault("auth_mode", "oidc")
        entry.setdefault("oidc_issuer", ISSUER)
        entry.setdefault("oidc_client_id", CLIENT_ID)
        root[latest.slot] = entry
        return root

    def _atomic_write_payload(self, payload: dict[str, Any]) -> None:
        self.auth_path.parent.mkdir(parents=True, exist_ok=True)
        with tempfile.NamedTemporaryFile(
            "w",
            encoding="utf-8",
            dir=self.auth_path.parent,
            delete=False,
        ) as handle:
            handle.write(json.dumps(payload, indent=2) + "\n")
            os.fchmod(handle.fileno(), 0o600)
            temp_name = handle.name
        os.replace(temp_name, self.auth_path)

    def _adjacent_lock(self):
        lock_path = self.auth_path.with_suffix(self.auth_path.suffix + ".lock")
        lock_path.parent.mkdir(parents=True, exist_ok=True)
        handle = lock_path.open("a+", encoding="utf-8")
        fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
        return _LockHandle(handle)


class _LockHandle:
    def __init__(self, handle: Any) -> None:
        self._handle = handle

    def __enter__(self) -> _LockHandle:
        return self

    def __exit__(self, *_args: object) -> None:
        try:
            fcntl.flock(self._handle.fileno(), fcntl.LOCK_UN)
        finally:
            self._handle.close()


def write_auth_file(
    path: Path,
    *,
    access: str,
    refresh: str,
    expires_in: float | None = None,
    extra: dict[str, Any] | None = None,
) -> None:
    entry: dict[str, Any] = {
        "key": access,
        "refresh_token": refresh,
        "auth_mode": "oidc",
        "oidc_issuer": ISSUER,
        "oidc_client_id": CLIENT_ID,
    }
    if expires_in is not None:
        entry["expires_at"] = time.strftime(
            "%Y-%m-%dT%H:%M:%SZ",
            time.gmtime(time.time() + float(expires_in)),
        )
    if extra:
        for key, value in extra.items():
            if value is not None:
                entry[key] = value
    payload = {AUTH_SLOT: entry}
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(
        "w",
        encoding="utf-8",
        dir=path.parent,
        delete=False,
    ) as handle:
        handle.write(json.dumps(payload, indent=2) + "\n")
        os.fchmod(handle.fileno(), 0o600)
        temp_name = handle.name
    os.replace(temp_name, path)


def start_device_code(timeout_secs: float = 15.0) -> dict[str, Any]:
    body = urllib.parse.urlencode(
        {"client_id": CLIENT_ID, "scope": SCOPES}
    ).encode("utf-8")
    request = urllib.request.Request(
        DEVICE_CODE_URL,
        data=body,
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=timeout_secs) as response:
        payload = json.loads(response.read().decode("utf-8"))
    if not isinstance(payload, dict) or not payload.get("device_code"):
        raise GrokOAuthError("device code start failed")
    return payload


def poll_device_code(device: dict[str, Any], timeout_secs: float = 900.0) -> dict[str, Any]:
    interval = max(3, int(device.get("interval") or 5))
    deadline = time.time() + min(float(device.get("expires_in") or 1800), timeout_secs)
    while time.time() < deadline:
        body = urllib.parse.urlencode(
            {
                "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
                "device_code": device["device_code"],
                "client_id": CLIENT_ID,
            }
        ).encode("utf-8")
        request = urllib.request.Request(
            TOKEN_URL,
            data=body,
            headers={"Content-Type": "application/x-www-form-urlencoded"},
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=20) as response:
                return json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            text = exc.read().decode("utf-8", errors="replace")
            try:
                err = json.loads(text).get("error", "")
            except json.JSONDecodeError:
                err = text
            if err in ("authorization_pending", "slow_down"):
                time.sleep(interval * (2 if err == "slow_down" else 1))
                continue
            raise GrokOAuthAuthError(f"device poll failed: {err}") from exc
        time.sleep(interval)
    raise GrokOAuthError("device login timed out")


def _decode_jwt_payload(access: str) -> dict[str, Any]:
    parts = access.split(".")
    if len(parts) < 2:
        raise ValueError("not a jwt")
    segment = parts[1]
    padding = "=" * (-len(segment) % 4)
    raw = base64.urlsafe_b64decode(segment + padding)
    payload = json.loads(raw.decode("utf-8"))
    if not isinstance(payload, dict):
        raise ValueError("jwt payload not object")
    return payload