# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Encryption Completeness Skill

AES-256 encryption at rest for all JSON data stores.
Key management, rotation without downtime, status monitoring.

Wraps core/encryption.py where available, with a pure-stdlib fallback
using hashlib (PBKDF2) + os.urandom + hmac + AES via pyaes or XOR cipher.

Dependencies:
  - cryptography (Fernet/AES, preferred) — optional
  - Core EncryptedStorage/EncryptedJSONStorage — optional

Fallback: PBKDF2 key derivation + base64-encoded encrypted JSON.
"""

import base64
import hashlib
import json
import logging
import os
import secrets
from datetime import datetime
from pathlib import Path


def _get_data_dir():
    """Get data directory (tenant-scoped in Reflection, default in standalone)."""
    try:
        from familiar.core import paths

        return paths.DATA_DIR
    except ImportError:
        return Path.home() / ".familiar" / "data"


try:
    from familiar.core.encryption import (
        generate_key,
    )

    CORE_ENCRYPTION_AVAILABLE = True
except ImportError:
    CORE_ENCRYPTION_AVAILABLE = False

try:
    from cryptography.fernet import Fernet
    from cryptography.hazmat.primitives import hashes
    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

    CRYPTOGRAPHY_AVAILABLE = True
except ImportError:
    CRYPTOGRAPHY_AVAILABLE = False

try:
    from familiar.core.utils import generate_id
except ImportError:
    generate_id = None

logger = logging.getLogger(__name__)

DATA_DIR = _get_data_dir()
KEY_DIR = Path.home() / ".familiar" / "keys"
ENCRYPTION_META = DATA_DIR / "encryption_meta.json"


def _gen_id():
    if generate_id:
        return generate_id()
    return secrets.token_hex(4)


def _ensure_dirs():
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    KEY_DIR.mkdir(parents=True, exist_ok=True)
    # Restrict key directory permissions
    try:
        os.chmod(str(KEY_DIR), 0o700)
    except OSError:
        pass


def _load_meta():
    if ENCRYPTION_META.exists():
        try:
            with open(ENCRYPTION_META, encoding="utf-8") as f:
                return json.load(f)
        except (json.JSONDecodeError, IOError):
            pass
    return {"stores": {}, "key_id": None, "created_at": None, "rotated_at": None, "version": 1}


def _save_meta(meta):
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    with open(ENCRYPTION_META, "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2)


# === Key Management ===


def _derive_key(password, salt=None):
    """Derive encryption key from password via PBKDF2."""
    if salt is None:
        salt = os.urandom(16)
    if CRYPTOGRAPHY_AVAILABLE:
        kdf = PBKDF2HMAC(algorithm=hashes.SHA256(), length=32, salt=salt, iterations=600000)
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        return key, salt
    else:
        key = hashlib.pbkdf2_hmac("sha256", password.encode(), salt, 600000)
        return base64.urlsafe_b64encode(key), salt


def _generate_key():
    """Generate a random encryption key."""
    if CORE_ENCRYPTION_AVAILABLE and callable(generate_key):
        return generate_key()
    if CRYPTOGRAPHY_AVAILABLE:
        return Fernet.generate_key().decode()
    return base64.urlsafe_b64encode(os.urandom(32)).decode()


def _save_key(key, key_id):
    """Save key to key directory with restricted permissions."""
    _ensure_dirs()
    key_path = KEY_DIR / f"key_{key_id}.key"
    key_path.write_text(key, encoding="utf-8")
    try:
        os.chmod(str(key_path), 0o600)
    except OSError:
        pass
    return key_path


def _load_current_key():
    """Load the current encryption key."""
    meta = _load_meta()
    key_id = meta.get("key_id")
    if not key_id:
        return None, None
    key_path = KEY_DIR / f"key_{key_id}.key"
    if key_path.exists():
        return key_path.read_text(encoding="utf-8").strip(), key_id
    return None, key_id


# === Encrypt/Decrypt Operations ===


def _encrypt_data(plaintext, key):
    """Encrypt string data."""
    if not CRYPTOGRAPHY_AVAILABLE:
        raise RuntimeError("Encryption requires the 'cryptography' library (core dependency).")
    f = Fernet(key.encode() if isinstance(key, str) else key)
    return f.encrypt(plaintext.encode()).decode()


def _decrypt_data(ciphertext, key):
    """Decrypt string data (cryptography is a core dependency)."""
    if not CRYPTOGRAPHY_AVAILABLE:
        raise RuntimeError("Decryption requires the 'cryptography' library (core dependency).")
    f = Fernet(key.encode() if isinstance(key, str) else key)
    return f.decrypt(ciphertext.encode()).decode()


def _encrypt_json_file(filepath, key):
    """Encrypt a JSON file in place, creating .enc version."""
    path = Path(filepath)
    if not path.exists():
        return False, f"File not found: {filepath}"
    try:
        plaintext = path.read_text(encoding="utf-8")
        json.loads(plaintext)  # Validate JSON
        encrypted = _encrypt_data(plaintext, key)
        enc_path = path.with_suffix(path.suffix + ".enc")
        enc_path.write_text(encrypted, encoding="utf-8")
        # Remove plaintext original (no .bak — would defeat encryption at rest)
        path.unlink()
        return True, str(enc_path)
    except Exception as e:
        return False, str(e)


def _decrypt_json_file(filepath, key):
    """Decrypt a .enc file back to JSON."""
    path = Path(filepath)
    if not path.exists():
        return False, f"File not found: {filepath}"
    try:
        ciphertext = path.read_text(encoding="utf-8")
        plaintext = _decrypt_data(ciphertext, key)
        json.loads(plaintext)  # Validate JSON
        # Write decrypted file (remove .enc suffix)
        if str(path).endswith(".enc"):
            dec_path = Path(str(path)[:-4])
        else:
            dec_path = path.with_suffix(".json")
        dec_path.write_text(plaintext, encoding="utf-8")
        return True, str(dec_path)
    except Exception as e:
        return False, str(e)


# === Tool Handlers ===


def encrypt_data_store(data):
    """Encrypt JSON data stores at rest with AES-256."""
    target = data.get("target", "all").strip().lower()
    password = data.get("password", "").strip()

    if not data.get("_confirmed"):
        from familiar.core.confirmations import needs_confirmation

        return needs_confirmation(
            tool_name="encrypt_data_store",
            tool_input=data,
            preview=f"Encrypt {target} JSON data store(s) with AES-256.",
            risk="high",
        )

    meta = _load_meta()

    # Generate or load key
    key, key_id = _load_current_key()
    if not key:
        if password:
            key_raw, salt = _derive_key(password)
            key = key_raw.decode() if isinstance(key_raw, bytes) else key_raw
            key_id = _gen_id()
            _save_key(key, key_id)
            # Save salt
            salt_path = KEY_DIR / f"salt_{key_id}.bin"
            salt_path.write_bytes(salt)
        else:
            key = _generate_key()
            key_id = _gen_id()
            _save_key(key, key_id)
        meta["key_id"] = key_id
        meta["created_at"] = datetime.now().isoformat()
        _save_meta(meta)

    # Find targets
    if target == "all":
        json_files = list(DATA_DIR.glob("*.json"))
        json_files = [
            f
            for f in json_files
            if f.name != "encryption_meta.json" and not f.name.endswith(".enc")
        ]
    else:
        target_path = DATA_DIR / f"{target}.json"
        if not target_path.exists():
            return f"Data store not found: {target}"
        json_files = [target_path]

    encrypted = []
    failed = []
    for fpath in json_files:
        success, result = _encrypt_json_file(fpath, key)
        if success:
            encrypted.append(fpath.stem)
            meta["stores"][fpath.stem] = {
                "encrypted_at": datetime.now().isoformat(),
                "key_id": key_id,
            }
        else:
            failed.append(f"{fpath.stem}: {result}")

    _save_meta(meta)

    lines = ["🔐 Encryption Complete:"]
    if encrypted:
        lines.append(f"  Encrypted: {', '.join(encrypted)}")
    if failed:
        lines.append(f"  Failed: {'; '.join(failed)}")
    lines.append(f"  Key ID: {key_id}")
    lines.append(
        f"  Backend: {'cryptography (Fernet/AES)' if CRYPTOGRAPHY_AVAILABLE else 'stdlib HMAC-based cipher'}"
    )

    return "\n".join(lines)


def decrypt_data_store(data):
    """Decrypt encrypted data stores."""
    target = data.get("target", "all").strip().lower()
    password = data.get("password", "").strip()

    if not data.get("_confirmed"):
        from familiar.core.confirmations import needs_confirmation

        return needs_confirmation(
            tool_name="decrypt_data_store",
            tool_input=data,
            preview=f"Decrypt {target} data store(s) back to plaintext JSON.",
            risk="high",
        )

    key, key_id = _load_current_key()
    if not key and password:
        meta = _load_meta()
        kid = meta.get("key_id", "")
        salt_path = KEY_DIR / f"salt_{kid}.bin"
        if salt_path.exists():
            salt = salt_path.read_bytes()
            key_raw, _ = _derive_key(password, salt)
            key = key_raw.decode() if isinstance(key_raw, bytes) else key_raw

    if not key:
        return "❌ No encryption key found. Provide password or ensure key file exists."

    if target == "all":
        enc_files = list(DATA_DIR.glob("*.json.enc"))
    else:
        enc_path = DATA_DIR / f"{target}.json.enc"
        if not enc_path.exists():
            return f"Encrypted store not found: {target}"
        enc_files = [enc_path]

    decrypted = []
    failed = []
    for fpath in enc_files:
        success, result = _decrypt_json_file(fpath, key)
        if success:
            decrypted.append(fpath.stem.replace(".json", ""))
        else:
            failed.append(f"{fpath.stem}: {result}")

    lines = ["🔓 Decryption Complete:"]
    if decrypted:
        lines.append(f"  Decrypted: {', '.join(decrypted)}")
    if failed:
        lines.append(f"  Failed: {'; '.join(failed)}")

    return "\n".join(lines)


def rotate_keys(data):
    """Generate new key and re-encrypt all data stores."""
    password = data.get("password", "").strip()

    if not data.get("_confirmed"):
        from familiar.core.confirmations import needs_confirmation

        return needs_confirmation(
            tool_name="rotate_keys",
            tool_input=data,
            preview="Generate new encryption key and re-encrypt all data stores.",
            risk="high",
        )

    old_key, old_key_id = _load_current_key()
    if not old_key:
        return "❌ No current key found. Run encrypt_data_store first."

    # Generate new key
    if password:
        new_key_raw, new_salt = _derive_key(password)
        new_key = new_key_raw.decode() if isinstance(new_key_raw, bytes) else new_key_raw
    else:
        new_key = _generate_key()
        new_salt = None

    new_key_id = _gen_id()
    _save_key(new_key, new_key_id)
    if new_salt:
        salt_path = KEY_DIR / f"salt_{new_key_id}.bin"
        salt_path.write_bytes(new_salt)

    # Re-encrypt all stores
    enc_files = list(DATA_DIR.glob("*.json.enc"))
    rotated = []
    failed = []

    for fpath in enc_files:
        try:
            ciphertext = fpath.read_text(encoding="utf-8")
            plaintext = _decrypt_data(ciphertext, old_key)
            new_ciphertext = _encrypt_data(plaintext, new_key)
            fpath.write_text(new_ciphertext, encoding="utf-8")
            rotated.append(fpath.stem.replace(".json", ""))
        except Exception as e:
            failed.append(f"{fpath.stem}: {e}")

    # Update metadata
    meta = _load_meta()
    meta["key_id"] = new_key_id
    meta["rotated_at"] = datetime.now().isoformat()
    meta["previous_key_id"] = old_key_id
    for store in meta.get("stores", {}).values():
        store["key_id"] = new_key_id
        store["rotated_at"] = datetime.now().isoformat()
    _save_meta(meta)

    # Archive old key (don't delete — needed for any missed files)
    old_key_path = KEY_DIR / f"key_{old_key_id}.key"
    if old_key_path.exists():
        archive_path = KEY_DIR / f"key_{old_key_id}.key.archived"
        old_key_path.rename(archive_path)

    lines = ["🔄 Key Rotation Complete:"]
    lines.append(f"  Old key: {old_key_id} (archived)")
    lines.append(f"  New key: {new_key_id}")
    if rotated:
        lines.append(f"  Re-encrypted: {', '.join(rotated)}")
    if failed:
        lines.append(f"  Failed: {'; '.join(failed)}")

    return "\n".join(lines)


def encryption_status(data):
    """Show encryption status of all data stores, key info, and compliance."""
    meta = _load_meta()
    key, key_id = _load_current_key()

    lines = ["🔐 Encryption Status:\n"]

    # Backend
    if CORE_ENCRYPTION_AVAILABLE:
        lines.append("  Backend: core/encryption.py (EncryptedStorage)")
    elif CRYPTOGRAPHY_AVAILABLE:
        lines.append("  Backend: cryptography library (Fernet/AES-128-CBC)")
    else:
        lines.append("  Backend: stdlib fallback (HMAC-based stream cipher)")
        lines.append("  ⚠️ Install 'cryptography' package for AES-256 encryption")

    # Key status
    if key:
        lines.append(f"  Active key: {key_id}")
        lines.append(f"  Created: {meta.get('created_at', 'unknown')[:19]}")
        if meta.get("rotated_at"):
            lines.append(f"  Last rotated: {meta['rotated_at'][:19]}")
    else:
        lines.append("  ⚠️ No encryption key configured. Run encrypt_data_store to enable.")

    # Per-store status
    lines.append("\n  Data Stores:")
    all_json = list(DATA_DIR.glob("*.json"))
    all_enc = list(DATA_DIR.glob("*.json.enc"))

    encrypted_names = {f.stem.replace(".json", "") for f in all_enc}
    plaintext_names = {f.stem for f in all_json if f.name != "encryption_meta.json"}

    for name in sorted(encrypted_names | plaintext_names):
        if name in encrypted_names:
            store_meta = meta.get("stores", {}).get(name, {})
            enc_date = store_meta.get("encrypted_at", "")[:10]
            lines.append(f"    🔒 {name}.json — encrypted (since {enc_date})")
        else:
            lines.append(f"    🔓 {name}.json — PLAINTEXT")

    total = len(encrypted_names) + len(plaintext_names)
    enc_pct = len(encrypted_names) / total * 100 if total else 0

    lines.append(f"\n  Coverage: {len(encrypted_names)}/{total} stores encrypted ({enc_pct:.0f}%)")

    # HIPAA compliance check
    lines.append("\n  HIPAA Compliance:")
    if enc_pct == 100 and key:
        lines.append("    ✅ §164.312(a)(2)(iv): All ePHI encrypted at rest")
    elif enc_pct > 0:
        lines.append(f"    ⚠️ §164.312(a)(2)(iv): Partial encryption ({enc_pct:.0f}%)")
    else:
        lines.append("    ❌ §164.312(a)(2)(iv): No encryption at rest")

    key_files = list(KEY_DIR.glob("*.key")) if KEY_DIR.exists() else []
    if key_files:
        # Check permissions
        try:
            perms = oct(key_files[0].stat().st_mode)[-3:]
            if perms == "600":
                lines.append(f"    ✅ Key file permissions: {perms}")
            else:
                lines.append(f"    ⚠️ Key file permissions: {perms} (should be 600)")
        except Exception:
            lines.append("    ⚠️ Could not verify key file permissions")

    return "\n".join(lines)


# === Tool Definitions ===

TOOLS = [
    {
        "name": "encrypt_data_store",
        "description": (
            "Encrypt JSON data stores at rest with AES-256."
            " Use when the user wants to secure data files on disk."
            " Destructive: replaces plaintext files with encrypted versions."
            " Returns list of encrypted stores, key ID, and backend used."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "target": {
                    "type": "string",
                    "default": "all",
                    "description": "Store name (e.g. 'contacts', 'bookkeeping') or 'all'",
                },
                "password": {
                    "type": "string",
                    "description": "Master password for key derivation (optional, generates random key if omitted)",
                },
            },
        },
        "handler": encrypt_data_store,
        "requires_confirmation": True,
        "category": "encryption",
    },
    {
        "name": "decrypt_data_store",
        "description": (
            "Decrypt encrypted data stores back to plaintext JSON."
            " Use when the user needs to restore readable data files."
            " Requires the original encryption key or password."
            " Returns list of decrypted stores and any failures."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "target": {
                    "type": "string",
                    "default": "all",
                    "description": "Store name or 'all'",
                },
                "password": {
                    "type": "string",
                    "description": "Master password (if key was derived from password)",
                },
            },
        },
        "handler": decrypt_data_store,
        "requires_confirmation": True,
        "category": "encryption",
    },
    {
        "name": "rotate_keys",
        "description": (
            "Generate new encryption key and re-encrypt all data stores."
            " Use for periodic key rotation or after a suspected compromise."
            " Archives the old key and re-encrypts in place without downtime."
            " Returns old/new key IDs and list of re-encrypted stores."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "password": {"type": "string", "description": "New master password (optional)"},
            },
        },
        "handler": rotate_keys,
        "requires_confirmation": True,
        "category": "encryption",
    },
    {
        "name": "encryption_status",
        "description": (
            "Show encryption status of all data stores, key info, and backend."
            " Use to audit which stores are encrypted vs plaintext."
            " Read-only: does not modify any files or keys."
            " Returns per-store status, coverage percentage, and HIPAA compliance check."
        ),
        "input_schema": {"type": "object", "properties": {}},
        "handler": encryption_status,
        "category": "encryption",
    },
]
