"""MCP bearer token store — authentication for remote MCP clients.

The SSE transport (``/mcp/sse`` on the running FastAPI server) is reachable
over the network, unlike the stdio launcher which only accepts a subprocess
parent. Exposing tool execution to anyone who can reach port 8000 is
unacceptable, so every SSE request must carry an ``Authorization: Bearer <token>``
header (or ``?token=...`` query param) that resolves to a live, non-revoked
token in this store.

Tokens are stored as **sha256 hashes of the secret half**, never in plaintext.
The id half is public (used for lookup); the secret half is what the user
pastes into their MCP client config. We show the full raw token exactly once
at creation time — after that, we can only revoke, never read.

Schema lives in ``mcp_tokens`` on the per-vault SQLite DB.
"""

from __future__ import annotations

import hashlib
import hmac
import logging
import secrets
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any

import aiosqlite

log = logging.getLogger(__name__)


_RAW_PREFIX = "agnt_mcp_"  # Makes leaked tokens grep-able / identifiable.
_ID_LEN = 16  # hex chars — public, used for lookup
_SECRET_LEN = 32  # hex chars — private, only the hash lands in DB
_TOTAL_LEN = len(_RAW_PREFIX) + _ID_LEN + _SECRET_LEN


# Supported token scopes. Keep this set closed — unknown scopes fall back
# to ``full`` in ``_normalize_scope()`` with a logged warning, never crash.
# Adding a new scope is additive: update ``SCOPES``, add enforcement
# in mcp/server.py:_is_tool_allowed_for_scope(), expose in the UI.
SCOPES: frozenset[str] = frozenset({"full", "read_only"})
DEFAULT_SCOPE = "full"


def _normalize_scope(raw: str | None) -> str:
    """Coerce stored / incoming scope strings to a known value. Unknown → full."""
    if raw in SCOPES:
        return raw
    if raw:  # non-empty but unrecognized — log so it surfaces in debugging
        log.warning("Unknown token scope %r, defaulting to %r", raw, DEFAULT_SCOPE)
    return DEFAULT_SCOPE


@dataclass(slots=True, frozen=True)
class MintedToken:
    """Returned once at creation — the user must save this value themselves."""

    token_id: str
    raw_token: str  # Full "agnt_mcp_<id><secret>" — DO NOT log.
    name: str
    scope: str
    created_at: str


@dataclass(slots=True, frozen=True)
class VerifiedToken:
    """Returned by ``TokenStore.verify`` on success — scope travels with identity.

    Callers use the scope to decide which tools the bearer can invoke. The
    enforcement layer is in ``mcp/server.py``, not here — this module only
    reports what was signed.
    """

    token_id: str
    scope: str


def _hash_secret(secret: str) -> str:
    """Constant-time-safe hash of the secret half of a token."""
    return hashlib.sha256(secret.encode("utf-8")).hexdigest()


def _parse_raw_token(raw: str) -> tuple[str, str] | None:
    """Split ``agnt_mcp_<id><secret>`` into ``(id, secret)``.

    Returns ``None`` when the input shape is wrong — callers treat that
    as an invalid token, not an error condition.
    """
    if not isinstance(raw, str) or not raw.startswith(_RAW_PREFIX):
        return None
    body = raw[len(_RAW_PREFIX):]
    if len(body) != _ID_LEN + _SECRET_LEN:
        return None
    # Character whitelist — reject anything non-hex (prevents SQLi vectors
    # even though we use parameterized queries; defense in depth).
    if not all(c in "0123456789abcdef" for c in body.lower()):
        return None
    return body[:_ID_LEN], body[_ID_LEN:]


async def init_schema(db: aiosqlite.Connection) -> None:
    """Create the ``mcp_tokens`` table. Idempotent.

    Includes an idempotent migration to add the ``scope`` column on DBs
    minted before scoped tokens existed — existing rows get ``full`` so
    the prior behavior is preserved exactly.
    """
    await db.execute(
        """
        CREATE TABLE IF NOT EXISTS mcp_tokens (
            id TEXT PRIMARY KEY,
            secret_hash TEXT NOT NULL,
            name TEXT NOT NULL DEFAULT '',
            scope TEXT NOT NULL DEFAULT 'full',
            created_at TEXT NOT NULL,
            last_used_at TEXT DEFAULT NULL,
            revoked INTEGER NOT NULL DEFAULT 0,
            revoked_at TEXT DEFAULT NULL
        )
        """,
    )
    # Migration: add scope column to tables minted before scoping shipped.
    # ALTER TABLE ADD COLUMN is idempotent via the except guard below.
    try:
        await db.execute(
            "ALTER TABLE mcp_tokens ADD COLUMN scope TEXT NOT NULL DEFAULT 'full'",
        )
    except Exception:
        pass  # column already exists — idempotent
    await db.commit()


class TokenStore:
    """Persistent MCP bearer-token store."""

    def __init__(self, db: aiosqlite.Connection) -> None:
        self._db = db

    async def create(self, name: str = "", scope: str = DEFAULT_SCOPE) -> MintedToken:
        """Mint a new token. The raw token is returned ONCE and never persisted.

        ``scope`` decides which tools the bearer can invoke. Unknown
        values fall back to ``full`` via ``_normalize_scope``, so the API
        layer can safely pass user input through.
        """
        scope = _normalize_scope(scope)
        token_id = secrets.token_hex(_ID_LEN // 2)
        secret = secrets.token_hex(_SECRET_LEN // 2)
        raw = f"{_RAW_PREFIX}{token_id}{secret}"
        assert len(raw) == _TOTAL_LEN  # sanity
        created_at = datetime.now(UTC).isoformat()
        await self._db.execute(
            "INSERT INTO mcp_tokens (id, secret_hash, name, scope, created_at) VALUES (?, ?, ?, ?, ?)",
            (token_id, _hash_secret(secret), name or "", scope, created_at),
        )
        await self._db.commit()
        return MintedToken(
            token_id=token_id,
            raw_token=raw,
            name=name or "",
            scope=scope,
            created_at=created_at,
        )

    async def list(self) -> list[dict[str, Any]]:
        """List tokens (public fields only — never exposes hashes or raw secrets)."""
        cursor = await self._db.execute(
            """
            SELECT id, name, scope, created_at, last_used_at, revoked, revoked_at
            FROM mcp_tokens
            ORDER BY created_at DESC
            """,
            # NOTE: raw string ordering, NOT SQLite's datetime() function —
            # the latter truncates to second precision and loses ordering
            # for tokens minted in quick succession. Our ISO-8601 UTC
            # timestamps (with microseconds + fixed +00:00 suffix) sort
            # correctly lexicographically.
        )
        rows = await cursor.fetchall()
        return [
            {
                "id": row["id"],
                "name": row["name"],
                "scope": _normalize_scope(row["scope"]),
                "created_at": row["created_at"],
                "last_used_at": row["last_used_at"],
                "revoked": bool(row["revoked"]),
                "revoked_at": row["revoked_at"],
                # Convenience field for UIs — the visible prefix is enough
                # to recognize which token is which without leaking secrets.
                "prefix": f"{_RAW_PREFIX}{row['id']}\u2026",
            }
            for row in rows
        ]

    async def revoke(self, token_id: str) -> bool:
        """Revoke a token by id. Returns True if a live token was revoked."""
        now = datetime.now(UTC).isoformat()
        cursor = await self._db.execute(
            """
            UPDATE mcp_tokens
            SET revoked = 1, revoked_at = ?
            WHERE id = ? AND revoked = 0
            """,
            (now, token_id),
        )
        await self._db.commit()
        return cursor.rowcount > 0

    async def verify(self, raw_token: str) -> VerifiedToken | None:
        """Verify a raw token. Returns a ``VerifiedToken`` (id + scope) on success.

        Constant-time comparison on the hash prevents timing side channels
        from revealing which tokens exist. Also bumps ``last_used_at``
        on success (best-effort — any DB error is swallowed).

        Callers use ``.scope`` to decide tool authorization. The scope
        value is always one of ``SCOPES`` — unknown stored values coerce
        to ``DEFAULT_SCOPE`` via ``_normalize_scope``.
        """
        parsed = _parse_raw_token(raw_token)
        if parsed is None:
            return None
        token_id, secret = parsed
        cursor = await self._db.execute(
            "SELECT secret_hash, scope, revoked FROM mcp_tokens WHERE id = ?",
            (token_id,),
        )
        row = await cursor.fetchone()
        if row is None:
            return None
        if row["revoked"]:
            return None
        stored_hash = row["secret_hash"]
        provided_hash = _hash_secret(secret)
        if not hmac.compare_digest(stored_hash, provided_hash):
            return None
        # Best-effort touch — never raise into the auth path.
        try:
            await self._db.execute(
                "UPDATE mcp_tokens SET last_used_at = ? WHERE id = ?",
                (datetime.now(UTC).isoformat(), token_id),
            )
            await self._db.commit()
        except Exception:
            log.debug("Failed to update last_used_at for token %s", token_id, exc_info=True)
        return VerifiedToken(token_id=token_id, scope=_normalize_scope(row["scope"]))

    async def any_live(self) -> bool:
        """True if at least one non-revoked token exists. Used to decide
        whether to auto-reject SSE requests (no tokens = SSE disabled)."""
        cursor = await self._db.execute(
            "SELECT 1 FROM mcp_tokens WHERE revoked = 0 LIMIT 1",
        )
        row = await cursor.fetchone()
        return row is not None
