from __future__ import annotations

import fcntl
import json
import math
import os
import socket
import tempfile
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Final, Mapping


USAGE_URL: Final[str] = "https://api.anthropic.com/api/oauth/usage"
TOKEN_URL: Final[str] = "https://platform.claude.com/v1/oauth/token"
CLAUDE_CODE_CLIENT_ID: Final[str] = "9d1c250a-e61b-44d9-88ed-5944d1962f5e"
OAUTH_BETA_HEADER: Final[str] = "oauth-2025-04-20"
# platform.claude.com is behind a Cloudflare rule that answers 1010 to any request without one.
USER_AGENT: Final[str] = "systray-ai/1.0"
ACCESS_TOKEN_EXPIRY_SKEW_SECONDS: Final[float] = 120.0
# A refresh token is consumed the moment the token endpoint processes the request,
# so a timeout that abandons an in-flight rotation can orphan the new grant. Give
# the token endpoint far more time than ordinary usage polling gets.
TOKEN_REFRESH_TIMEOUT_SECONDS: Final[float] = 30.0
ROTATE_JOURNAL_SUFFIX: Final[str] = ".rotate-journal.json"


class ClaudeOAuthError(RuntimeError):
    pass


class ClaudeOAuthCredentialsError(ClaudeOAuthError):
    pass


class ClaudeOAuthAuthError(ClaudeOAuthError):
    pass


class ClaudeOAuthLoggedOutError(ClaudeOAuthAuthError):
    pass


class ClaudeOAuthProtocolError(ClaudeOAuthError):
    pass


class ClaudeOAuthTimeoutError(ClaudeOAuthError):
    pass


class ClaudeOAuthTransportError(ClaudeOAuthError):
    def __init__(self, message: str, *, retry_after_seconds: float | None = None) -> None:
        super().__init__(message)
        self.retry_after_seconds = retry_after_seconds


@dataclass(frozen=True)
class ClaudeOAuthRequest:
    method: str
    url: str
    headers: Mapping[str, str]
    body: bytes | None = None


@dataclass(frozen=True)
class ClaudeOAuthResponse:
    status_code: int
    body: str
    headers: Mapping[str, str] = field(default_factory=dict)


ClaudeOAuthTransport = Callable[[ClaudeOAuthRequest, float], ClaudeOAuthResponse]


@dataclass(frozen=True)
class UsageReading:
    """Usage payload plus the access token that fetched it — the only durable
    link between a reading and the account it describes, because the
    credentials file may be rewritten by another process mid-fetch."""

    payload: dict[str, object]
    access_token: str


@dataclass(frozen=True)
class _LoadedCredentials:
    payload: dict[str, object]
    oauth_payload: dict[str, object]
    access_token: str
    refresh_token: str


class ClaudeOAuthSession:
    """Consumer of a Claude CLI credentials file.

    By default the session is strictly read-only: Anthropic refresh tokens are
    one-time-use, and a Claude CLI session that has the grant in memory rotates it
    itself — a second rotator replays a consumed refresh token and the provider
    revokes the entire grant chain. ``allow_rotation=True`` is only correct when the
    caller has proven no live CLI session owns this credentials file; then this
    session becomes the sole rotator, and it persists the rotated grant atomically
    so the next CLI launch picks it up."""

    def __init__(
        self,
        credentials_path: Path | str,
        transport: ClaudeOAuthTransport | None = None,
        timeout_secs: float = 5.0,
        allow_rotation: bool = False,
        rotation_backup_path: Path | str | None = None,
    ) -> None:
        if timeout_secs <= 0:
            raise ClaudeOAuthTimeoutError("Claude OAuth timeout must be greater than zero")
        self.credentials_path = Path(credentials_path)
        self._transport = default_transport if transport is None else transport
        self._timeout_secs = timeout_secs
        self._allow_rotation = allow_rotation
        self._rotation_backup_path = (
            None if rotation_backup_path is None else Path(rotation_backup_path)
        )

    def get_usage(self) -> dict[str, object]:
        return self.get_usage_reading().payload

    def get_usage_reading(self) -> UsageReading:
        credentials = self._load_credentials()
        if self._is_expired(credentials):
            credentials = self._freshen(credentials)
        response = self._send(self._build_usage_request(credentials.access_token))
        if response.status_code == 401:
            credentials = self._freshen(credentials)
            response = self._send(self._build_usage_request(credentials.access_token))
            if response.status_code == 401:
                raise ClaudeOAuthLoggedOutError(
                    "Claude OAuth usage request failed with status 401 with the latest on-disk token"
                )
        return UsageReading(
            payload=self._parse_usage_response(response),
            access_token=credentials.access_token,
        )

    @staticmethod
    def _is_expired(credentials: _LoadedCredentials) -> bool:
        expires_at_ms = credentials.oauth_payload.get("expiresAt")
        if isinstance(expires_at_ms, bool) or not isinstance(expires_at_ms, (int, float)):
            return False
        return expires_at_ms / 1000.0 <= time.time() + ACCESS_TOKEN_EXPIRY_SKEW_SECONDS

    def _freshen(self, credentials: _LoadedCredentials) -> _LoadedCredentials:
        latest = self._load_credentials()
        if latest.access_token != credentials.access_token:
            return latest
        if self._allow_rotation:
            return self._rotate(latest)
        if self._journal_path().exists():
            recovered = self._with_rotate_lock(lambda: self._replay_rotate_journal())
            if recovered is not None:
                return recovered
        raise ClaudeOAuthTransportError(
            "Claude OAuth access token is stale and only the Claude CLI may rotate it"
        )

    def _rotate(self, credentials: _LoadedCredentials) -> _LoadedCredentials:
        def rotate_locked() -> _LoadedCredentials:
            recovered = self._replay_rotate_journal()
            if recovered is not None:
                return recovered
            latest = self._load_credentials()
            if latest.access_token != credentials.access_token:
                return latest
            self._preflight_writable()
            rotated = self._request_rotated_grant(latest)
            self._write_rotate_journal(rotated.payload)
            self._atomic_write_payload(rotated.payload)
            self._clear_rotate_journal()
            self._write_rotation_backup(rotated.payload)
            return rotated

        return self._with_rotate_lock(rotate_locked)

    def _with_rotate_lock(self, action: Callable[[], _LoadedCredentials | None]):
        lock_path = self.credentials_path.with_name(self.credentials_path.name + ".rotate.lock")
        try:
            lock_file = open(lock_path, "w", encoding="utf-8")
        except OSError as exc:
            raise ClaudeOAuthCredentialsError(
                "Claude OAuth rotation lock could not be created"
            ) from exc
        with lock_file:
            fcntl.flock(lock_file, fcntl.LOCK_EX)
            return action()

    def _journal_path(self) -> Path:
        return self.credentials_path.with_name(self.credentials_path.name + ROTATE_JOURNAL_SUFFIX)

    def _replay_rotate_journal(self) -> _LoadedCredentials | None:
        """A journal left behind means a rotated grant arrived from the provider but
        the credentials file was never replaced — the journal holds the only copy of
        the live refresh token. Replay it before anything else touches the grant."""
        journal_path = self._journal_path()
        try:
            text = journal_path.read_text(encoding="utf-8")
        except FileNotFoundError:
            return None
        except OSError as exc:
            raise ClaudeOAuthCredentialsError(
                "Claude OAuth rotation journal exists but could not be read"
            ) from exc
        try:
            journaled = self._loaded_credentials_from_raw_payload(json.loads(text))
        except (json.JSONDecodeError, ClaudeOAuthCredentialsError):
            self._clear_rotate_journal()
            return None
        latest = self._load_credentials()
        if journaled.access_token == latest.access_token:
            self._clear_rotate_journal()
            return None
        self._atomic_write_payload(journaled.payload)
        self._clear_rotate_journal()
        self._write_rotation_backup(journaled.payload)
        return journaled

    def _preflight_writable(self) -> None:
        """Prove the grant can be persisted BEFORE the one-time refresh token is
        spent; a full disk or read-only mount must never consume a token."""
        try:
            fd, temp_name = tempfile.mkstemp(
                dir=self.credentials_path.parent, prefix=".credentials.preflight-"
            )
            try:
                os.write(fd, b"{}")
                os.fsync(fd)
            finally:
                os.close(fd)
            os.unlink(temp_name)
        except OSError as exc:
            raise ClaudeOAuthCredentialsError(
                "Claude OAuth credentials directory is not writable; refusing to rotate"
            ) from exc

    def _write_rotate_journal(self, payload: dict[str, object]) -> None:
        journal_path = self._journal_path()
        try:
            fd, temp_name = tempfile.mkstemp(
                dir=self.credentials_path.parent, prefix=".credentials.journal-"
            )
            try:
                os.fchmod(fd, 0o600)
                with os.fdopen(fd, "w", encoding="utf-8") as handle:
                    handle.write(json.dumps(payload))
                    handle.flush()
                    os.fsync(handle.fileno())
                os.replace(temp_name, journal_path)
            except OSError:
                try:
                    os.unlink(temp_name)
                except OSError:
                    pass
                raise
        except OSError as exc:
            raise ClaudeOAuthCredentialsError(
                "rotated Claude OAuth grant could not be journaled; this account needs /login"
            ) from exc

    def _clear_rotate_journal(self) -> None:
        try:
            os.unlink(self._journal_path())
        except OSError:
            pass

    def _write_rotation_backup(self, payload: dict[str, object]) -> None:
        if self._rotation_backup_path is None:
            return
        try:
            backup = self._rotation_backup_path
            backup.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
            fd, temp_name = tempfile.mkstemp(dir=backup.parent, prefix=".backup-")
            try:
                os.fchmod(fd, 0o600)
                with os.fdopen(fd, "w", encoding="utf-8") as handle:
                    handle.write(json.dumps(payload))
                    handle.flush()
                    os.fsync(handle.fileno())
                os.replace(temp_name, backup)
            except OSError:
                try:
                    os.unlink(temp_name)
                except OSError:
                    pass
                raise
        except OSError:
            pass

    def _request_rotated_grant(self, credentials: _LoadedCredentials) -> _LoadedCredentials:
        body = urllib.parse.urlencode(
            {
                "grant_type": "refresh_token",
                "refresh_token": credentials.refresh_token,
                "client_id": CLAUDE_CODE_CLIENT_ID,
            }
        ).encode("utf-8")
        response = self._send(
            ClaudeOAuthRequest(
                method="POST",
                url=TOKEN_URL,
                headers={
                    "Accept": "application/json",
                    "Content-Type": "application/x-www-form-urlencoded",
                    "User-Agent": USER_AGENT,
                },
                body=body,
            ),
            timeout_secs=max(self._timeout_secs, TOKEN_REFRESH_TIMEOUT_SECONDS),
        )
        if response.status_code != 200:
            if self._oauth_error_code(response.body) == "invalid_grant":
                raise ClaudeOAuthLoggedOutError(
                    "Claude OAuth refresh grant was rejected as invalid_grant"
                )
            raise ClaudeOAuthTransportError(
                f"Claude OAuth token refresh failed with status {response.status_code}",
                retry_after_seconds=self._retry_after_seconds(response.headers),
            )
        grant = self._parse_json_object(response.body, "token response")
        access_token = self._require_non_empty_string(grant, "access_token")
        refresh_token = grant.get("refresh_token")
        expires_in = grant.get("expires_in")
        oauth_payload = dict(credentials.oauth_payload)
        oauth_payload["accessToken"] = access_token
        if isinstance(refresh_token, str) and refresh_token:
            oauth_payload["refreshToken"] = refresh_token
        if not isinstance(expires_in, bool) and isinstance(expires_in, (int, float)):
            oauth_payload["expiresAt"] = int((time.time() + expires_in) * 1000)
        payload = dict(credentials.payload)
        payload["claudeAiOauth"] = oauth_payload
        return self._loaded_credentials_from_raw_payload(payload)

    def _atomic_write_payload(self, payload: dict[str, object]) -> None:
        directory = self.credentials_path.parent
        try:
            mode = self.credentials_path.stat().st_mode & 0o777
        except OSError:
            mode = 0o600
        fd, temp_name = tempfile.mkstemp(dir=directory, prefix=".credentials.rotate-")
        try:
            os.fchmod(fd, mode)
            with os.fdopen(fd, "w", encoding="utf-8") as handle:
                handle.write(json.dumps(payload))
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(temp_name, self.credentials_path)
        except OSError as exc:
            try:
                os.unlink(temp_name)
            except OSError:
                pass
            raise ClaudeOAuthCredentialsError(
                "rotated Claude OAuth grant could not be persisted; this account needs /login"
            ) from exc
        try:
            dir_fd = os.open(directory, os.O_RDONLY)
            try:
                os.fsync(dir_fd)
            finally:
                os.close(dir_fd)
        except OSError:
            pass

    @staticmethod
    def _oauth_error_code(body: str) -> str | None:
        try:
            payload = json.loads(body)
        except json.JSONDecodeError:
            return None
        if not isinstance(payload, dict):
            return None
        error = payload.get("error")
        return error if isinstance(error, str) else None

    def _parse_usage_response(self, response: ClaudeOAuthResponse) -> dict[str, object]:
        if response.status_code in {401, 403, 404}:
            raise ClaudeOAuthAuthError(
                f"Claude OAuth usage request failed with status {response.status_code}"
            )
        if response.status_code != 200:
            raise ClaudeOAuthTransportError(
                f"Claude OAuth usage request failed with status {response.status_code}",
                retry_after_seconds=self._retry_after_seconds(response.headers),
            )
        payload = self._parse_json_object(response.body, "usage response")
        return payload

    @staticmethod
    def _retry_after_seconds(headers: Mapping[str, str]) -> float | None:
        value = next(
            (header for key, header in headers.items() if key.lower() == "retry-after"),
            None,
        )
        if value is None:
            return None
        try:
            seconds = float(value)
        except (TypeError, ValueError):
            return None
        return seconds if math.isfinite(seconds) and seconds >= 0 else None

    def _build_usage_request(self, access_token: str) -> ClaudeOAuthRequest:
        return ClaudeOAuthRequest(
            method="GET",
            url=USAGE_URL,
            headers={
                "Accept": "application/json",
                "Authorization": f"Bearer {access_token}",
                "User-Agent": USER_AGENT,
                "anthropic-beta": OAUTH_BETA_HEADER,
            },
        )

    def _send(
        self, request: ClaudeOAuthRequest, timeout_secs: float | None = None
    ) -> ClaudeOAuthResponse:
        try:
            response = self._transport(
                request, self._timeout_secs if timeout_secs is None else timeout_secs
            )
        except ClaudeOAuthError:
            raise
        except (TimeoutError, socket.timeout) as exc:
            raise ClaudeOAuthTimeoutError("Claude OAuth request timed out") from exc
        except Exception as exc:
            raise ClaudeOAuthTransportError("Claude OAuth transport failed") from exc

        if not isinstance(response, ClaudeOAuthResponse):
            raise ClaudeOAuthProtocolError("Claude OAuth transport returned an invalid response")
        if isinstance(response.status_code, bool) or not isinstance(response.status_code, int):
            raise ClaudeOAuthProtocolError(
                "Claude OAuth transport response status_code must be an integer"
            )
        if not isinstance(response.body, str):
            raise ClaudeOAuthProtocolError("Claude OAuth transport response body must be text")
        return response

    def _load_credentials(self) -> _LoadedCredentials:
        try:
            text = self.credentials_path.read_text(encoding="utf-8")
        except FileNotFoundError as exc:
            raise ClaudeOAuthCredentialsError("Claude OAuth credentials file is missing") from exc
        except OSError as exc:
            raise ClaudeOAuthCredentialsError("Claude OAuth credentials file could not be read") from exc

        try:
            payload = json.loads(text)
        except json.JSONDecodeError as exc:
            raise ClaudeOAuthCredentialsError(
                "Claude OAuth credentials file must contain a valid JSON object"
            ) from exc
        return self._loaded_credentials_from_raw_payload(payload)

    def _loaded_credentials_from_raw_payload(self, payload: object) -> _LoadedCredentials:
        if not isinstance(payload, dict):
            raise ClaudeOAuthCredentialsError(
                "Claude OAuth credentials file must contain a valid JSON object"
            )
        oauth_payload = payload.get("claudeAiOauth")
        if not isinstance(oauth_payload, dict):
            raise ClaudeOAuthCredentialsError(
                "Claude OAuth credentials file is missing claudeAiOauth"
            )

        access_token = self._require_non_empty_string(
            oauth_payload,
            "accessToken",
            failure_cls=ClaudeOAuthCredentialsError,
            prefix="Claude OAuth credentials file is missing",
        )
        refresh_token = self._require_non_empty_string(
            oauth_payload,
            "refreshToken",
            failure_cls=ClaudeOAuthCredentialsError,
            prefix="Claude OAuth credentials file is missing",
        )
        return _LoadedCredentials(
            payload=dict(payload),
            oauth_payload=dict(oauth_payload),
            access_token=access_token,
            refresh_token=refresh_token,
        )

    @staticmethod
    def _parse_json_object(body: str, label: str) -> dict[str, object]:
        try:
            payload = json.loads(body)
        except json.JSONDecodeError as exc:
            raise ClaudeOAuthProtocolError(
                f"Claude OAuth {label} must be valid JSON"
            ) from exc
        if not isinstance(payload, dict):
            raise ClaudeOAuthProtocolError(
                f"Claude OAuth {label} must be a JSON object"
            )
        return payload

    @staticmethod
    def _require_non_empty_string(
        payload: Mapping[str, object],
        key: str,
        *,
        failure_cls: type[ClaudeOAuthError] = ClaudeOAuthProtocolError,
        prefix: str = "Claude OAuth response field",
    ) -> str:
        value = payload.get(key)
        if isinstance(value, str) and value:
            return value
        raise failure_cls(f"{prefix} {key} must be a non-empty string")


def default_transport(request: ClaudeOAuthRequest, timeout_secs: float) -> ClaudeOAuthResponse:
    raw_request = urllib.request.Request(
        request.url,
        data=request.body,
        headers=dict(request.headers),
        method=request.method,
    )
    try:
        with urllib.request.urlopen(raw_request, timeout=timeout_secs) as response:
            body = response.read().decode("utf-8", errors="replace")
            return ClaudeOAuthResponse(
                status_code=response.getcode(),
                body=body,
                headers=dict(response.headers.items()),
            )
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        return ClaudeOAuthResponse(
            status_code=exc.code,
            body=body,
            headers=dict(exc.headers.items()) if exc.headers is not None else {},
        )
    except urllib.error.URLError as exc:
        reason = getattr(exc, "reason", None)
        if isinstance(reason, (TimeoutError, socket.timeout)):
            raise TimeoutError("timed out") from exc
        raise OSError("transport failure") from exc


__all__ = [
    "ACCESS_TOKEN_EXPIRY_SKEW_SECONDS",
    "CLAUDE_CODE_CLIENT_ID",
    "OAUTH_BETA_HEADER",
    "ROTATE_JOURNAL_SUFFIX",
    "TOKEN_REFRESH_TIMEOUT_SECONDS",
    "TOKEN_URL",
    "USAGE_URL",
    "USER_AGENT",
    "ClaudeOAuthAuthError",
    "ClaudeOAuthCredentialsError",
    "ClaudeOAuthError",
    "ClaudeOAuthLoggedOutError",
    "ClaudeOAuthProtocolError",
    "ClaudeOAuthRequest",
    "ClaudeOAuthResponse",
    "ClaudeOAuthSession",
    "ClaudeOAuthTimeoutError",
    "ClaudeOAuthTransport",
    "ClaudeOAuthTransportError",
    "default_transport",
]