#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Orchestrator Security: Boot verification and quarantine enforcement.

Extracted from pulse.py per Law 7 (files <= 300 lines).

Functions:
  - verify_sealed_environment(): RSA signature check on .env
  - verify_immutable_core(): SHA-256 hash check on Constitution + Auth
  - enforce_quarantine(): Kill switch for quarantined agents
  - load_env_file(): Load .env into os.environ
"""
import hashlib
import json
import os
import subprocess
import sys
import time
from pathlib import Path

from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.exceptions import InvalidSignature

SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent.parent

from pulse.core.paths import _is_dev_mode, get_root_dir

QUARANTINE_FILE = ROOT_DIR / "state" / "quarantined_entities.json"

from pulse.core.brain_utils import load_json


def _resolve_core_paths():
    """Resolve Constitution, Auth, .env paths for dev or installed mode."""
    root = get_root_dir()
    return {
        "constitution": root / "Corpus_Callosum" / "CONSTITUTION.md",
        "auth": root / "Corpus_Callosum" / "authorized_humans.json",
        "env": root / ".env",
        "env_sig": root / ".env.sig",
    }


def compute_script_hash(script_path: Path) -> str:
    """C7: Compute SHA-256 hash of a daemon script for identity verification."""
    try:
        with open(script_path, "rb") as f:
            return hashlib.sha256(f.read()).hexdigest()
    except (OSError, IOError):
        return ""


def verify_sealed_environment():
    """Verify the .env file against its RSA signature. Exits on failure."""
    paths = _resolve_core_paths()
    env_file = paths["env"]
    env_sig_file = paths["env_sig"]
    auth_file = paths["auth"]

    print("[PULSE] Verifying Sealed Environment...")

    # Installed mode: no RSA sealing expected — verify .env exists, skip sig check
    if not _is_dev_mode():
        if not env_file.exists():
            print("[PULSE] Installed mode — no .env found. Skipping seal check.")
        else:
            print("[PULSE] Installed mode — .env present (no RSA seal required).")
        return

    if not env_file.exists() or not env_sig_file.exists():
        print("CRITICAL ERROR: .env or .env.sig MISSING!")
        print("Run 'seal_env.py' to secure your environment.")
        sys.exit(1)

    founder_keys = []
    try:
        with open(auth_file, "r") as f:
            auth_data = json.load(f)
            founder_keys = [
                user["public_key"] for user in auth_data.values()
                if user.get("role") == "FOUNDER"
            ]
        if not founder_keys:
            print("CRITICAL ERROR: No FOUNDER public keys found!")
            sys.exit(1)
    except Exception as e:
        print(f"CRITICAL ERROR: Could not load authorization keys. {e}")
        sys.exit(1)

    try:
        with open(env_sig_file, "r") as f:
            signature_hex = f.read().strip()
        with open(env_file, "rb") as f:
            env_hash = hashlib.sha256(f.read()).hexdigest()

        verified = False
        for pem_key in founder_keys:
            try:
                public_key = serialization.load_pem_public_key(pem_key.encode())
                public_key.verify(
                    bytes.fromhex(signature_hex),
                    env_hash.encode(),
                    padding.PSS(
                        mgf=padding.MGF1(hashes.SHA256()),
                        salt_length=padding.PSS.MAX_LENGTH
                    ),
                    hashes.SHA256()
                )
                verified = True
                break
            except (InvalidSignature, ValueError):
                continue

        if verified:
            print("[PULSE] Environment Seal Verified.")
        else:
            raise InvalidSignature("No authorized key could verify the signature.")

    except (InvalidSignature, ValueError):
        print("\n" + "!" * 60)
        print("CRITICAL SECURITY ALERT: ENVIRONMENT (.env) TAMPERED")
        print("The signature in .env.sig does not match the content of .env.")
        print("The system refuses to boot.")
        print("!" * 60 + "\n")
        sys.exit(1)
    except Exception as e:
        print(f"CRITICAL ERROR: Unexpected error during env verification: {e}")
        sys.exit(1)


def verify_immutable_core():
    """Verify Constitution and Authorization file hashes. Exits on mismatch."""
    paths = _resolve_core_paths()
    constitution_file = paths["constitution"]
    auth_file = paths["auth"]

    print("[PULSE] Verifying Immutable Core...")

    files_to_check = {
        "CONSTITUTION": (constitution_file, os.getenv("CONSTITUTION_HASH")),
    }
    # Auth file only exists in dev mode (installed mode has no authorized_humans.json)
    if _is_dev_mode():
        files_to_check["AUTHORIZATION"] = (auth_file, os.getenv("AUTH_HASH"))

    for name, (file_path, expected_hash) in files_to_check.items():
        if not file_path.exists():
            if not _is_dev_mode():
                print(f"WARNING: {name} file not found at {file_path} — skipping.")
                continue
            print(f"CRITICAL ERROR: {name} FILE MISSING!")
            sys.exit(1)
        with open(file_path, "rb") as f:
            file_hash = hashlib.sha256(f.read()).hexdigest()
        if not expected_hash:
            print(f"WARNING: No {name}_HASH set. Calculated: {file_hash}")
            continue
        if file_hash != expected_hash:
            print("\n" + "!" * 60)
            print(f"CRITICAL SECURITY ALERT: {name} FILE TAMPERED")
            print("The system refuses to boot until integrity is restored.")
            print("!" * 60 + "\n")
            sys.exit(1)

    print("[PULSE] Immutable Core Verified.")


def load_env_file():
    """Load .env variables into os.environ manually."""
    env_file = _resolve_core_paths()["env"]
    if not env_file.exists():
        return
    try:
        with open(env_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                if "=" in line:
                    key, value = line.split("=", 1)
                    value = value.strip()
                    if len(value) >= 2 and value[0] == value[-1] and value[0] in ('"', "'"):
                        value = value[1:-1]
                    os.environ[key.strip()] = value
    except Exception as e:
        print(f"[PULSE] Warning: Could not load .env file: {e}")


def enforce_quarantine(processes: dict, quarantine_events: dict, daemon_map: dict):
    """C7: Quarantine by agent_id AND script hash to prevent rename bypass.

    Args:
        processes: dict of {agent_id: Popen} — mutated in place
        quarantine_events: dict of {agent_id: [timestamps]} — mutated in place
        daemon_map: dict of {agent_id: Path} — script paths
    """
    if not QUARANTINE_FILE.exists():
        return

    try:
        quarantined_list = load_json(QUARANTINE_FILE)
        if not isinstance(quarantined_list, list):
            return
        quarantined_set = set(quarantined_list)
        now = time.time()

        for agent_id, process in list(processes.items()):
            name_match = agent_id in quarantined_set

            script = daemon_map.get(agent_id)
            hash_match = False
            if script and script.exists():
                script_hash = compute_script_hash(script)
                if script_hash and script_hash in quarantined_set:
                    hash_match = True

            if name_match or hash_match:
                match_type = "hash" if hash_match and not name_match else "name"
                print(f"\n[KILL SWITCH] Terminating Quarantined Agent: {agent_id} (matched by {match_type})")
                process.terminate()
                try:
                    process.wait(timeout=1)
                except subprocess.TimeoutExpired:
                    process.kill()
                del processes[agent_id]
                print(f"[KILL SWITCH] Agent {agent_id} terminated.")

                quarantine_events[agent_id].append(now)
                recent_kills = [t for t in quarantine_events[agent_id] if now - t < 300]
                if len(recent_kills) >= 3:
                    print("\n" + "!" * 60)
                    print(f"COMPANY-LEVEL ALERT: Agent {agent_id} respawning rapidly.")
                    print("!" * 60 + "\n")

    except Exception as e:
        print(f"[PULSE] Quarantine Check Error: {e}")
