"""Encrypted on-disk store for OAuth clients, sessions and token hashes.

One file, ``DATA_DIR/oauth-store.json``, holds the whole state as a single
Fernet-encrypted JSON document. The Fernet key is derived from
``TOKEN_ENCRYPTION_KEY`` with HKDF-SHA256, so the setting can be any long random
string (the setup command writes a ``secrets.token_urlsafe`` value, not a Fernet
key).

What is stored:

* registered MCP clients, including their client secrets (the SDK compares the
  presented secret with the stored one at ``/token``);
* sessions: the Gitea login and the user's Gitea access and refresh tokens;
* our own access and refresh tokens, only as SHA-256 hashes.

File safety: the directory is created with mode 0700 (an existing one must be owned by the
running user and is tightened to 0700, otherwise startup stops) and the file is written
with mode 0600 through a fresh ``O_EXCL`` temporary file and an atomic
``os.replace``. A symlink at the store path is refused. A file that cannot be
decrypted or parsed stops startup with a clear message instead of being wiped,
because wiping would silently sign every user out and hide a wrong key.

Mutations happen synchronously in memory (so a check-and-delete is atomic under
asyncio) and are persisted with :meth:`OAuthStore.save`.
"""

from __future__ import annotations

import asyncio
import base64
import contextlib
import hashlib
import json
import logging
import os
import secrets
import stat
import time
from collections.abc import Callable, Iterator
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Literal

import anyio.to_thread
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

logger = logging.getLogger(__name__)

STORE_FILENAME = "oauth-store.json"
STORE_VERSION = 1
HKDF_INFO = b"aegis-gitea-mcp oauth store"

#: A session that never received tokens (an abandoned sign-in) is dropped after this long.
ABANDONED_SESSION_SECONDS = 600.0

Clock = Callable[[], float]
TokenKind = Literal["access", "refresh"]


class OAuthStoreError(Exception):
    """The OAuth store cannot be used (startup error; the message is safe to print)."""


def derive_fernet_key(secret: str) -> bytes:
    """Derive the Fernet key for the store from ``TOKEN_ENCRYPTION_KEY`` (HKDF-SHA256)."""
    hkdf = HKDF(algorithm=hashes.SHA256(), length=32, salt=None, info=HKDF_INFO)
    return base64.urlsafe_b64encode(hkdf.derive(secret.encode("utf-8")))


def token_hash(value: str) -> str:
    """SHA-256 hex digest used as the stored form of our tokens, codes and states."""
    return hashlib.sha256(value.encode("utf-8")).hexdigest()


@dataclass
class ClientRecord:
    """A dynamically registered MCP client.

    Attributes:
        info: ``OAuthClientInformationFull`` as JSON (includes the client secret).
        registered_at: Registration time (epoch seconds).
        last_used_at: Last time the client was looked up (epoch seconds).
    """

    info: dict[str, Any] = field(repr=False)
    registered_at: float
    last_used_at: float


@dataclass
class SessionRecord:
    """One signed-in grant: a Gitea user, the client it was granted to and the Gitea tokens.

    Attributes:
        session_id: Random identifier (never sent to clients).
        client_id: MCP client the grant belongs to.
        login: Gitea login of the user.
        is_admin: Whether Gitea reported the user as site admin at sign-in
            (informational only; Gitea enforces permissions).
        gitea_access_token: The user's Gitea access token.
        gitea_refresh_token: The user's Gitea refresh token, if Gitea issued one.
        gitea_expires_at: Expiry of the Gitea access token (epoch seconds), if known.
        scopes: Our scopes granted to the client.
        resource: The resource (MCP endpoint URL) the grant is bound to.
        created_at: Sign-in time (epoch seconds).
        previous_refresh_hash: Hash of the last rotated-out refresh token (reuse detection).
    """

    session_id: str
    client_id: str
    login: str
    is_admin: bool
    gitea_access_token: str = field(repr=False)
    gitea_refresh_token: str | None = field(repr=False)
    gitea_expires_at: float | None
    scopes: list[str]
    resource: str
    created_at: float
    previous_refresh_hash: str | None = field(default=None, repr=False)


@dataclass
class TokenRecord:
    """One of our tokens, stored under its SHA-256 hash."""

    kind: TokenKind
    session_id: str
    client_id: str
    scopes: list[str]
    expires_at: float


def _restrict_directory(directory: Path) -> None:
    """Create the state directory with mode 0700, or check and tighten an existing one.

    Whoever controls the directory can swap the store for an older copy (which
    resurrects revoked grants) or delete it, so an existing directory must belong
    to the effective user and must not stay writable by group or others.

    Raises:
        OAuthStoreError: The directory is owned by another user or remains writable
            by group or others after an attempt to tighten it.
    """
    if not directory.exists():
        directory.mkdir(mode=0o700, parents=True, exist_ok=True)
        return
    info = os.stat(directory)
    if not stat.S_ISDIR(info.st_mode):
        raise OAuthStoreError(f"{directory} is not a directory")
    owner = getattr(os, "geteuid", lambda: info.st_uid)()
    if info.st_uid != owner:
        raise OAuthStoreError(
            f"{directory} is owned by another user (uid {info.st_uid}, this process runs as "
            f"uid {owner}); refusing to keep the OAuth token store there. Use a DATA_DIR "
            "owned by the service user"
        )
    mode = stat.S_IMODE(info.st_mode)
    if mode & 0o077:
        try:
            directory.chmod(0o700)
        except OSError:
            logger.warning("could not restrict the permissions of the state directory to 0700")
        if stat.S_IMODE(os.stat(directory).st_mode) & 0o022:
            raise OAuthStoreError(
                f"{directory} is writable by group or others and could not be restricted; "
                "refusing to keep the OAuth token store there. Run 'chmod 700' on it"
            )


def _parse_document(document: Any) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
    if not isinstance(document, dict) or document.get("version") != STORE_VERSION:
        raise ValueError("unsupported store version")
    clients_raw = document.get("clients", {})
    sessions_raw = document.get("sessions", {})
    tokens_raw = document.get("tokens", {})
    if not all(isinstance(part, dict) for part in (clients_raw, sessions_raw, tokens_raw)):
        raise ValueError("malformed store sections")
    clients = {str(key): ClientRecord(**value) for key, value in clients_raw.items()}
    sessions = {str(key): SessionRecord(**value) for key, value in sessions_raw.items()}
    tokens = {str(key): TokenRecord(**value) for key, value in tokens_raw.items()}
    for record in clients.values():
        if not isinstance(record.info, dict):
            raise ValueError("malformed client record")
    for token in tokens.values():
        if token.kind not in ("access", "refresh"):
            raise ValueError("malformed token record")
    return clients, sessions, tokens


class OAuthStore:
    """In-memory OAuth state backed by one encrypted file.

    Use :meth:`load` to create it. All getters and mutators are synchronous;
    :meth:`save` persists the current state (pruning expired entries first).

    Args:
        directory: State directory (``DATA_DIR``).
        secret: ``TOKEN_ENCRYPTION_KEY``.
        clock: Wall clock (epoch seconds); tests inject a fake one.
    """

    def __init__(self, directory: Path, secret: str, clock: Clock = time.time) -> None:
        self._directory = directory
        self._path = directory / STORE_FILENAME
        self._fernet = Fernet(derive_fernet_key(secret))
        self._clock = clock
        self._lock = asyncio.Lock()
        self.clients: dict[str, ClientRecord] = {}
        self.sessions: dict[str, SessionRecord] = {}
        self.tokens: dict[str, TokenRecord] = {}

    @classmethod
    def load(cls, directory: Path, secret: str, clock: Clock = time.time) -> OAuthStore:
        """Open the store, reading and decrypting the file when it exists.

        Raises:
            OAuthStoreError: The directory belongs to another user or stays group or
                world writable, the path is a symlink or not a regular file, or the
                file cannot be decrypted or parsed.
        """
        store = cls(directory, secret, clock)
        _restrict_directory(directory)
        store._read()
        store.prune()
        return store

    @property
    def path(self) -> Path:
        """Location of the store file."""
        return self._path

    # Reading and writing the file

    def _read(self) -> None:
        try:
            info = os.lstat(self._path)
        except FileNotFoundError:
            return
        if stat.S_ISLNK(info.st_mode):
            raise OAuthStoreError(
                f"{self._path} is a symlink; refusing to use it as the OAuth token store"
            )
        if not stat.S_ISREG(info.st_mode):
            raise OAuthStoreError(f"{self._path} is not a regular file")
        fd = os.open(self._path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0))
        with os.fdopen(fd, "rb") as handle:
            raw = handle.read()
        if stat.S_IMODE(info.st_mode) & 0o077:
            logger.warning("OAuth store file had loose permissions; restricting it to 0600")
            os.chmod(self._path, 0o600)
        try:
            plain = self._fernet.decrypt(raw)
        except InvalidToken:
            raise OAuthStoreError(
                f"cannot decrypt {self._path}: TOKEN_ENCRYPTION_KEY does not match the key the "
                "store was written with, or the file is damaged. Restore the original key, or "
                "move the file away to start over (this signs every OAuth user out)."
            ) from None
        try:
            self.clients, self.sessions, self.tokens = _parse_document(json.loads(plain))
        except (ValueError, TypeError, KeyError):
            raise OAuthStoreError(
                f"{self._path} is not a valid OAuth store document; move the file away to "
                "start over (this signs every OAuth user out)"
            ) from None

    def _document(self) -> bytes:
        document = {
            "version": STORE_VERSION,
            "clients": {key: asdict(value) for key, value in self.clients.items()},
            "sessions": {key: asdict(value) for key, value in self.sessions.items()},
            "tokens": {key: asdict(value) for key, value in self.tokens.items()},
        }
        return json.dumps(document, separators=(",", ":")).encode("utf-8")

    def _write(self, data: bytes) -> None:
        _restrict_directory(self._directory)
        with contextlib.suppress(FileNotFoundError):
            if stat.S_ISLNK(os.lstat(self._path).st_mode):
                raise OAuthStoreError(
                    f"{self._path} is a symlink; refusing to write the OAuth token store"
                )
        temporary = self._directory / f".{STORE_FILENAME}.{secrets.token_hex(8)}.tmp"
        flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
        fd = os.open(temporary, flags, 0o600)
        try:
            with os.fdopen(fd, "wb") as handle:
                handle.write(data)
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(temporary, self._path)
        except BaseException:
            with contextlib.suppress(FileNotFoundError):
                os.unlink(temporary)
            raise

    async def save(self) -> None:
        """Prune expired entries and persist the state atomically (encrypted, mode 0600)."""
        async with self._lock:
            self.prune()
            data = self._fernet.encrypt(self._document())
            # The write fsyncs: keep that off the event loop (the bytes are already built,
            # so the thread touches no shared state).
            await anyio.to_thread.run_sync(self._write, data)

    # Maintenance

    def prune(self) -> None:
        """Drop expired tokens, and sessions that have no tokens left."""
        now = self._clock()
        for key in [key for key, token in self.tokens.items() if token.expires_at <= now]:
            del self.tokens[key]
        live = {token.session_id for token in self.tokens.values()}
        for session_id, session in list(self.sessions.items()):
            if session_id in live:
                continue
            if session.created_at + ABANDONED_SESSION_SECONDS < now:
                del self.sessions[session_id]

    # Sessions and tokens

    def session_token_hashes(self, session_id: str) -> list[str]:
        """Hashes of every token issued for ``session_id``."""
        return [key for key, token in self.tokens.items() if token.session_id == session_id]

    def end_session(self, session_id: str) -> bool:
        """Delete a session and all of its tokens; return whether it existed."""
        for key in self.session_token_hashes(session_id):
            del self.tokens[key]
        return self.sessions.pop(session_id, None) is not None

    def sessions_of(self, login: str) -> Iterator[SessionRecord]:
        """Every session of one Gitea login."""
        return (session for session in self.sessions.values() if session.login == login)

    def session_by_previous_refresh(self, refresh_hash: str) -> SessionRecord | None:
        """The session whose last rotated-out refresh token has ``refresh_hash``."""
        for session in self.sessions.values():
            if session.previous_refresh_hash is not None and secrets.compare_digest(
                session.previous_refresh_hash, refresh_hash
            ):
                return session
        return None

    def client_has_sessions(self, client_id: str) -> bool:
        """Whether any session belongs to ``client_id``."""
        return any(session.client_id == client_id for session in self.sessions.values())