"""
Encrypted Credential Vault

AES-256-GCM encrypted local credential store.
Credentials arrive via WSS from the frontend and are stored
on the user's machine — never on the backend.
"""

import json
import os
import hashlib
import logging
from dataclasses import dataclass, asdict
from typing import Optional, Dict, Any
from pathlib import Path

logger = logging.getLogger(__name__)

try:
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
    from cryptography.hazmat.primitives import hashes
    HAS_CRYPTO = True
except ImportError:
    HAS_CRYPTO = False
    logger.warning("cryptography not installed — vault will store credentials in plaintext")


@dataclass
class StoredCredential:
    """A single credential entry in the vault."""
    service_slug: str
    auth_type: str               # "bearer", "apikey", "basic", "custom"
    header_name: str             # e.g. "X-Shopify-Access-Token"
    header_value_pattern: str    # e.g. "Bearer {value}" or just "{value}"
    secret_value: str            # The actual API key/token
    query_param_name: str = ""   # For query param injection
    username: str = ""           # For basic auth
    updated_at: str = ""         # ISO timestamp


class CredentialVault:
    """
    Encrypted credential store on the workspace machine.
    
    Credentials flow:
    1. User enters in frontend
    2. Frontend sends via WSS → backend relays (pass-through) → workspace agent
    3. Workspace agent encrypts + stores here
    4. On tool execution, vault resolves credentials locally
    """

    def __init__(self, vault_dir: str, encryption_key: Optional[str] = None):
        """
        Args:
            vault_dir: Directory to store the vault file (e.g., /workspace/.vault)
            encryption_key: Key for AES-256-GCM encryption. If None, derived from TOKEN env var.
        """
        self._vault_dir = Path(vault_dir)
        self._vault_dir.mkdir(parents=True, exist_ok=True)
        self._vault_file = self._vault_dir / "credentials.enc"
        self._credentials: Dict[str, StoredCredential] = {}
        self._key = self._derive_key(encryption_key or os.environ.get("TOKEN", "default-key"))
        self._load()

    def _derive_key(self, passphrase: str) -> bytes:
        """Derive AES-256 key from passphrase using PBKDF2."""
        if not HAS_CRYPTO:
            return b""
        salt = b"workspace-agent-vault-v1"  # Fixed salt — key is per-token anyway
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100_000,
        )
        return kdf.derive(passphrase.encode("utf-8"))

    def _load(self):
        """Load and decrypt vault from disk."""
        if not self._vault_file.exists():
            self._credentials = {}
            return

        try:
            raw = self._vault_file.read_bytes()
            if HAS_CRYPTO and len(raw) > 12:
                nonce = raw[:12]
                ciphertext = raw[12:]
                aesgcm = AESGCM(self._key)
                plaintext = aesgcm.decrypt(nonce, ciphertext, None)
                data = json.loads(plaintext.decode("utf-8"))
            else:
                data = json.loads(raw.decode("utf-8"))

            self._credentials = {
                slug: StoredCredential(**cred)
                for slug, cred in data.items()
            }
            logger.info(f"Vault loaded: {len(self._credentials)} credentials")
        except Exception as e:
            logger.error(f"Failed to load vault: {e}")
            self._credentials = {}

    def _save(self):
        """Encrypt and write vault to disk."""
        try:
            data = {
                slug: asdict(cred)
                for slug, cred in self._credentials.items()
            }
            plaintext = json.dumps(data).encode("utf-8")

            if HAS_CRYPTO:
                nonce = os.urandom(12)
                aesgcm = AESGCM(self._key)
                ciphertext = aesgcm.encrypt(nonce, plaintext, None)
                self._vault_file.write_bytes(nonce + ciphertext)
            else:
                self._vault_file.write_bytes(plaintext)

            logger.info(f"Vault saved: {len(self._credentials)} credentials")
        except Exception as e:
            logger.error(f"Failed to save vault: {e}")

    def store(self, credential: StoredCredential):
        """Store or update a credential."""
        from datetime import datetime, timezone
        credential.updated_at = datetime.now(timezone.utc).isoformat()
        self._credentials[credential.service_slug] = credential
        self._save()
        logger.info(f"Credential stored: {credential.service_slug} ({credential.auth_type})")

    def get(self, service_slug: str) -> Optional[StoredCredential]:
        """Retrieve a credential by service slug."""
        return self._credentials.get(service_slug)

    def delete(self, service_slug: str) -> bool:
        """Delete a credential."""
        if service_slug in self._credentials:
            del self._credentials[service_slug]
            self._save()
            logger.info(f"Credential deleted: {service_slug}")
            return True
        return False

    def list_services(self) -> list:
        """List all stored service slugs (no secrets)."""
        return [
            {
                "service_slug": cred.service_slug,
                "auth_type": cred.auth_type,
                "updated_at": cred.updated_at,
            }
            for cred in self._credentials.values()
        ]

    def resolve_headers(self, service_slug: str) -> Dict[str, str]:
        """
        Resolve credential into HTTP headers ready for injection.
        
        This is what the http_executor calls before firing a service request.
        """
        cred = self.get(service_slug)
        if not cred:
            logger.warning(f"No credential found for service: {service_slug}")
            return {}

        if cred.auth_type == "bearer":
            return {cred.header_name: f"Bearer {cred.secret_value}"}
        elif cred.auth_type == "apikey":
            return {cred.header_name: cred.secret_value}
        elif cred.auth_type == "basic":
            import base64
            auth_string = f"{cred.username}:{cred.secret_value}"
            encoded = base64.b64encode(auth_string.encode()).decode()
            return {"Authorization": f"Basic {encoded}"}
        elif cred.auth_type == "custom":
            value = cred.header_value_pattern.replace("{value}", cred.secret_value)
            return {cred.header_name: value}
        else:
            return {}

    def resolve_query_params(self, service_slug: str) -> Dict[str, str]:
        """Resolve credential into query parameters."""
        cred = self.get(service_slug)
        if not cred or not cred.query_param_name:
            return {}
        return {cred.query_param_name: cred.secret_value}
