"""
MCP Key Service — Ed25519 capability tokens
============================================
Manages the node's Ed25519 keypair and issues/verifies signed MCP tokens.

Token format:
    <base64url(payload_json)>.<base64url(ed25519_signature)>

Payload fields:
    v, jti, task_id, permission, issued_to, node_id, expires
"""

import os
import json
import hashlib
import base64
import secrets
from datetime import datetime, timezone, timedelta
from pathlib import Path
from typing import Optional

from cryptography.hazmat.primitives.asymmetric.ed25519 import (
    Ed25519PrivateKey, Ed25519PublicKey
)
from cryptography.hazmat.primitives.serialization import (
    Encoding, PublicFormat, PrivateFormat, NoEncryption,
    load_pem_private_key, load_pem_public_key
)

# Key storage directory relative to project root
_IDENTITY_DIR = Path(__file__).parent.parent / "identity"
_PRIVATE_KEY_PATH = _IDENTITY_DIR / "node_private.pem"
_PUBLIC_KEY_PATH = _IDENTITY_DIR / "node_public.pem"

_cached_private_key: Optional[Ed25519PrivateKey] = None
_cached_public_key: Optional[Ed25519PublicKey] = None


def _b64url_encode(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def _b64url_decode(s: str) -> bytes:
    padding = 4 - len(s) % 4
    if padding != 4:
        s += "=" * padding
    return base64.urlsafe_b64decode(s)


def get_or_create_keypair() -> tuple[Ed25519PrivateKey, Ed25519PublicKey]:
    """Load keypair from disk, generating it if it doesn't exist."""
    global _cached_private_key, _cached_public_key

    if _cached_private_key is not None:
        return _cached_private_key, _cached_public_key

    _IDENTITY_DIR.mkdir(parents=True, exist_ok=True)

    if _PRIVATE_KEY_PATH.exists():
        pem = _PRIVATE_KEY_PATH.read_bytes()
        private_key = load_pem_private_key(pem, password=None)
    else:
        private_key = Ed25519PrivateKey.generate()
        pem = private_key.private_bytes(Encoding.PEM, PrivateFormat.PKCS8, NoEncryption())
        _PRIVATE_KEY_PATH.write_bytes(pem)
        _PRIVATE_KEY_PATH.chmod(0o600)
        # Auto-trust admin on first run (if admin_public.pem exists in package)
        try:
            auto_trust_admin()
        except Exception:
            pass

    public_key = private_key.public_key()
    pub_pem = public_key.public_bytes(Encoding.PEM, PublicFormat.SubjectPublicKeyInfo)
    _PUBLIC_KEY_PATH.write_bytes(pub_pem)

    _cached_private_key = private_key
    _cached_public_key = public_key
    return private_key, public_key


def get_public_key_b64() -> str:
    """Return the node's public key as base64url-encoded raw bytes."""
    _, public_key = get_or_create_keypair()
    raw = public_key.public_bytes(Encoding.Raw, PublicFormat.Raw)
    return _b64url_encode(raw)


def get_node_id() -> str:
    """Return first 16 hex chars of SHA256(public_key_raw_bytes)."""
    _, public_key = get_or_create_keypair()
    raw = public_key.public_bytes(Encoding.Raw, PublicFormat.Raw)
    return hashlib.sha256(raw).hexdigest()[:16]


def generate_mcp_key(
    permission: str,
    issued_to: str,
    task_id: Optional[str] = None,
    expires_days: int = 30,
) -> tuple[str, str]:
    """
    Create a signed MCP capability token.

    Returns (jti, token_str) where token_str is the full signed token.
    """
    private_key, _ = get_or_create_keypair()
    node_id = get_node_id()
    jti = secrets.token_hex(16)

    expires_dt = datetime.now(timezone.utc) + timedelta(days=expires_days)

    payload = {
        "v": 1,
        "jti": jti,
        "task_id": task_id,
        "permission": permission,
        "issued_to": issued_to,
        "node_id": node_id,
        "expires": expires_dt.isoformat(),
    }

    payload_bytes = json.dumps(payload, separators=(",", ":")).encode()
    payload_b64 = _b64url_encode(payload_bytes)

    signature = private_key.sign(payload_bytes)
    sig_b64 = _b64url_encode(signature)

    token_str = f"{payload_b64}.{sig_b64}"
    return jti, token_str


def verify_mcp_key(token_str: str) -> Optional[dict]:
    """
    Verify an MCP token signature and expiry.

    Returns the decoded payload dict on success, or None on failure.
    Automatically loads the correct public key from trusted_nodes if the
    token's node_id differs from our own.
    """
    try:
        parts = token_str.strip().split(".")
        if len(parts) != 2:
            return None

        payload_b64, sig_b64 = parts
        payload_bytes = _b64url_decode(payload_b64)
        signature = _b64url_decode(sig_b64)
        payload = json.loads(payload_bytes)

        token_node_id = payload.get("node_id", "")
        if token_node_id == get_node_id():
            _, public_key = get_or_create_keypair()
        else:
            # Look up the peer's public key
            public_key = _load_peer_public_key(token_node_id)
            if public_key is None:
                return None

        public_key.verify(signature, payload_bytes)  # raises on bad sig

        # Check expiry
        expires_str = payload.get("expires")
        if expires_str:
            expires_dt = datetime.fromisoformat(expires_str)
            if datetime.now(timezone.utc) > expires_dt:
                return None

        return payload

    except Exception:
        return None


def _load_peer_public_key(node_id: str) -> Optional[Ed25519PublicKey]:
    """Fetch a peer's raw public key from trusted_nodes DB and return an Ed25519PublicKey."""
    try:
        import importlib.util, sys
        db_path = Path(__file__).parent.parent / "memory" / "database.py"
        spec = importlib.util.spec_from_file_location("_db", db_path)
        db = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(db)
        node = db.get_trusted_node(node_id)
        if not node:
            return None
        raw = _b64url_decode(node["public_key_b64"])
        return Ed25519PublicKey.from_public_bytes(raw)
    except Exception:
        return None


def auto_trust_admin():
    """
    Auto-import the admin node's public key if admin_public.pem exists.
    This key is embedded during package build so customer nodes automatically
    trust the admin for MCP remote control and support.
    Called once during node startup.
    """
    admin_pem = _IDENTITY_DIR / "admin_public.pem"
    if not admin_pem.exists():
        return False

    try:
        # Load admin's public key
        pem_bytes = admin_pem.read_bytes()
        admin_pub = load_pem_public_key(pem_bytes)
        raw = admin_pub.public_bytes(Encoding.Raw, PublicFormat.Raw)
        admin_pub_b64 = _b64url_encode(raw)
        admin_node_id = hashlib.sha256(raw).hexdigest()[:16]

        # Don't import our own key
        if admin_node_id == get_node_id():
            return False

        # Import into trusted_nodes DB
        try:
            import importlib.util
            db_path = Path(__file__).parent.parent / "memory" / "database.py"
            spec = importlib.util.spec_from_file_location("_db", db_path)
            db = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(db)

            db.register_trusted_node(
                node_id=admin_node_id,
                public_key_b64=admin_pub_b64,
                display_name="BizNode-Admin",
                endpoint="",
                capabilities='["admin","mcp_auth","remote_control"]',
            )
            return True
        except Exception:
            return False
    except Exception:
        return False
