# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Licensed under the MIT License

"""
Mesh Trust — Phase 2

User-facing trust management layer on top of the existing
SecureSessionManager (core/secure_transport.py, 1324 LOC).

Does NOT implement crypto. All crypto is in secure_transport.py:
  - X3DH key exchange (forward secrecy)
  - Double Ratchet (per-message keys, post-compromise recovery)
  - SecureSessionManager (key persistence, session lifecycle)

This module adds:
  - TrustRecord: per-peer trust state with identity fingerprint
  - PermissionMatrix: granular per-peer capability control
  - Verification codes: 6-digit codes from identity key material
  - Trust lifecycle: unknown → pending → trusted | blocked
  - Integration with PeerStore (Phase 1)
"""

from __future__ import annotations

import hashlib
import json
import logging
import threading
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# Try to import secure transport — graceful if pynacl not installed
try:
    from ..secure_transport import (
        NACL_AVAILABLE,
        PreKeyBundle,
        SecureSessionManager,
    )

    HAS_SECURE_TRANSPORT = NACL_AVAILABLE
except ImportError:
    HAS_SECURE_TRANSPORT = False
    logger.debug("secure_transport not available — trust requires pynacl")


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


# ============================================================
# PERMISSION MATRIX
# ============================================================


@dataclass
class PermissionMatrix:
    """
    Per-peer permission controls.

    After trusting a peer's identity (verification code match),
    the user grants specific capabilities. Default is all-deny
    except list_skills.
    """

    list_skills: bool = True  # Can see our skill names
    invoke_skills: bool = False  # Can call our skills remotely
    request_memory: bool = False  # Can query our shared memories (Phase 3)
    delegate_tasks: bool = False  # Can send delegation requests (Phase 4)
    share_memory_to: bool = False  # We can push memories to them (Phase 3)
    share_markdown_memories: bool = False  # Can exchange markdown memories
    share_models: bool = False  # Can exchange .familiar-model packages
    send_messages: bool = False  # Can send peer-to-peer chat messages

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "PermissionMatrix":
        known = {f.name for f in cls.__dataclass_fields__.values()}
        return cls(**{k: v for k, v in d.items() if k in known})

    def grant(self, permission: str) -> bool:
        """Grant a permission. Returns True if valid permission name."""
        if permission in self.__dataclass_fields__:
            setattr(self, permission, True)
            return True
        return False

    def revoke(self, permission: str) -> bool:
        """Revoke a permission. Returns True if valid permission name."""
        if permission in self.__dataclass_fields__:
            setattr(self, permission, False)
            return True
        return False

    def check(self, permission: str) -> bool:
        """Check if a permission is granted."""
        if permission not in self.__dataclass_fields__:
            return False
        return getattr(self, permission, False)

    @classmethod
    def all_deny(cls) -> "PermissionMatrix":
        """All permissions denied (HIPAA default)."""
        return cls(
            list_skills=False,
            invoke_skills=False,
            request_memory=False,
            delegate_tasks=False,
            share_memory_to=False,
            share_markdown_memories=False,
            share_models=False,
            send_messages=False,
        )

    @classmethod
    def all_grant(cls) -> "PermissionMatrix":
        """All permissions granted (convenience, not recommended)."""
        return cls(
            list_skills=True,
            invoke_skills=True,
            request_memory=True,
            delegate_tasks=True,
            share_memory_to=True,
            share_markdown_memories=True,
            share_models=True,
            send_messages=True,
        )

    @property
    def granted_list(self) -> List[str]:
        """List of granted permission names."""
        return [name for name in self.__dataclass_fields__ if getattr(self, name, False)]


# ============================================================
# TRUST RECORD
# ============================================================


@dataclass
class TrustRecord:
    """
    Trust state for a single peer.

    Tracks identity fingerprint, trust level, permissions,
    and the verification code used during establishment.
    """

    node_id: str
    node_name: str
    identity_fingerprint: str  # SHA-256 of their Ed25519 identity key
    trust_level: str = "unknown"  # 'unknown', 'pending', 'trusted', 'blocked'
    permissions: PermissionMatrix = field(default_factory=PermissionMatrix)
    trusted_at: Optional[str] = None
    blocked_at: Optional[str] = None
    verification_code: str = ""  # 6-digit code for this pair
    last_verified: Optional[str] = None

    def to_dict(self) -> dict:
        d = asdict(self)
        d["permissions"] = self.permissions.to_dict()
        return d

    @classmethod
    def from_dict(cls, d: dict) -> "TrustRecord":
        d = d.copy()
        if "permissions" in d and isinstance(d["permissions"], dict):
            d["permissions"] = PermissionMatrix.from_dict(d["permissions"])
        known = {f.name for f in cls.__dataclass_fields__.values()}
        return cls(**{k: v for k, v in d.items() if k in known})


# ============================================================
# VERIFICATION CODE
# ============================================================


def compute_verification_code(our_identity: bytes, their_identity: bytes) -> str:
    """
    Compute a 6-digit verification code from two identity public keys.

    Deterministic: same keys always produce the same code.
    Order-independent: A↔B produces the same code as B↔A.

    Both users compare this code out-of-band (phone, in person, etc.)
    to confirm they're talking to the right peer, not a MITM.
    """
    # Sort keys so A→B and B→A produce the same code
    sorted_keys = b"".join(sorted([our_identity, their_identity]))
    digest = hashlib.sha256(sorted_keys).hexdigest()
    # Take first 6 decimal digits from hex digest
    numeric = int(digest[:8], 16) % 1000000
    return f"{numeric:06d}"


def format_verification_code(code: str) -> str:
    """Format as '847 291' for easier reading."""
    return f"{code[:3]} {code[3:]}"


def compute_fingerprint(identity_key: bytes) -> str:
    """Compute SHA-256 fingerprint of an identity public key."""
    return hashlib.sha256(identity_key).hexdigest()


# ============================================================
# TRUST STORE — persistence
# ============================================================


class TrustStore:
    """
    Persists TrustRecords to disk.
    Thread-safe. Stored in MESH_DIR/trust.json.
    """

    def __init__(self, storage_path: Path):
        self._path = storage_path / "trust.json"
        self._lock = threading.Lock()
        self._records: Dict[str, TrustRecord] = {}
        self._load()

    def _load(self):
        if self._path.exists():
            try:
                data = json.loads(self._path.read_text())
                for entry in data.get("records", []):
                    rec = TrustRecord.from_dict(entry)
                    self._records[rec.node_id] = rec
            except (json.JSONDecodeError, IOError, OSError) as e:
                logger.warning(f"Failed to load trust store: {e}")

    def _save(self):
        try:
            self._path.parent.mkdir(parents=True, exist_ok=True)
            data = {"records": [r.to_dict() for r in self._records.values()]}
            self._path.write_text(json.dumps(data, indent=2))
        except (IOError, OSError) as e:
            logger.warning(f"Failed to save trust store: {e}")

    def get(self, node_id: str) -> Optional[TrustRecord]:
        with self._lock:
            return self._records.get(node_id)

    def upsert(self, record: TrustRecord):
        with self._lock:
            self._records[record.node_id] = record
            self._save()

    def remove(self, node_id: str) -> bool:
        with self._lock:
            if node_id in self._records:
                del self._records[node_id]
                self._save()
                return True
            return False

    def get_all(self) -> List[TrustRecord]:
        with self._lock:
            return list(self._records.values())

    def get_trusted(self) -> List[TrustRecord]:
        with self._lock:
            return [r for r in self._records.values() if r.trust_level == "trusted"]

    def count(self) -> int:
        with self._lock:
            return len(self._records)


# ============================================================
# MESH TRUST MANAGER — the main class
# ============================================================


class MeshTrustManager:
    """
    Manages trust relationships between mesh peers.

    Wraps SecureSessionManager (existing crypto) with:
    - Trust records and permission matrices
    - Verification code generation
    - User-facing trust lifecycle commands

    Usage:
        trust = MeshTrustManager(mesh_dir=Path("~/.familiar/data/mesh"))

        # On peer connect: establish encrypted session
        init_bytes = trust.initiate_handshake(peer_id, their_prekey_bundle)
        # → send init_bytes to peer

        # On receiving handshake from peer:
        trust.accept_handshake(peer_id, peer_name, init_message_bytes)
        code = trust.get_verification_code(peer_id)
        # → display code to user

        # User confirms:
        trust.approve(peer_id)
        trust.grant_permission(peer_id, "invoke_skills")

        # Check permissions before allowing actions:
        if trust.check_permission(peer_id, "delegate_tasks"):
            handle_delegation(...)
    """

    def __init__(self, mesh_dir: Optional[Path] = None):
        if mesh_dir is None:
            mesh_dir = Path.home() / ".familiar" / "data" / "mesh"

        self._mesh_dir = mesh_dir
        self._trust_store = TrustStore(mesh_dir)

        # Secure transport (crypto layer)
        self._session_manager: Optional[Any] = None
        if HAS_SECURE_TRANSPORT:
            secure_dir = mesh_dir / "secure"
            self._session_manager = SecureSessionManager(storage_path=secure_dir)
        else:
            logger.warning(
                "SecureSessionManager unavailable (pynacl not installed). "
                "Mesh trust will work without encryption."
            )

    @property
    def has_crypto(self) -> bool:
        """Whether encrypted sessions are available."""
        return self._session_manager is not None

    def get_our_fingerprint(self) -> str:
        """Get our identity key fingerprint (for mDNS advertisement)."""
        if not self._session_manager:
            return ""
        identity_key = self._session_manager.local_keys.identity_key.public
        return compute_fingerprint(identity_key)[:16]

    def get_prekey_bundle_dict(self) -> Optional[dict]:
        """Get our PreKeyBundle as a dict for PEER_HELLO payloads."""
        if not self._session_manager:
            return None
        bundle = self._session_manager.get_prekey_bundle()
        # Serialize to a dict for JSON transport (not raw bytes)
        return {
            "ik": bundle.identity_key.hex(),
            "spk": bundle.signed_prekey.hex(),
            "spk_sig": bundle.signed_prekey_signature.hex(),
            "spk_id": bundle.signed_prekey_id,
            "opk": bundle.one_time_prekey.hex() if bundle.one_time_prekey else None,
            "opk_id": bundle.one_time_prekey_id,
        }

    # ── Handshake ──

    def initiate_handshake(
        self, peer_id: str, peer_name: str, their_bundle_dict: dict
    ) -> Optional[str]:
        """
        Initiate a secure session with a peer.

        Args:
            peer_id: Remote node's ID
            peer_name: Remote node's display name
            their_bundle_dict: Their PreKeyBundle as a dict (from PEER_HELLO)

        Returns:
            X3DH init message as hex string (to send to peer), or None if no crypto
        """
        if not self._session_manager:
            # No crypto — create trust record without encryption
            their_fp = hashlib.sha256(peer_id.encode()).hexdigest()
            self._create_trust_record(peer_id, peer_name, their_fp, "")
            return None

        # Reconstruct their PreKeyBundle
        bundle = PreKeyBundle(
            identity_key=bytes.fromhex(their_bundle_dict["ik"]),
            signed_prekey=bytes.fromhex(their_bundle_dict["spk"]),
            signed_prekey_signature=bytes.fromhex(their_bundle_dict["spk_sig"]),
            signed_prekey_id=their_bundle_dict["spk_id"],
            one_time_prekey=bytes.fromhex(their_bundle_dict["opk"])
            if their_bundle_dict.get("opk")
            else None,
            one_time_prekey_id=their_bundle_dict.get("opk_id"),
        )

        # Create session via existing SecureSessionManager
        session, init_bytes = self._session_manager.create_session(peer_id, bundle)

        # Create trust record
        their_fp = compute_fingerprint(bundle.identity_key)
        our_identity = self._session_manager.local_keys.identity_key.public
        code = compute_verification_code(our_identity, bundle.identity_key)
        self._create_trust_record(peer_id, peer_name, their_fp, code)

        return init_bytes.hex()

    def accept_handshake(
        self,
        peer_id: str,
        peer_name: str,
        init_message_hex: str,
        their_identity_key_hex: str,
    ) -> Optional[bytes]:
        """
        Accept a secure session initiated by a peer.

        Args:
            peer_id: Remote node's ID
            peer_name: Remote node's display name
            init_message_hex: X3DH init message as hex string
            their_identity_key_hex: Their Ed25519 identity key as hex

        Returns:
            Decrypted initial message (if any), or None
        """
        if not self._session_manager:
            their_fp = hashlib.sha256(peer_id.encode()).hexdigest()
            self._create_trust_record(peer_id, peer_name, their_fp, "")
            return None

        init_bytes = bytes.fromhex(init_message_hex)
        session, initial_plaintext = self._session_manager.accept_session(peer_id, init_bytes)

        their_identity = bytes.fromhex(their_identity_key_hex)
        their_fp = compute_fingerprint(their_identity)
        our_identity = self._session_manager.local_keys.identity_key.public
        code = compute_verification_code(our_identity, their_identity)
        self._create_trust_record(peer_id, peer_name, their_fp, code)

        return initial_plaintext

    def _create_trust_record(self, peer_id: str, peer_name: str, fingerprint: str, code: str):
        """Create or update a trust record for a peer."""
        existing = self._trust_store.get(peer_id)

        if existing and existing.identity_fingerprint != fingerprint:
            # Key changed! Revoke trust.
            logger.warning(
                f"Identity key changed for peer {peer_name} ({peer_id[:8]}...). "
                f"Trust revoked — re-verification required."
            )
            existing.trust_level = "unknown"
            existing.permissions = PermissionMatrix()
            existing.identity_fingerprint = fingerprint
            existing.verification_code = code
            existing.trusted_at = None
            self._trust_store.upsert(existing)
            return

        if existing:
            # Update name and code but preserve trust
            existing.node_name = peer_name
            existing.verification_code = code
            self._trust_store.upsert(existing)
            return

        # New peer
        record = TrustRecord(
            node_id=peer_id,
            node_name=peer_name,
            identity_fingerprint=fingerprint,
            trust_level="pending",
            verification_code=code,
        )
        self._trust_store.upsert(record)
        logger.info(
            f"New peer pending trust: {peer_name} ({peer_id[:8]}...) "
            f"verification code: {format_verification_code(code)}"
        )

    # ── Trust lifecycle ──

    def approve(self, node_id: str) -> bool:
        """
        Approve a peer after verification code confirmation.
        Transitions: pending/unknown → trusted
        """
        record = self._trust_store.get(node_id)
        if not record:
            logger.warning(f"No trust record for {node_id}")
            return False

        if record.trust_level == "blocked":
            logger.warning(f"Cannot approve blocked peer {record.node_name}. Unblock first.")
            return False

        record.trust_level = "trusted"
        record.trusted_at = _utcnow().isoformat()
        record.last_verified = _utcnow().isoformat()
        record.permissions.send_messages = True
        self._trust_store.upsert(record)
        logger.info(f"Peer trusted: {record.node_name} ({node_id[:8]}...)")

        # Also update PeerStore trust level if available
        try:
            from .discovery import PeerStore

            peer_store = PeerStore(self._mesh_dir)
            peer_store.set_trust(node_id, "trusted")
        except Exception as e:
            logger.debug("Error suppressed: %s", e)
        return True

    def block(self, node_id: str) -> bool:
        """Block a peer. Revokes all permissions and trust."""
        record = self._trust_store.get(node_id)
        if not record:
            logger.warning(f"No trust record for {node_id}")
            return False

        record.trust_level = "blocked"
        record.blocked_at = _utcnow().isoformat()
        record.permissions = PermissionMatrix.all_deny()
        self._trust_store.upsert(record)

        # Remove encrypted session
        if self._session_manager and node_id in self._session_manager.sessions:
            del self._session_manager.sessions[node_id]
            self._session_manager._save_sessions()

        logger.info(f"Peer blocked: {record.node_name} ({node_id[:8]}...)")

        try:
            from .discovery import PeerStore

            PeerStore(self._mesh_dir).set_trust(node_id, "blocked")
        except Exception as e:
            logger.debug("Error suppressed: %s", e)
        return True

    def unblock(self, node_id: str) -> bool:
        """Unblock a peer. Sets to unknown (requires re-verification)."""
        record = self._trust_store.get(node_id)
        if not record:
            return False

        record.trust_level = "unknown"
        record.blocked_at = None
        self._trust_store.upsert(record)
        logger.info(f"Peer unblocked: {record.node_name} ({node_id[:8]}...)")
        return True

    # ── Permissions ──

    def grant_permission(self, node_id: str, permission: str) -> bool:
        """Grant a permission to a trusted peer."""
        record = self._trust_store.get(node_id)
        if not record:
            return False
        if record.trust_level != "trusted":
            logger.warning(f"Cannot grant permissions to non-trusted peer {record.node_name}")
            return False
        if record.permissions.grant(permission):
            self._trust_store.upsert(record)
            return True
        logger.warning(f"Unknown permission: {permission}")
        return False

    def revoke_permission(self, node_id: str, permission: str) -> bool:
        """Revoke a permission from a peer."""
        record = self._trust_store.get(node_id)
        if not record:
            return False
        if record.permissions.revoke(permission):
            self._trust_store.upsert(record)
            return True
        return False

    def check_permission(self, node_id: str, permission: str) -> bool:
        """Check if a peer has a specific permission."""
        record = self._trust_store.get(node_id)
        if not record or record.trust_level != "trusted":
            return False
        return record.permissions.check(permission)

    # ── Encryption ──

    def encrypt_for_peer(self, peer_id: str, plaintext: bytes) -> Optional[bytes]:
        """Encrypt a message for a peer using the Double Ratchet session."""
        if not self._session_manager:
            return None
        try:
            return self._session_manager.encrypt(peer_id, plaintext)
        except (ValueError, RuntimeError) as e:
            logger.warning(f"Encryption failed for {peer_id[:8]}: {e}")
            return None

    def decrypt_from_peer(self, peer_id: str, ciphertext: bytes) -> Optional[bytes]:
        """Decrypt a message from a peer."""
        if not self._session_manager:
            return None
        try:
            return self._session_manager.decrypt(peer_id, ciphertext)
        except (ValueError, RuntimeError) as e:
            logger.warning(f"Decryption failed for {peer_id[:8]}: {e}")
            return None

    def has_session(self, peer_id: str) -> bool:
        """Check if we have an encrypted session with a peer."""
        if not self._session_manager:
            return False
        session = self._session_manager.get_session(peer_id)
        return session is not None and session.established

    # ── Query ──

    def get_trust_record(self, node_id: str) -> Optional[TrustRecord]:
        return self._trust_store.get(node_id)

    def get_verification_code(self, node_id: str) -> Optional[str]:
        """Get the formatted verification code for a peer."""
        record = self._trust_store.get(node_id)
        if record and record.verification_code:
            return format_verification_code(record.verification_code)
        return None

    def get_all_records(self) -> List[TrustRecord]:
        return self._trust_store.get_all()

    def get_trusted_peers(self) -> List[TrustRecord]:
        return self._trust_store.get_trusted()

    def is_trusted(self, node_id: str) -> bool:
        record = self._trust_store.get(node_id)
        return record is not None and record.trust_level == "trusted"

    def is_blocked(self, node_id: str) -> bool:
        record = self._trust_store.get(node_id)
        return record is not None and record.trust_level == "blocked"

    def get_status(self) -> dict:
        """Get trust status for CLI/dashboard."""
        records = self._trust_store.get_all()
        return {
            "has_crypto": self.has_crypto,
            "our_fingerprint": self.get_our_fingerprint(),
            "total_peers": len(records),
            "trusted": len([r for r in records if r.trust_level == "trusted"]),
            "pending": len([r for r in records if r.trust_level == "pending"]),
            "blocked": len([r for r in records if r.trust_level == "blocked"]),
            "peers": [
                {
                    "node_id": r.node_id[:8] + "...",
                    "name": r.node_name,
                    "trust": r.trust_level,
                    "permissions": r.permissions.granted_list,
                    "fingerprint": r.identity_fingerprint[:16] + "...",
                    "has_session": self.has_session(r.node_id),
                }
                for r in records
            ],
        }
