"""
Encrypted local token storage for OAuth credentials.

Uses the same Fernet (AES-128-CBC + HMAC-SHA256) encryption pattern as the
KeyVault, with machine-bound PBKDF2 key derivation. Tokens are stored at
``~/.brynq/tokens.json`` and are NEVER sent to Brynq cloud.

Storage format (after decryption)::

    {
        "anthropic": {
            "access_token": "...",
            "refresh_token": "...",
            "expires_at": 1711234567,
            "scopes": ["user:read", "messages:write"],
            "user_email": "user@example.com"
        },
        ...
    }
"""

from __future__ import annotations

import base64
import hashlib
import json
import logging
import os
import secrets
import socket
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional

from cryptography.fernet import Fernet, InvalidToken

log = logging.getLogger("brynq.runtime.auth.token_store")

# ── Constants ────────────────────────────────────────────────────────────

TOKENS_FILE = Path.home() / ".brynq" / "tokens.json"

_KDF_ITERATIONS = 480_000  # OWASP 2023 recommendation for PBKDF2-SHA256
_SALT_LEN = 32  # 256-bit salt
_KEY_LEN = 32  # 256-bit derived key

# Tokens expiring within this window are considered "about to expire"
_EXPIRY_BUFFER_SECONDS = 60


# ── Data Classes ─────────────────────────────────────────────────────────


@dataclass
class TokenSet:
    """Token data for a single provider."""

    access_token: str
    refresh_token: Optional[str] = None
    expires_at: Optional[float] = None  # Unix timestamp
    created_at: Optional[float] = None  # Unix timestamp of token creation
    last_used: Optional[float] = None  # Unix timestamp of last retrieval
    scopes: list[str] = field(default_factory=list)
    user_email: Optional[str] = None

    @property
    def is_expired(self) -> bool:
        """True if the access token has expired (or is about to).

        Also treats the token as expired if ``created_at`` is in the future,
        which indicates clock skew or a tampered token.
        """
        if self.created_at is not None and self.created_at > time.time():
            log.warning(
                "Token created_at (%s) is in the future — treating as expired",
                self.created_at,
            )
            return True
        if self.expires_at is None:
            return False  # No expiry info = assume valid
        return time.time() >= (self.expires_at - _EXPIRY_BUFFER_SECONDS)

    def to_dict(self) -> dict[str, Any]:
        """Serialize to a JSON-compatible dict."""
        d: dict[str, Any] = {"access_token": self.access_token}
        if self.refresh_token is not None:
            d["refresh_token"] = self.refresh_token
        if self.expires_at is not None:
            d["expires_at"] = self.expires_at
        if self.created_at is not None:
            d["created_at"] = self.created_at
        if self.last_used is not None:
            d["last_used"] = self.last_used
        if self.scopes:
            d["scopes"] = self.scopes
        if self.user_email is not None:
            d["user_email"] = self.user_email
        return d

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> TokenSet:
        """Deserialize from a dict."""
        return cls(
            access_token=data["access_token"],
            refresh_token=data.get("refresh_token"),
            expires_at=data.get("expires_at"),
            created_at=data.get("created_at"),
            last_used=data.get("last_used"),
            scopes=data.get("scopes", []),
            user_email=data.get("user_email"),
        )


# ── Machine Fingerprint (same pattern as KeyVault) ─────────────────────


def _get_machine_id() -> bytes:
    """Derive a stable machine identifier from hostname + MAC address."""
    hostname = socket.gethostname()
    mac = uuid.getnode()
    mac_hex = format(mac, "012x")
    raw = f"brynq-tokens-v1:{hostname}:{mac_hex}"
    return hashlib.sha256(raw.encode("utf-8")).digest()


def _derive_fernet_key(salt: bytes) -> bytes:
    """Derive a Fernet-compatible key from machine ID.

    Returns url-safe base64-encoded 32-byte key as Fernet requires.
    """
    machine_id = _get_machine_id()
    derived = hashlib.pbkdf2_hmac(
        "sha256",
        machine_id,
        salt,
        iterations=_KDF_ITERATIONS,
        dklen=_KEY_LEN,
    )
    return base64.urlsafe_b64encode(derived)


# ── Token Store ──────────────────────────────────────────────────────────


class TokenStore:
    """Encrypted local storage for OAuth token sets.

    Tokens are encrypted at rest with machine-specific Fernet encryption.
    Tokens NEVER leave the user's machine.

    Usage::

        store = TokenStore()
        store.save_tokens("anthropic", token_set)
        ts = store.get_tokens("anthropic")
        if ts and ts.is_expired:
            ...  # trigger refresh
    """

    def __init__(self, tokens_path: Optional[Path] = None) -> None:
        self._tokens_path = tokens_path or TOKENS_FILE

    def save_tokens(self, provider: str, tokens: TokenSet) -> None:
        """Encrypt and persist a token set for a provider.

        Args:
            provider: Provider name (e.g. "anthropic", "google", "openai").
            tokens: The token set to store.

        Raises:
            ValueError: If provider is empty.
        """
        provider = provider.strip().lower()
        if not provider:
            raise ValueError("Provider name cannot be empty")

        envelope = self._load_envelope()
        all_tokens = self._decrypt_payload(envelope)

        all_tokens[provider] = tokens.to_dict()

        self._encrypt_and_save(envelope, all_tokens)
        # NEVER log token values
        log.info("Saved OAuth tokens for provider: %s", provider)

    def get_tokens(self, provider: str) -> Optional[TokenSet]:
        """Retrieve and decrypt a token set for a provider.

        Returns:
            TokenSet if found, None if not stored or decryption fails.
        """
        provider = provider.strip().lower()
        envelope = self._load_envelope()
        all_tokens = self._decrypt_payload(envelope)

        data = all_tokens.get(provider)
        if data is None:
            return None

        try:
            token_set = TokenSet.from_dict(data)
        except (KeyError, TypeError) as exc:
            log.warning("Corrupt token data for %s: %s", provider, exc)
            return None

        # Update last_used timestamp and persist
        token_set.last_used = time.time()
        all_tokens[provider] = token_set.to_dict()
        self._encrypt_and_save(envelope, all_tokens)

        return token_set

    def delete_tokens(self, provider: str) -> bool:
        """Remove stored tokens for a provider.

        Returns:
            True if tokens were deleted, False if provider was not found.
        """
        provider = provider.strip().lower()
        envelope = self._load_envelope()
        all_tokens = self._decrypt_payload(envelope)

        if provider not in all_tokens:
            return False

        del all_tokens[provider]
        self._encrypt_and_save(envelope, all_tokens)
        log.info("Deleted OAuth tokens for provider: %s", provider)
        return True

    def list_providers(self) -> list[str]:
        """Return a sorted list of providers that have stored tokens."""
        envelope = self._load_envelope()
        all_tokens = self._decrypt_payload(envelope)
        return sorted(all_tokens.keys())

    def has_tokens(self, provider: str) -> bool:
        """Check if tokens exist for a provider (without decrypting them
        into a TokenSet)."""
        provider = provider.strip().lower()
        envelope = self._load_envelope()
        all_tokens = self._decrypt_payload(envelope)
        return provider in all_tokens

    def clear_all(self) -> int:
        """Remove all stored tokens.

        Returns:
            The number of providers whose tokens were cleared.
        """
        envelope = self._load_envelope()
        all_tokens = self._decrypt_payload(envelope)
        count = len(all_tokens)

        if count > 0:
            self._encrypt_and_save(envelope, {})

        log.info("Cleared all stored tokens (count: %d)", count)
        return count

    # ── Private Helpers ──────────────────────────────────────────────

    def _load_envelope(self) -> dict[str, Any]:
        """Load the envelope from disk, creating a new one if needed."""
        if not self._tokens_path.exists():
            return self._create_empty_envelope()

        try:
            raw = self._tokens_path.read_text("utf-8")
            data = json.loads(raw)
        except (json.JSONDecodeError, OSError):
            log.warning("Corrupted tokens file -- creating new one")
            return self._create_empty_envelope()

        if "version" not in data or "salt" not in data or "payload" not in data:
            log.warning("Invalid tokens structure -- creating new one")
            return self._create_empty_envelope()

        return data

    def _decrypt_payload(self, envelope: dict[str, Any]) -> dict[str, Any]:
        """Decrypt the payload from an envelope. Returns empty dict on failure."""
        payload_b64 = envelope.get("payload")
        if not payload_b64:
            return {}

        try:
            salt = base64.b64decode(envelope["salt"])
            fernet_key = _derive_fernet_key(salt)
            f = Fernet(fernet_key)
            encrypted = base64.b64decode(payload_b64)
            decrypted = f.decrypt(encrypted)
            return json.loads(decrypted.decode("utf-8"))
        except (InvalidToken, json.JSONDecodeError, ValueError, KeyError) as exc:
            log.warning("Vault corruption detected — failed to decrypt tokens: %s", exc)
            return {}
        except Exception as exc:
            log.warning("Unexpected error decrypting token vault: %s", exc)
            return {}

    def _encrypt_and_save(
        self, envelope: dict[str, Any], payload: dict[str, Any]
    ) -> None:
        """Encrypt payload and write the envelope to disk."""
        salt = base64.b64decode(envelope["salt"])
        fernet_key = _derive_fernet_key(salt)
        f = Fernet(fernet_key)

        plaintext = json.dumps(payload).encode("utf-8")
        encrypted = f.encrypt(plaintext)
        envelope["payload"] = base64.b64encode(encrypted).decode("utf-8")

        self._save_envelope(envelope)

    def _save_envelope(self, envelope: dict[str, Any]) -> None:
        """Write the envelope to disk with restrictive permissions."""
        self._tokens_path.parent.mkdir(parents=True, exist_ok=True)
        content = json.dumps(envelope, indent=2) + "\n"
        self._tokens_path.write_text(content, encoding="utf-8")

        try:
            if os.name != "nt":
                self._tokens_path.chmod(0o600)
        except OSError:
            pass

    def _create_empty_envelope(self) -> dict[str, Any]:
        """Create and persist a new empty envelope with a fresh salt."""
        salt = secrets.token_bytes(_SALT_LEN)
        envelope: dict[str, Any] = {
            "version": 1,
            "salt": base64.b64encode(salt).decode("utf-8"),
            "payload": "",
        }
        self._save_envelope(envelope)
        return envelope
