"""
Brynq Per-User API Key Manager (BYOK — Bring Your Own Keys)

Encrypts and stores per-user API keys for LLM providers using Fernet
(via vault.py's encryption primitives). Each user gets a separate
encrypted JSON file in state/broker/keys/{user_id}.json.

Keys are encrypted at rest. Never log or return unmasked keys.

Usage:
    from brynq.key_manager import store_user_key, get_user_key, list_user_keys, delete_user_key

    store_user_key("user123", "claude", "sk-ant-api03-xxxx")
    key = get_user_key("user123", "claude")  # returns plaintext key
    keys = list_user_keys("user123")          # returns masked entries
    delete_user_key("user123", "claude")
"""

import json
import logging
import re
from pathlib import Path
from typing import Optional

from brynq.vault import _encrypt, _decrypt

log = logging.getLogger("brynq.key_manager")

# ── Configuration ────────────────────────────────────────────────────────

KEYS_DIR = Path("state/broker/keys")

# Providers we accept keys for. Extensible — add new providers here.
VALID_PROVIDERS = frozenset({
    "claude",
    "openai",
    "gemini",
    "cohere",
    "mistral",
    "groq",
    "deepseek",
    "perplexity",
    "together",
    "fireworks",
    "replicate",
    "huggingface",
})

# Map provider names to the env var the llm_router expects
PROVIDER_ENV_MAP = {
    "claude": "ANTHROPIC_API_KEY",
    "openai": "OPENAI_API_KEY",
    "gemini": "GEMINI_API_KEY",
}


def _ensure_keys_dir() -> None:
    """Create the keys directory if it doesn't exist."""
    KEYS_DIR.mkdir(parents=True, exist_ok=True)


def _user_key_path(user_id: str) -> Path:
    """Return the path to a user's encrypted key file."""
    # Sanitize user_id to prevent path traversal
    safe_id = re.sub(r"[^a-zA-Z0-9_-]", "", user_id)
    if not safe_id:
        raise ValueError("Invalid user_id")
    return KEYS_DIR / f"{safe_id}.json"


def _load_user_keys(user_id: str) -> dict[str, str]:
    """Load a user's encrypted key store from disk.

    Returns a dict of {provider: encrypted_key_string}.
    """
    path = _user_key_path(user_id)
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text("utf-8"))
        return data.get("keys", {})
    except (json.JSONDecodeError, OSError) as exc:
        log.warning("Failed to load keys for user %s: %s", user_id, exc)
        return {}


def _save_user_keys(user_id: str, keys: dict[str, str]) -> None:
    """Save a user's encrypted key store to disk."""
    _ensure_keys_dir()
    path = _user_key_path(user_id)
    path.write_text(
        json.dumps({"keys": keys}, indent=2) + "\n",
        encoding="utf-8",
    )


def _mask_key(api_key: str) -> str:
    """Mask an API key for safe display.

    Examples:
        "sk-ant-api03-abcdef...xyz" -> "sk-a...xyz"
        "short" -> "sho...ort"
    """
    if len(api_key) <= 8:
        return api_key[:3] + "..." + api_key[-3:] if len(api_key) > 6 else "***"
    prefix = api_key[:4]
    suffix = api_key[-4:]
    return f"{prefix}...{suffix}"


# ── Public API ───────────────────────────────────────────────────────────


def store_user_key(user_id: str, provider: str, api_key: str) -> dict:
    """Encrypt and store an API key for a user + provider.

    Args:
        user_id: The user's unique identifier.
        provider: LLM provider name (e.g. "claude", "openai", "gemini").
        api_key: The raw API key to store.

    Returns:
        {"ok": True, "provider": provider, "masked_key": "sk-a...xxxx"}
    """
    provider = provider.strip().lower()
    if provider not in VALID_PROVIDERS:
        return {
            "ok": False,
            "error": f"Unknown provider '{provider}'. Valid: {sorted(VALID_PROVIDERS)}",
        }

    api_key = api_key.strip()
    if not api_key:
        return {"ok": False, "error": "API key cannot be empty"}

    if len(api_key) < 8:
        return {"ok": False, "error": "API key seems too short (minimum 8 characters)"}

    keys = _load_user_keys(user_id)
    keys[provider] = _encrypt(api_key)
    _save_user_keys(user_id, keys)

    log.info("Stored %s key for user %s", provider, user_id)
    return {
        "ok": True,
        "provider": provider,
        "masked_key": _mask_key(api_key),
    }


def get_user_key(user_id: str, provider: str) -> Optional[str]:
    """Decrypt and return a user's API key for a provider.

    Args:
        user_id: The user's unique identifier.
        provider: LLM provider name.

    Returns:
        The decrypted API key string, or None if not found.
    """
    provider = provider.strip().lower()
    keys = _load_user_keys(user_id)
    encrypted = keys.get(provider)
    if not encrypted:
        return None

    decrypted = _decrypt(encrypted)
    if not decrypted:
        log.warning("Failed to decrypt %s key for user %s", provider, user_id)
        return None

    return decrypted


def list_user_keys(user_id: str) -> list[dict]:
    """List all stored providers for a user with masked keys.

    Returns:
        List of {"provider": "claude", "masked_key": "sk-a...xxxx"}
    """
    keys = _load_user_keys(user_id)
    result = []
    for provider, encrypted in sorted(keys.items()):
        decrypted = _decrypt(encrypted)
        masked = _mask_key(decrypted) if decrypted else "***"
        result.append({
            "provider": provider,
            "masked_key": masked,
        })
    return result


def delete_user_key(user_id: str, provider: str) -> dict:
    """Remove a stored API key for a user + provider.

    Args:
        user_id: The user's unique identifier.
        provider: LLM provider name.

    Returns:
        {"ok": True} on success, {"ok": False, "error": ...} on failure.
    """
    provider = provider.strip().lower()
    keys = _load_user_keys(user_id)
    if provider not in keys:
        return {"ok": False, "error": f"No key stored for provider '{provider}'"}

    del keys[provider]
    _save_user_keys(user_id, keys)

    log.info("Deleted %s key for user %s", provider, user_id)
    return {"ok": True, "provider": provider}
