#!/usr/bin/env python3
"""
LockStock Guard Daemon - Secure Enclave in Software

This daemon holds decrypted secrets in memory and serves them via Unix socket.
Agent applications NEVER have direct access to the vault file.

Security Architecture:
    liberty-user: Owns the vault, runs the daemon
    agent-user:   Runs application code, connects via socket

Filesystem Layout:
    /var/lib/liberty/secrets.enc   (liberty-user:liberty-user, mode 600)
    /var/run/liberty/liberty.sock  (liberty-user:agents, mode 660)

Usage:
    # Start daemon (as liberty-user)
    lockstock-guard start

    # Agent code (as agent-user)
    from lockstock_guard import client
    secret = client.get("DATABASE_URL")
"""

import os
import sys
import json
import socket
import signal
import struct
import threading
import logging
import hashlib
import hmac
import grp
import pwd
import time
from pathlib import Path
from datetime import datetime, timezone
from typing import Dict, Optional, Any, List
from dataclasses import dataclass, field

# Optional httpx for server sync (non-blocking import)
try:
    import httpx
    HTTPX_AVAILABLE = True
except ImportError:
    HTTPX_AVAILABLE = False
    httpx = None

# Import vault from liberty v1.0.0 (stable, personal use)
# Import enterprise fingerprinting from our module (cloud-native)
try:
    from liberty import SecretVault
    from lockstock_guard.enterprise_fingerprint import EnterpriseFingerprint
except ImportError:
    # Fallback for standalone testing
    print("Warning: liberty module not found, using mock vault")
    SecretVault = None
    EnterpriseFingerprint = None

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("lockstock-guard")


# Default paths
DEFAULT_SOCKET_PATH = "/var/run/liberty/liberty.sock"
DEFAULT_VAULT_PATH = "/var/lib/liberty"
DEFAULT_PID_FILE = "/var/run/liberty/liberty.pid"

# Protocol constants
PROTOCOL_VERSION = 1
MAX_MESSAGE_SIZE = 64 * 1024  # 64KB max message


class DaemonProtocol:
    """Wire protocol for daemon communication."""

    # Message types
    MSG_GET_SECRET = 0x01
    MSG_SECRET_RESPONSE = 0x02
    MSG_LIST_KEYS = 0x03
    MSG_KEYS_RESPONSE = 0x04
    MSG_PING = 0x05
    MSG_PONG = 0x06
    MSG_SET_SECRET = 0x07      # Store a secret
    MSG_SET_RESPONSE = 0x08    # Set operation response
    MSG_DELETE_SECRET = 0x09   # Delete a secret
    MSG_DELETE_RESPONSE = 0x0A # Delete operation response
    # Chain Authority Mode (v1.1)
    MSG_SIGN_AND_ADVANCE = 0x0B   # Sign action AND advance chain state
    MSG_SIGN_RESPONSE = 0x0C      # Signature + new state
    MSG_SYNC_CHAIN = 0x0D         # Request chain sync from server
    MSG_SYNC_RESPONSE = 0x0E      # Chain sync result
    MSG_GET_CHAIN_STATE = 0x0F    # Get current chain state for agent
    MSG_CHAIN_STATE_RESPONSE = 0x10
    MSG_ERROR = 0xFF


# =============================================================================
# CRYPTOGRAPHIC CONSTANTS AND MATRIX MATH
# =============================================================================

FIELD_MODULUS = 65537  # Fermat prime F4

# Task generators - MUST MATCH server's src/generator.rs EXACTLY (v1.0.8+)
# These are Burau representation generators for the braid group B3
# Matrix format: {"a11": row1_col1, "a12": row1_col2, "a21": row2_col1, "a22": row2_col2}
GENERATORS = {
    # === Operations (original) ===
    "deploy": {"a11": 1, "a12": 1, "a21": 0, "a22": 1},
    "restart": {"a11": 1, "a12": 0, "a21": 1, "a22": 1},
    "backup": {"a11": 2, "a12": 1, "a21": 1, "a22": 1},
    "rollback": {"a11": 1, "a12": 1, "a21": 1, "a22": 2},

    # === Agent Lifecycle ===
    "heartbeat": {"a11": 0, "a12": 65536, "a21": 1, "a22": 0},  # Rotation matrix (order 4)
    "checkpoint": {"a11": 2, "a12": 1, "a21": 1, "a22": 2},
    "teleport": {"a11": 3, "a12": 1, "a21": 1, "a22": 1},

    # === MCP Tool Categories ===
    "fileread": {"a11": 1, "a12": 2, "a21": 0, "a22": 1},    # Upper triangular
    "filewrite": {"a11": 1, "a12": 3, "a21": 0, "a22": 1},   # Different upper triangular
    "network": {"a11": 1, "a12": 0, "a21": 2, "a22": 1},     # Lower triangular
    "shell": {"a11": 1, "a12": 4, "a21": 0, "a22": 1},       # Distinct from Deploy (v1.0.8+)
    "database": {"a11": 1, "a12": 0, "a21": 3, "a22": 1},    # Different lower triangular
    "chat_completion": {"a11": 1, "a12": 5, "a21": 0, "a22": 1},  # LLM inference (request)
    "chat_completion_response": {"a11": 1, "a12": 6, "a21": 0, "a22": 1},  # LLM inference (response)

    # === A2A Protocol ===
    "handoff": {"a11": 3, "a12": 2, "a21": 1, "a22": 1},
    "delegate": {"a11": 1, "a12": 1, "a21": 2, "a22": 3},

    # === Wildcard ===
    # Identity matrix - admin capability doesn't mutate braid state
    "all": {"a11": 1, "a12": 0, "a21": 0, "a22": 1},
}

# Task name to PascalCase mapping - MUST MATCH Rust's {:?} Debug format for Task enum
# Rust's Debug trait prints enum variants exactly as written in code (PascalCase)
# This mapping ensures HMAC signatures match between Python daemon and Rust server
TASK_PASCALCASE = {
    # Operations
    "deploy": "Deploy",
    "restart": "Restart",
    "backup": "Backup",
    "rollback": "Rollback",
    # Agent Lifecycle
    "heartbeat": "Heartbeat",
    "checkpoint": "Checkpoint",
    "teleport": "Teleport",
    # MCP Tool Categories
    "fileread": "FileRead",
    "filewrite": "FileWrite",
    "network": "Network",
    "shell": "Shell",
    "database": "Database",
    "chat_completion": "ChatCompletion",
    "chat_completion_response": "ChatCompletionResponse",
    # A2A Protocol
    "handoff": "Handoff",
    "delegate": "Delegate",
    # Wildcard
    "all": "All",
}

# Task name aliases - maps alternative names to canonical form (keys in GENERATORS)
# This ensures consistency regardless of how task name is formatted in input
TASK_ALIASES = {
    # No-underscore variants → canonical underscore form
    "chatcompletion": "chat_completion",
    "chatcompletionresponse": "chat_completion_response",
    # Underscore variants (already canonical, but explicit for clarity)
    "chat_completion": "chat_completion",
    "chat_completion_response": "chat_completion_response",
    # SDK compatibility aliases (match server's types.rs)
    "file_read": "fileread",
    "file_write": "filewrite",
    "read": "fileread",
    "query": "fileread",
    "write": "filewrite",
    "edit": "filewrite",
    "http": "network",
    "fetch": "network",
    "bash": "shell",
    "exec": "shell",
    "execute": "shell",
    "db": "database",
    "sql": "database",
    "hand_off": "handoff",
    "*": "all",
    "admin": "all",
}


def normalize_task(task: str) -> str:
    """
    Normalize task name to canonical form (lowercase, as used in GENERATORS).

    Handles:
    - Case insensitivity (HEARTBEAT → heartbeat)
    - Alias resolution (chatcompletion → chat_completion)
    - SDK compatibility (query → fileread)
    """
    task_lower = task.lower()
    # Check if it's an alias
    if task_lower in TASK_ALIASES:
        return TASK_ALIASES[task_lower]
    # Otherwise return as-is (will be validated against GENERATORS later)
    return task_lower


def matrix_multiply(m1: Dict[str, int], m2: Dict[str, int]) -> Dict[str, int]:
    """Multiply two 2x2 matrices over F_65537."""
    return {
        "a11": (m1["a11"] * m2["a11"] + m1["a12"] * m2["a21"]) % FIELD_MODULUS,
        "a12": (m1["a11"] * m2["a12"] + m1["a12"] * m2["a22"]) % FIELD_MODULUS,
        "a21": (m1["a21"] * m2["a11"] + m1["a22"] * m2["a21"]) % FIELD_MODULUS,
        "a22": (m1["a21"] * m2["a12"] + m1["a22"] * m2["a22"]) % FIELD_MODULUS,
    }


def compute_state_hash(matrix: Dict[str, int]) -> str:
    """Compute SHA-256 hash of matrix state (matches Rust server)."""
    # Format must match src/matrix.rs hash() method
    data = f"{matrix['a11']},{matrix['a12']},{matrix['a21']},{matrix['a22']}"
    return hashlib.sha256(data.encode()).hexdigest()


def generate_signature(
    secret: str,
    client_id: str,
    task: str,
    parent_hash: str,
    matrix: Dict[str, int],
    timestamp: int,
) -> str:
    """
    Generate HMAC-SHA256 signature for a verification request.

    MUST match server's src/crypto.rs generate_signature() format:
    format!("{}|{:?}|{}|{},{},{},{}|{}", client_id, task, parent_hash, a11, a12, a21, a22, timestamp)
    """
    # Normalize task to canonical form, then convert to PascalCase
    # This handles aliases (chatcompletion → chat_completion) and case insensitivity
    task_canonical = normalize_task(task)
    task_pascalcase = TASK_PASCALCASE.get(task_canonical)
    if not task_pascalcase:
        raise ValueError(
            f"Unknown task '{task}' (normalized: '{task_canonical}') for signature. "
            f"Valid tasks: {', '.join(sorted(TASK_PASCALCASE.keys()))}. "
            f"Add task to TASK_PASCALCASE dict in daemon.py."
        )
    message = f"{client_id}|{task_pascalcase}|{parent_hash}|{matrix['a11']},{matrix['a12']},{matrix['a21']},{matrix['a22']}|{timestamp}"
    return hmac.new(secret.encode(), message.encode(), hashlib.sha256).hexdigest()


@dataclass
class AgentChainState:
    """
    Authoritative chain state for an agent.

    Guard tracks this to validate sign requests - preventing compromised
    MCP instances from forging chain history with fake parent_hash values.
    """
    agent_id: str
    current_hash: str
    sequence: int
    state_matrix: Dict[str, int]
    last_server_timestamp: int = 0
    last_synced: float = field(default_factory=time.time)
    sync_failures: int = 0


def _encode_message(msg_type: int, payload: bytes) -> bytes:
    """Encode message with length prefix and type."""
    header = struct.pack("!BBL", PROTOCOL_VERSION, msg_type, len(payload))
    return header + payload


def _decode_header(data: bytes) -> tuple:
    """Decode message header. Returns (version, msg_type, length)."""
    if len(data) < 6:
        raise ValueError("Invalid header: too short")
    version, msg_type, length = struct.unpack("!BBL", data[:6])
    if version != PROTOCOL_VERSION:
        raise ValueError(f"Unsupported protocol version: {version}")
    if length > MAX_MESSAGE_SIZE:
        raise ValueError(f"Message too large: {length}")
    return version, msg_type, length


# Attach to DaemonProtocol for backwards compatibility
DaemonProtocol.encode_message = staticmethod(_encode_message)
DaemonProtocol.decode_header = staticmethod(_decode_header)


class SecretCache:
    """
    In-memory secret cache with access control, fingerprint binding, and chain authority.

    Secrets are decrypted once at daemon startup and held in memory.
    Never written to disk in plaintext.

    Security:
    - Each secret is bound to a hardware/cloud fingerprint
    - Chain state is tracked to prevent fork attacks from compromised MCP
    - Access is denied if current fingerprint doesn't match bound fingerprint
    """

    def __init__(self, server_url: str = None):
        self._secrets: Dict[str, str] = {}
        self._fingerprints: Dict[str, str] = {}  # key -> bound fingerprint
        self._chain_state: Dict[str, AgentChainState] = {}  # agent_id -> chain state
        self._access_log: list = []
        self._lock = threading.RLock()
        self._server_url = server_url or os.getenv(
            "LOCKSTOCK_URL", "https://lockstock-api-i9kp.onrender.com"
        )
        self._sync_interval = 300  # Sync every 5 minutes
        self._chain_state_file = Path("/var/lib/liberty/chain_state.json")

    def load_from_vault(self, vault: "SecretVault") -> int:
        """Load all secrets from vault into memory."""
        with self._lock:
            try:
                raw_secrets = vault._load_secrets()
                # Extract values and fingerprints (vault stores metadata too)
                self._secrets = {}
                self._fingerprints = {}
                for key, data in raw_secrets.items():
                    # Handle JSON-encoded secret data (new format with fingerprint)
                    if isinstance(data, str):
                        try:
                            parsed = json.loads(data)
                            if isinstance(parsed, dict) and 'value' in parsed:
                                self._secrets[key] = parsed['value']
                                if 'fingerprint' in parsed:
                                    self._fingerprints[key] = parsed['fingerprint']
                                continue
                        except (json.JSONDecodeError, TypeError):
                            pass
                        # Plain string secret (legacy format)
                        self._secrets[key] = data
                    elif isinstance(data, dict):
                        self._secrets[key] = data.get('value', str(data))
                        if 'fingerprint' in data:
                            self._fingerprints[key] = data['fingerprint']
                    else:
                        self._secrets[key] = str(data)

                logger.info(f"Loaded {len(self._secrets)} secrets into cache")
                if self._fingerprints:
                    logger.info(f"Loaded {len(self._fingerprints)} fingerprint bindings")
                return len(self._secrets)
            except Exception as e:
                logger.error(f"Failed to load vault: {e}")
                return 0

    def get(self, key: str, client_info: dict = None, verify_fingerprint: bool = True) -> Optional[str]:
        """Get a secret, verifying fingerprint binding and logging access."""
        with self._lock:
            value = self._secrets.get(key)

            if value is None:
                self._log_access(key, client_info, found=False, reason="not_found")
                return None

            # Verify fingerprint binding if enabled and fingerprint exists
            if verify_fingerprint and key in self._fingerprints:
                if not self._verify_binding(key, client_info):
                    self._log_access(key, client_info, found=True, reason="fingerprint_mismatch")
                    return None

            self._log_access(key, client_info, found=True, reason="success")
            return value

    def _verify_binding(self, key: str, client_info: dict) -> bool:
        """
        Verify current environment matches bound fingerprint.

        Uses EnterpriseFingerprint to detect cloud/container/bare-metal
        and compares against the fingerprint stored at bind time.
        """
        bound_fingerprint = self._fingerprints.get(key)
        if not bound_fingerprint:
            # No binding - allow access (legacy secrets)
            logger.warning(f"No fingerprint binding for key '{key}' - allowing access")
            return True

        if EnterpriseFingerprint is None:
            logger.error("EnterpriseFingerprint not available - cannot verify binding")
            return False

        try:
            current_fingerprint = EnterpriseFingerprint.generate()

            if current_fingerprint == bound_fingerprint:
                logger.debug(f"Fingerprint verified for key '{key}'")
                return True
            else:
                logger.warning(
                    f"Fingerprint mismatch for key '{key}': "
                    f"bound={bound_fingerprint[:16]}... current={current_fingerprint[:16]}..."
                )
                return False
        except Exception as e:
            logger.error(f"Failed to generate fingerprint: {e}")
            return False

    def get_fingerprint(self, key: str) -> Optional[str]:
        """Get the bound fingerprint for a key."""
        with self._lock:
            return self._fingerprints.get(key)

    def list_keys(self) -> list:
        """List all secret keys (not values)."""
        with self._lock:
            return list(self._secrets.keys())

    def set(self, key: str, value: str, fingerprint: Optional[str] = None,
            client_info: dict = None, vault: "SecretVault" = None) -> bool:
        """
        Store a secret in cache and optionally persist to vault.

        Args:
            key: Secret key name
            value: Secret value
            fingerprint: Optional hardware fingerprint to bind the secret to
            client_info: Client information for audit logging
            vault: Optional vault instance for persistence

        Returns:
            True if successful
        """
        with self._lock:
            # Store in memory cache
            self._secrets[key] = value
            if fingerprint:
                self._fingerprints[key] = fingerprint

            # Persist to vault if provided
            if vault is not None:
                try:
                    # Store with fingerprint metadata
                    vault_data = json.dumps({
                        'value': value,
                        'fingerprint': fingerprint,
                        'created_at': datetime.now(timezone.utc).isoformat(),
                    })
                    vault.set(key, vault_data)
                    logger.info(f"Secret '{key}' stored and persisted to vault")
                except Exception as e:
                    logger.error(f"Failed to persist secret '{key}' to vault: {e}")
                    # Still return True since it's in memory cache
            else:
                logger.info(f"Secret '{key}' stored in memory (not persisted)")

            self._log_access(key, client_info, found=True, reason="set")
            return True

    def delete(self, key: str, client_info: dict = None, vault: "SecretVault" = None) -> bool:
        """
        Delete a secret from cache and optionally from vault.

        Args:
            key: Secret key name
            client_info: Client information for audit logging
            vault: Optional vault instance for persistence

        Returns:
            True if secret existed and was deleted
        """
        with self._lock:
            if key not in self._secrets:
                self._log_access(key, client_info, found=False, reason="delete_not_found")
                return False

            # Remove from memory cache
            del self._secrets[key]
            if key in self._fingerprints:
                del self._fingerprints[key]

            # Remove from vault if provided
            if vault is not None:
                try:
                    vault.delete(key)
                    logger.info(f"Secret '{key}' deleted from cache and vault")
                except Exception as e:
                    logger.warning(f"Failed to delete secret '{key}' from vault: {e}")
            else:
                logger.info(f"Secret '{key}' deleted from memory")

            self._log_access(key, client_info, found=True, reason="deleted")
            return True

    def _log_access(self, key: str, client_info: dict, found: bool, reason: str = ""):
        """Log secret access for audit."""
        entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "key": key,
            "found": found,
            "reason": reason,
            "client": client_info or {},
        }
        self._access_log.append(entry)
        # Keep last 1000 entries
        if len(self._access_log) > 1000:
            self._access_log = self._access_log[-1000:]

        # Log security-relevant events
        if reason == "fingerprint_mismatch":
            logger.warning(f"ACCESS DENIED: Fingerprint mismatch for '{key}' from {client_info}")

    # =========================================================================
    # CHAIN AUTHORITY MODE - Guard tracks and validates chain state
    # =========================================================================

    def get_chain_state(self, agent_id: str) -> Optional[AgentChainState]:
        """Get current chain state for an agent."""
        with self._lock:
            return self._chain_state.get(agent_id)

    def set_chain_state(self, agent_id: str, state: AgentChainState) -> None:
        """Set chain state for an agent."""
        with self._lock:
            self._chain_state[agent_id] = state
            self._persist_chain_state()

    def sign_and_advance(
        self,
        agent_id: str,
        task: str,
        parent_hash: str,
        client_info: dict = None,
        payload: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """
        Sign an action AND advance chain state atomically.

        This is the core of Chain Authority Mode:
        1. Validate parent_hash matches Guard's known chain head
        2. Compute new state (matrix multiplication)
        3. Generate HMAC signature
        4. Advance local chain state
        5. Return signature + new state

        Returns dict with signature, new_hash, new_sequence, new_matrix, timestamp
        Raises ValueError if parent_hash doesn't match or agent not initialized
        """
        with self._lock:
            # Get current chain state
            state = self._chain_state.get(agent_id)
            if state is None:
                # Try to sync from server first
                if self._sync_agent_from_server(agent_id):
                    state = self._chain_state.get(agent_id)
                if state is None:
                    raise ValueError(f"Agent {agent_id} not initialized. Call sync_chain first.")

            # CRITICAL: Validate parent_hash matches Guard's known head
            if parent_hash != state.current_hash:
                logger.warning(
                    f"CHAIN FORK ATTEMPT BLOCKED: {agent_id} "
                    f"sent parent={parent_hash[:16]}... "
                    f"but Guard knows head={state.current_hash[:16]}..."
                )
                raise ValueError(
                    f"parent_hash mismatch: expected {state.current_hash[:16]}..., "
                    f"got {parent_hash[:16]}..."
                )

            # Get secret for this agent
            secret = self._secrets.get(agent_id)
            if not secret:
                # Try alternate key formats
                for key_format in [f"{agent_id}_SECRET", f"{agent_id.upper()}_SECRET"]:
                    secret = self._secrets.get(key_format)
                    if secret:
                        break
            if not secret:
                raise ValueError(f"Secret not found for agent {agent_id}")

            # Get task generator - NO FALLBACK!
            # Unknown tasks must fail loudly to prevent HMAC mismatch
            task_normalized = normalize_task(task)
            if task_normalized not in GENERATORS:
                raise ValueError(
                    f"Unknown task '{task}'. Valid tasks: {', '.join(sorted(GENERATORS.keys()))}. "
                    f"This error prevents silent HMAC mismatch. If you need this task, "
                    f"add it to src/generator.rs AND update Guard daemon GENERATORS dict."
                )
            generator = GENERATORS[task_normalized]

            # Compute new state
            new_matrix = matrix_multiply(state.state_matrix, generator)
            new_sequence = state.sequence + 1
            timestamp = int(time.time())

            # Compute chain hash using v2 formula (matches Rust server)
            # Hash preimage: parent_hash | agent_id | task | matrix | payload
            # Must use canonical JSON (sorted keys, compact) for cross-language determinism
            def canonical_json(obj):
                """Canonical JSON: sorted keys, compact format, no whitespace."""
                return json.dumps(obj, sort_keys=True, separators=(',', ':'))

            matrix_str = canonical_json(new_matrix)

            if payload:
                # v2 format: include payload in hash preimage
                payload_str = canonical_json(payload)
                chain_message = f"{parent_hash}|{agent_id}|{task.upper()}|{matrix_str}|{payload_str}"
            else:
                # v1 format: backward compatible (no payload)
                chain_message = f"{parent_hash}|{agent_id}|{task.upper()}|{matrix_str}"

            new_hash = hashlib.sha256(chain_message.encode()).hexdigest()

            # Generate signature using current_hash as HMAC key (NO SECRETS architecture)
            # The chain hash IS the authentication credential - it evolves with each action
            # This must match server's verify_signature_with_hash() which uses current_hash
            signature = generate_signature(
                parent_hash,  # Use current chain hash as HMAC key, NOT vault secret
                agent_id,
                task_normalized,
                parent_hash,
                new_matrix,
                timestamp,
            )

            # ADVANCE STATE (Guard is now authoritative)
            state.current_hash = new_hash
            state.sequence = new_sequence
            state.state_matrix = new_matrix
            state.last_synced = time.time()

            # Persist to disk
            self._persist_chain_state()

            logger.info(
                f"Chain advanced: {agent_id} seq={new_sequence} hash={new_hash[:16]}..."
            )

            # Best-effort sync to server (background, non-blocking)
            matrix_fingerprint = compute_state_hash(new_matrix)
            threading.Thread(
                target=self._sync_chain_to_server,
                args=(agent_id, new_sequence, new_hash, matrix_fingerprint, parent_hash, task, new_matrix, payload),
                daemon=True,
                name=f"chain-sync-{agent_id}-{new_sequence}"
            ).start()

            return {
                "signature": signature,
                "new_hash": new_hash,
                "new_sequence": new_sequence,
                "new_matrix": new_matrix,
                "timestamp": timestamp,
                "parent_hash": parent_hash,
            }

    def sync_chain_from_server(self, agent_id: str) -> bool:
        """
        Sync chain state from LockStock server.

        Called on daemon startup and periodically to ensure Guard
        stays in sync with server's authoritative state.
        """
        return self._sync_agent_from_server(agent_id)

    def _sync_chain_to_server(
        self,
        agent_id: str,
        sequence: int,
        chain_hash: str,
        state_hash: str,
        parent_hash: str,
        task: str,
        matrix: Dict[str, int],
        payload: Optional[Dict[str, Any]] = None
    ) -> bool:
        """
        Best-effort sync: Push chain entry to server.

        Called after local chain advances. If sync fails, it's logged
        but doesn't affect the local chain (already advanced).

        Args:
            chain_hash: Daemon's computed chain hash (for cross-verification)
            state_hash: Matrix fingerprint (for audit_log.matrix_state_hash)
        """
        if not HTTPX_AVAILABLE:
            logger.warning("httpx not available - cannot sync to server")
            return False

        try:
            url = f"{self._server_url}/api/guard/chain/sync"

            sync_request = {
                "agent_id": agent_id,
                "sequence": sequence,
                "chain_hash": chain_hash,  # Daemon's computation (for verification)
                "state_hash": state_hash,  # Matrix fingerprint
                "parent_hash": parent_hash,
                "task": task,
                "matrix": matrix,
            }

            if payload:
                sync_request["payload"] = payload

            with httpx.Client(timeout=5) as client:
                response = client.post(url, json=sync_request)

            if response.status_code == 200:
                logger.info(f"✅ Chain synced: {agent_id} seq={sequence} hash={chain_hash[:16]}...")
                return True
            else:
                logger.warning(
                    f"⚠️  Chain sync failed: {response.status_code} - {response.text[:200]}"
                )
                return False

        except Exception as e:
            logger.error(f"❌ Chain sync error for {agent_id}: {e}")
            return False

    def _sync_agent_from_server(self, agent_id: str) -> bool:
        """Internal: Fetch chain state from server."""
        if not HTTPX_AVAILABLE:
            logger.warning("httpx not available - cannot sync from server")
            return False

        try:
            url = f"{self._server_url}/api/guard/agent/{agent_id}/state"
            logger.info(f"Syncing chain state from {url}")

            with httpx.Client(timeout=10) as client:
                response = client.get(url)

            if response.status_code == 404:
                logger.warning(f"Agent {agent_id} not found on server")
                return False

            if response.status_code != 200:
                logger.error(f"Server sync failed: {response.status_code}")
                return False

            data = response.json()

            # Update local state
            state = AgentChainState(
                agent_id=agent_id,
                current_hash=data["last_hash"],
                sequence=data["last_sequence"],
                state_matrix=data["state_matrix"],
                last_server_timestamp=data.get("last_server_timestamp", 0),
                last_synced=time.time(),
            )

            with self._lock:
                existing = self._chain_state.get(agent_id)
                if existing and existing.current_hash != state.current_hash:
                    logger.warning(
                        f"Chain divergence detected for {agent_id}: "
                        f"local={existing.current_hash[:16]}... "
                        f"server={state.current_hash[:16]}... "
                        f"(server wins)"
                    )
                self._chain_state[agent_id] = state
                self._persist_chain_state()

            logger.info(
                f"Chain synced: {agent_id} seq={state.sequence} hash={state.current_hash[:16]}..."
            )
            return True

        except Exception as e:
            logger.error(f"Failed to sync chain state for {agent_id}: {e}")
            return False

    def _persist_chain_state(self) -> None:
        """Persist chain state to disk for daemon restart recovery."""
        try:
            self._chain_state_file.parent.mkdir(parents=True, exist_ok=True)
            data = {
                agent_id: {
                    "agent_id": state.agent_id,
                    "current_hash": state.current_hash,
                    "sequence": state.sequence,
                    "state_matrix": state.state_matrix,
                    "last_server_timestamp": state.last_server_timestamp,
                    "last_synced": state.last_synced,
                }
                for agent_id, state in self._chain_state.items()
            }
            with open(self._chain_state_file, "w") as f:
                json.dump(data, f, indent=2)
        except Exception as e:
            logger.error(f"Failed to persist chain state: {e}")

    def load_chain_state(self) -> int:
        """Load persisted chain state from disk."""
        if not self._chain_state_file.exists():
            return 0

        try:
            with open(self._chain_state_file) as f:
                data = json.load(f)

            count = 0
            for agent_id, state_data in data.items():
                self._chain_state[agent_id] = AgentChainState(
                    agent_id=state_data["agent_id"],
                    current_hash=state_data["current_hash"],
                    sequence=state_data["sequence"],
                    state_matrix=state_data["state_matrix"],
                    last_server_timestamp=state_data.get("last_server_timestamp", 0),
                    last_synced=state_data.get("last_synced", 0),
                )
                count += 1

            logger.info(f"Loaded {count} chain states from disk")
            return count
        except Exception as e:
            logger.error(f"Failed to load chain state: {e}")
            return 0


class ClientHandler(threading.Thread):
    """Handle a single client connection."""

    def __init__(self, conn: socket.socket, addr, cache: SecretCache,
                 allowed_keys: set = None, vault: "SecretVault" = None,
                 allow_write: bool = False):
        super().__init__(daemon=True)
        self.conn = conn
        self.addr = addr
        self.cache = cache
        self.allowed_keys = allowed_keys  # If set, restrict to these keys
        self.vault = vault  # For persistence of new secrets
        self.allow_write = allow_write  # Whether this client can set/delete secrets
        self._running = True

    def run(self):
        """Handle client requests."""
        try:
            # Get peer credentials (Unix socket only)
            client_info = self._get_peer_credentials()
            logger.info(f"Client connected: {client_info}")

            while self._running:
                # Read header
                header = self._recv_exact(6)
                if not header:
                    break

                version, msg_type, length = DaemonProtocol.decode_header(header)

                # Read payload
                payload = self._recv_exact(length) if length > 0 else b""

                # Handle message
                response = self._handle_message(msg_type, payload, client_info)
                self.conn.sendall(response)

        except Exception as e:
            logger.error(f"Client handler error: {e}")
        finally:
            self.conn.close()
            logger.info(f"Client disconnected")

    def _recv_exact(self, n: int) -> bytes:
        """Receive exactly n bytes."""
        data = b""
        while len(data) < n:
            chunk = self.conn.recv(n - len(data))
            if not chunk:
                return None
            data += chunk
        return data

    def _get_peer_credentials(self) -> dict:
        """Get peer credentials from Unix socket."""
        try:
            creds = self.conn.getsockopt(
                socket.SOL_SOCKET,
                socket.SO_PEERCRED,
                struct.calcsize("3i")
            )
            pid, uid, gid = struct.unpack("3i", creds)
            return {
                "pid": pid,
                "uid": uid,
                "gid": gid,
                "user": pwd.getpwuid(uid).pw_name if uid >= 0 else "unknown",
            }
        except Exception:
            return {"pid": 0, "uid": -1, "gid": -1, "user": "unknown"}

    def _handle_message(self, msg_type: int, payload: bytes, client_info: dict) -> bytes:
        """Handle a single message."""
        try:
            if msg_type == DaemonProtocol.MSG_PING:
                return DaemonProtocol.encode_message(DaemonProtocol.MSG_PONG, b"pong")

            elif msg_type == DaemonProtocol.MSG_GET_SECRET:
                key = payload.decode("utf-8")

                # Check access control
                if self.allowed_keys and key not in self.allowed_keys:
                    error = json.dumps({"error": "access_denied", "key": key})
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_ERROR,
                        error.encode("utf-8")
                    )

                # Get secret with fingerprint verification
                value = self.cache.get(key, client_info, verify_fingerprint=True)
                if value is None:
                    # Determine if it was not found or fingerprint mismatch
                    if key in self.cache._secrets:
                        # Secret exists but fingerprint failed
                        error = json.dumps({
                            "error": "fingerprint_mismatch",
                            "key": key,
                            "message": "Environment fingerprint does not match bound fingerprint"
                        })
                    else:
                        error = json.dumps({"error": "not_found", "key": key})
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_ERROR,
                        error.encode("utf-8")
                    )

                response = json.dumps({"key": key, "value": value})
                return DaemonProtocol.encode_message(
                    DaemonProtocol.MSG_SECRET_RESPONSE,
                    response.encode("utf-8")
                )

            elif msg_type == DaemonProtocol.MSG_LIST_KEYS:
                keys = self.cache.list_keys()
                # Filter by allowed keys if restricted
                if self.allowed_keys:
                    keys = [k for k in keys if k in self.allowed_keys]
                response = json.dumps({"keys": keys})
                return DaemonProtocol.encode_message(
                    DaemonProtocol.MSG_KEYS_RESPONSE,
                    response.encode("utf-8")
                )

            elif msg_type == DaemonProtocol.MSG_SET_SECRET:
                # Check if client has write permission
                if not self.allow_write:
                    error = json.dumps({
                        "error": "permission_denied",
                        "message": "Client does not have write permission"
                    })
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_ERROR,
                        error.encode("utf-8")
                    )

                try:
                    data = json.loads(payload.decode("utf-8"))
                    key = data.get("key")
                    value = data.get("value")
                    fingerprint = data.get("fingerprint")

                    if not key or value is None:
                        error = json.dumps({
                            "error": "invalid_request",
                            "message": "Missing 'key' or 'value' field"
                        })
                        return DaemonProtocol.encode_message(
                            DaemonProtocol.MSG_ERROR,
                            error.encode("utf-8")
                        )

                    # Auto-generate fingerprint if requested
                    if fingerprint == "auto" and EnterpriseFingerprint is not None:
                        fingerprint = EnterpriseFingerprint.generate()

                    success = self.cache.set(
                        key, value, fingerprint,
                        client_info=client_info,
                        vault=self.vault
                    )

                    response = json.dumps({
                        "success": success,
                        "key": key,
                        "fingerprint_bound": fingerprint is not None
                    })
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_SET_RESPONSE,
                        response.encode("utf-8")
                    )

                except json.JSONDecodeError as e:
                    error = json.dumps({
                        "error": "invalid_json",
                        "message": str(e)
                    })
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_ERROR,
                        error.encode("utf-8")
                    )

            elif msg_type == DaemonProtocol.MSG_DELETE_SECRET:
                # Check if client has write permission
                if not self.allow_write:
                    error = json.dumps({
                        "error": "permission_denied",
                        "message": "Client does not have write permission"
                    })
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_ERROR,
                        error.encode("utf-8")
                    )

                key = payload.decode("utf-8")
                success = self.cache.delete(key, client_info=client_info, vault=self.vault)

                if success:
                    response = json.dumps({"success": True, "key": key})
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_DELETE_RESPONSE,
                        response.encode("utf-8")
                    )
                else:
                    error = json.dumps({"error": "not_found", "key": key})
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_ERROR,
                        error.encode("utf-8")
                    )

            # =================================================================
            # CHAIN AUTHORITY MODE HANDLERS
            # =================================================================

            elif msg_type == DaemonProtocol.MSG_SIGN_AND_ADVANCE:
                # Sign action AND advance chain state (prevents fork attacks)
                try:
                    data = json.loads(payload.decode("utf-8"))
                    agent_id = data.get("agent_id")
                    task = data.get("task")
                    parent_hash = data.get("parent_hash")

                    if not all([agent_id, task, parent_hash]):
                        error = json.dumps({
                            "error": "invalid_request",
                            "message": "Missing required fields: agent_id, task, parent_hash"
                        })
                        return DaemonProtocol.encode_message(
                            DaemonProtocol.MSG_ERROR,
                            error.encode("utf-8")
                        )

                    result = self.cache.sign_and_advance(
                        agent_id=agent_id,
                        task=task,
                        parent_hash=parent_hash,
                        client_info=client_info,
                    )

                    response = json.dumps(result)
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_SIGN_RESPONSE,
                        response.encode("utf-8")
                    )

                except ValueError as e:
                    # Chain validation failed (parent_hash mismatch, etc.)
                    error = json.dumps({
                        "error": "chain_validation_failed",
                        "message": str(e)
                    })
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_ERROR,
                        error.encode("utf-8")
                    )
                except json.JSONDecodeError as e:
                    error = json.dumps({
                        "error": "invalid_json",
                        "message": str(e)
                    })
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_ERROR,
                        error.encode("utf-8")
                    )

            elif msg_type == DaemonProtocol.MSG_SYNC_CHAIN:
                # Sync chain state from server
                try:
                    data = json.loads(payload.decode("utf-8"))
                    agent_id = data.get("agent_id")

                    if not agent_id:
                        error = json.dumps({
                            "error": "invalid_request",
                            "message": "Missing agent_id"
                        })
                        return DaemonProtocol.encode_message(
                            DaemonProtocol.MSG_ERROR,
                            error.encode("utf-8")
                        )

                    success = self.cache.sync_chain_from_server(agent_id)
                    state = self.cache.get_chain_state(agent_id)

                    response = json.dumps({
                        "success": success,
                        "agent_id": agent_id,
                        "current_hash": state.current_hash if state else None,
                        "sequence": state.sequence if state else None,
                        "state_matrix": state.state_matrix if state else None,
                    })
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_SYNC_RESPONSE,
                        response.encode("utf-8")
                    )

                except json.JSONDecodeError as e:
                    error = json.dumps({
                        "error": "invalid_json",
                        "message": str(e)
                    })
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_ERROR,
                        error.encode("utf-8")
                    )

            elif msg_type == DaemonProtocol.MSG_GET_CHAIN_STATE:
                # Get current chain state for an agent
                try:
                    agent_id = payload.decode("utf-8")
                    state = self.cache.get_chain_state(agent_id)

                    if state:
                        response = json.dumps({
                            "agent_id": state.agent_id,
                            "current_hash": state.current_hash,
                            "sequence": state.sequence,
                            "state_matrix": state.state_matrix,
                            "last_server_timestamp": state.last_server_timestamp,
                            "last_synced": state.last_synced,
                        })
                        return DaemonProtocol.encode_message(
                            DaemonProtocol.MSG_CHAIN_STATE_RESPONSE,
                            response.encode("utf-8")
                        )
                    else:
                        error = json.dumps({
                            "error": "not_found",
                            "agent_id": agent_id,
                            "message": "Chain state not found. Call sync_chain first."
                        })
                        return DaemonProtocol.encode_message(
                            DaemonProtocol.MSG_ERROR,
                            error.encode("utf-8")
                        )

                except Exception as e:
                    error = json.dumps({
                        "error": "internal_error",
                        "message": str(e)
                    })
                    return DaemonProtocol.encode_message(
                        DaemonProtocol.MSG_ERROR,
                        error.encode("utf-8")
                    )

            else:
                error = json.dumps({"error": "unknown_message_type", "type": msg_type})
                return DaemonProtocol.encode_message(
                    DaemonProtocol.MSG_ERROR,
                    error.encode("utf-8")
                )

        except Exception as e:
            error = json.dumps({"error": "internal_error", "message": str(e)})
            return DaemonProtocol.encode_message(
                DaemonProtocol.MSG_ERROR,
                error.encode("utf-8")
            )


class LibertyDaemon:
    """
    Liberty Daemon - holds secrets in memory, serves via Unix socket.

    This implements the "Secure Enclave in Software" pattern:
    - Daemon runs as privileged user with vault access
    - Agent apps run as unprivileged user with socket access only
    - Secrets never touch agent process memory until explicitly requested
    """

    def __init__(
        self,
        socket_path: str = DEFAULT_SOCKET_PATH,
        vault_path: str = DEFAULT_VAULT_PATH,
        pid_file: str = DEFAULT_PID_FILE,
        socket_group: str = "agents",
        allow_write: bool = False,
    ):
        self.socket_path = Path(socket_path)
        self.vault_path = Path(vault_path)
        self.pid_file = Path(pid_file)
        self.socket_group = socket_group
        self.allow_write = allow_write

        self.cache = SecretCache()
        self.vault = None  # Set during _load_vault()
        self.server_socket = None
        self._running = False
        self._threads = []

    def start(self, foreground: bool = False):
        """Start the daemon."""
        # Check if already running
        if self._is_running():
            logger.error("Daemon already running")
            return False

        # Load secrets into cache
        if not self._load_vault():
            logger.error("Failed to load vault")
            return False

        # Create socket directory
        self.socket_path.parent.mkdir(parents=True, exist_ok=True)

        # Remove stale socket
        if self.socket_path.exists():
            self.socket_path.unlink()

        # Create Unix socket
        self.server_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self.server_socket.bind(str(self.socket_path))

        # Set socket permissions: owner + group can connect
        os.chmod(self.socket_path, 0o660)

        # Set group ownership if specified
        try:
            gid = grp.getgrnam(self.socket_group).gr_gid
            os.chown(self.socket_path, -1, gid)
        except KeyError:
            logger.warning(f"Group '{self.socket_group}' not found, using default")

        self.server_socket.listen(10)
        logger.info(f"Listening on {self.socket_path}")

        # Write PID file
        self.pid_file.parent.mkdir(parents=True, exist_ok=True)
        self.pid_file.write_text(str(os.getpid()))

        # Setup signal handlers (only works in main thread)
        try:
            signal.signal(signal.SIGTERM, self._handle_signal)
            signal.signal(signal.SIGINT, self._handle_signal)
            signal.signal(signal.SIGHUP, self._handle_signal)
        except ValueError:
            # Not in main thread, skip signal handling
            pass

        # Daemonize if not foreground
        if not foreground:
            self._daemonize()

        # Accept connections
        self._running = True
        try:
            while self._running:
                try:
                    self.server_socket.settimeout(1.0)
                    conn, addr = self.server_socket.accept()
                    handler = ClientHandler(
                        conn, addr, self.cache,
                        vault=self.vault if self.allow_write else None,
                        allow_write=self.allow_write
                    )
                    handler.start()
                    self._threads.append(handler)
                except socket.timeout:
                    continue
        finally:
            self._cleanup()

        return True

    def stop(self):
        """Stop the daemon."""
        if not self._is_running():
            logger.info("Daemon not running")
            return

        # Read PID and send SIGTERM
        try:
            pid = int(self.pid_file.read_text().strip())
            os.kill(pid, signal.SIGTERM)
            logger.info(f"Sent SIGTERM to PID {pid}")
        except (FileNotFoundError, ValueError, ProcessLookupError) as e:
            logger.error(f"Failed to stop daemon: {e}")

    def _load_vault(self) -> bool:
        """Load vault into memory cache."""
        if SecretVault is None:
            logger.error("SecretVault not available")
            return False

        self.vault = SecretVault(str(self.vault_path))

        # Initialize vault if it doesn't exist
        if not self.vault.secrets_file.exists():
            logger.info(f"Vault not found, creating empty vault at {self.vault_path}")
            result = self.vault.init()
            if not result:
                logger.error(f"Failed to initialize vault at {self.vault_path}")
                return False
            logger.info(f"Empty vault created at {self.vault.secrets_file}")

        count = self.cache.load_from_vault(self.vault)

        if count == 0:
            logger.info("Vault is empty - ready for agent binding")

        if self.allow_write:
            logger.info("Write operations ENABLED - clients can store/delete secrets")

        # Load persisted chain state (Chain Authority Mode)
        chain_count = self.cache.load_chain_state()
        if chain_count > 0:
            logger.info(f"Chain Authority Mode: {chain_count} agent chain states loaded")

        return True

    def _is_running(self) -> bool:
        """Check if daemon is already running."""
        if not self.pid_file.exists():
            return False

        try:
            pid = int(self.pid_file.read_text().strip())
            # Check if process exists
            os.kill(pid, 0)
            return True
        except (ValueError, ProcessLookupError):
            # Stale PID file
            self.pid_file.unlink(missing_ok=True)
            return False

    def _daemonize(self):
        """Fork into background daemon."""
        # First fork
        if os.fork() > 0:
            sys.exit(0)

        os.setsid()

        # Second fork
        if os.fork() > 0:
            sys.exit(0)

        # Redirect stdio
        sys.stdin = open(os.devnull, 'r')
        sys.stdout = open(os.devnull, 'w')
        sys.stderr = open(os.devnull, 'w')

    def _handle_signal(self, signum, frame):
        """Handle shutdown and reload signals."""
        if signum == signal.SIGHUP:
            logger.info("Received SIGHUP - reloading vault and chain state")
            self._reload_vault()
            self._reload_chain_state()
        else:
            logger.info(f"Received signal {signum}, shutting down")
            self._running = False

    def _reload_chain_state(self):
        """Reload chain state from disk, merging new agents without clobbering existing ones."""
        if not self.cache._chain_state_file.exists():
            return

        try:
            with open(self.cache._chain_state_file) as f:
                data = json.load(f)

            added = []
            with self.cache._lock:
                for agent_id, state_data in data.items():
                    if agent_id not in self.cache._chain_state:
                        self.cache._chain_state[agent_id] = AgentChainState(
                            agent_id=state_data["agent_id"],
                            current_hash=state_data["current_hash"],
                            sequence=state_data["sequence"],
                            state_matrix=state_data["state_matrix"],
                            last_server_timestamp=state_data.get("last_server_timestamp", 0),
                            last_synced=state_data.get("last_synced", 0),
                        )
                        added.append(agent_id)

            if added:
                logger.info(f"Merged {len(added)} new chain state(s): {', '.join(added)}")
        except Exception as e:
            logger.error(f"Failed to reload chain state: {e}")

    def _reload_vault(self):
        """Reload secrets from vault without restarting daemon."""
        logger.info("Reloading vault...")
        if self.vault is None:
            logger.error("Vault not initialized")
            return

        # Vault should exist from startup, but handle edge case
        if not self.vault.secrets_file.exists():
            logger.warning(f"Vault not found during reload, this shouldn't happen")
            return

        count = self.cache.load_from_vault(self.vault)
        logger.info(f"Reloaded {count} secrets from vault")

    def _cleanup(self):
        """Clean up resources."""
        if self.server_socket:
            self.server_socket.close()

        if self.socket_path.exists():
            self.socket_path.unlink()

        if self.pid_file.exists():
            self.pid_file.unlink()

        # Wait for handlers to finish
        for thread in self._threads:
            thread.join(timeout=1.0)

        logger.info("Daemon stopped")


def main():
    """CLI entry point."""
    import argparse

    parser = argparse.ArgumentParser(
        description="Liberty Daemon - Secure secret server",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
    # Start daemon (as liberty-user)
    liberty-daemon start

    # Start in foreground for debugging
    liberty-daemon start --foreground

    # Stop daemon
    liberty-daemon stop

    # Check status
    liberty-daemon status
"""
    )

    parser.add_argument(
        "command",
        choices=["start", "stop", "status"],
        help="Command to run"
    )
    parser.add_argument(
        "--socket",
        default=DEFAULT_SOCKET_PATH,
        help=f"Socket path (default: {DEFAULT_SOCKET_PATH})"
    )
    parser.add_argument(
        "--vault",
        default=DEFAULT_VAULT_PATH,
        help=f"Vault path (default: {DEFAULT_VAULT_PATH})"
    )
    parser.add_argument(
        "--group",
        default="agents",
        help="Socket group for agent access (default: agents)"
    )
    parser.add_argument(
        "--foreground", "-f",
        action="store_true",
        help="Run in foreground (don't daemonize)"
    )
    parser.add_argument(
        "--pid",
        default=None,
        help="PID file path (default: alongside socket)"
    )
    parser.add_argument(
        "--allow-write",
        action="store_true",
        help="Allow clients to store/delete secrets (required for teleport import)"
    )

    args = parser.parse_args()

    # Derive PID path from socket path if not specified
    pid_file = args.pid
    if pid_file is None:
        socket_dir = Path(args.socket).parent
        pid_file = str(socket_dir / "liberty.pid")

    daemon = LibertyDaemon(
        socket_path=args.socket,
        vault_path=args.vault,
        socket_group=args.group,
        pid_file=pid_file,
        allow_write=args.allow_write,
    )

    if args.command == "start":
        daemon.start(foreground=args.foreground)

    elif args.command == "stop":
        daemon.stop()

    elif args.command == "status":
        if daemon._is_running():
            pid = daemon.pid_file.read_text().strip()
            print(f"Liberty daemon is running (PID: {pid})")
            print(f"Socket: {daemon.socket_path}")
        else:
            print("Liberty daemon is not running")
            sys.exit(1)


if __name__ == "__main__":
    main()
