import os
import hashlib
import logging
from uuid import uuid4, getnode
from pathlib import Path
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes

from soagila.core.encryption.interface import EncryptionInterface

# Optional OS keyring integration. When available, the per-project
# ``master.salt`` is mirrored into Windows Credential Manager / macOS
# Keychain / Linux Secret Service so a stolen disk without the OS
# credentials is insufficient to recover memories.
try:
    import keyring  # type: ignore
    _HAVE_KEYRING = True
except Exception:  # pragma: no cover
    keyring = None  # type: ignore
    _HAVE_KEYRING = False

_KEYRING_SERVICE = "soagila-daemon"

logger = logging.getLogger("soagila.encryption.simulated_he")

# IV size for AES-GCM (96 bits recommended by NIST)
IV_SIZE = 12
# AES key size (256 bits)
KEY_SIZE = 32
# Salt size for PBKDF2
SALT_SIZE = 32
# PBKDF2 iterations
PBKDF2_ITERATIONS = 480_000


def _get_machine_id() -> str:
    """Get a deterministic machine identifier (MAC address + username)."""
    mac = getnode()
    user = os.environ.get("USERNAME", os.environ.get("USER", "soagila"))
    return f"{mac}-{user}"


def _derive_master_key(machine_id: str, project_hash: str, salt: bytes) -> bytes:
    """Derive a master key using PBKDF2 from machine identity + project hash."""
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=KEY_SIZE,
        salt=salt,
        iterations=PBKDF2_ITERATIONS,
    )
    material = f"{machine_id}:{project_hash}".encode("utf-8")
    return kdf.derive(material)


class SimulatedHEEncryption(EncryptionInterface):
    """
    AES-256-GCM encryption wrapper that simulates the Homomorphic Encryption interface.
    In production, this would be replaced by Microsoft SEAL or OpenFHE.

    Key hierarchy (simulated HSM):
    - master_key: Derived from machine identity via PBKDF2 (protects key store)
    - org_master_key: Organization-level master key
    - user_mek: User Memory Encryption Key (derived from org master)
    - session_key: Ephemeral session key

    Keys are persisted to SQLite `encryption_keys` table, encrypted with master_key.

    Multi-project isolation (v0.7.0): the engine keeps per-project caches
    so the daemon can serve multiple workspaces without one project's
    master key bleeding into another. All mutable state keyed by the
    current ``project_hash`` resolved from either the argument or the
    async-local ``sqlite_init.current_project_hash`` ContextVar.
    """

    def __init__(self):
        # Per-project caches. Key = project_hash, value = dict of data
        self._key_stores: dict[str, dict[str, bytes]] = {}
        self._key_purposes_stores: dict[str, dict[str, str]] = {}
        self._master_keys: dict[str, bytes] = {}
        logger.info("Simulated HE encryption engine initialized (AES-256-GCM, multi-project)")

    # ── Project scoping helpers ─────────────────────────────────────

    def _current_project_hash(self) -> str:
        """Resolve the active project_hash with a sane default."""
        try:
            from soagila.db.sqlite_init import current_project_hash as cv
            pinned = cv.get()
            if pinned:
                return pinned
        except Exception:
            pass
        # Fallback: ask the config module (legacy single-project mode)
        from soagila.config import get_project_hash
        return get_project_hash()

    @property
    def _key_store(self) -> dict[str, bytes]:  # backward-compat accessor
        ph = self._current_project_hash()
        return self._key_stores.setdefault(ph, {})

    @property
    def _key_purposes(self) -> dict[str, str]:
        ph = self._current_project_hash()
        return self._key_purposes_stores.setdefault(ph, {})

    def _ensure_master_key(self) -> bytes:
        """Get or create the master key for the ACTIVE project.

        Salt lookup order:
            1. OS keyring (Windows Credential Manager / macOS Keychain /
               Linux Secret Service) — keyed on ``salt:<project_hash>``
            2. Legacy on-disk ``master.salt`` file next to the project DB
            3. Freshly generated 32-byte random salt (persisted to both)
        """
        ph = self._current_project_hash()
        existing = self._master_keys.get(ph)
        if existing is not None:
            return existing

        from soagila.config import get_soagila_home
        project_dir = get_soagila_home() / "projects" / ph
        project_dir.mkdir(parents=True, exist_ok=True)
        salt_path = project_dir / "master.salt"

        machine_id = _get_machine_id()
        salt: bytes | None = None
        keyring_entry = f"salt:{ph}"

        # 1. Try OS keyring first
        if _HAVE_KEYRING:
            try:
                stored_hex = keyring.get_password(_KEYRING_SERVICE, keyring_entry)
                if stored_hex:
                    salt = bytes.fromhex(stored_hex)
            except Exception as exc:
                logger.debug(f"Keyring read failed for {keyring_entry}: {exc}")

        # 2. Fall back to on-disk salt (legacy installs)
        if salt is None and salt_path.exists():
            salt = salt_path.read_bytes()
            # Migrate into keyring for next time
            if _HAVE_KEYRING:
                try:
                    keyring.set_password(_KEYRING_SERVICE, keyring_entry, salt.hex())
                    logger.info(f"Migrated salt for {ph} to OS keyring")
                except Exception as exc:
                    logger.debug(f"Keyring migrate failed: {exc}")

        # 3. Generate a new one
        if salt is None:
            salt = os.urandom(SALT_SIZE)
            persisted = False
            if _HAVE_KEYRING:
                try:
                    keyring.set_password(_KEYRING_SERVICE, keyring_entry, salt.hex())
                    logger.info(f"Generated new master key salt in OS keyring for {ph}")
                    persisted = True
                except Exception as exc:
                    logger.warning(
                        f"Keyring unavailable ({exc}); falling back to on-disk salt"
                    )
            if not persisted:
                salt_path.write_bytes(salt)
                logger.info(f"Generated new master key salt (file): {salt_path}")

        master = _derive_master_key(machine_id, ph, salt)
        self._master_keys[ph] = master
        return master

    def _encrypt_key_for_storage(self, raw_key: bytes) -> bytes:
        """Encrypt a key with the master key for SQLite storage."""
        master = self._ensure_master_key()
        iv = os.urandom(IV_SIZE)
        aesgcm = AESGCM(master)
        ciphertext = aesgcm.encrypt(iv, raw_key, None)
        return iv + ciphertext

    def _decrypt_key_from_storage(self, encrypted_key: bytes) -> bytes:
        """Decrypt a key from SQLite storage using the master key."""
        master = self._ensure_master_key()
        iv = encrypted_key[:IV_SIZE]
        ciphertext = encrypted_key[IV_SIZE:]
        aesgcm = AESGCM(master)
        return aesgcm.decrypt(iv, ciphertext, None)

    async def persist_key(self, key_ref: str, raw_key: bytes, purpose: str):
        """Persist an encryption key to SQLite, encrypted with master key."""
        from soagila.db.sqlite_init import get_db
        encrypted = self._encrypt_key_for_storage(raw_key)
        db = await get_db()
        await db.execute(
            "INSERT OR REPLACE INTO encryption_keys (key_ref, encrypted_key, purpose) VALUES (?, ?, ?)",
            (key_ref, encrypted, purpose),
        )
        await db.commit()

    async def load_persisted_keys(self):
        """Load all persisted keys from SQLite into memory. Called on startup."""
        from soagila.db.sqlite_init import get_db
        try:
            db = await get_db()
            cursor = await db.execute("SELECT key_ref, encrypted_key, purpose FROM encryption_keys")
            rows = await cursor.fetchall()
            loaded = 0
            for row in rows:
                key_ref, encrypted_key, purpose = row[0], row[1], row[2]
                try:
                    raw_key = self._decrypt_key_from_storage(encrypted_key)
                    self._key_store[key_ref] = raw_key
                    self._key_purposes[key_ref] = purpose
                    loaded += 1
                except Exception as e:
                    logger.warning(f"Failed to decrypt key {key_ref}: {e}")
            if loaded > 0:
                logger.info(f"Loaded {loaded} encryption keys from SQLite")
        except Exception as e:
            logger.warning(f"Failed to load persisted keys: {e}")

    def encrypt(self, plaintext: str, key_ref: str) -> bytes:
        """
        Encrypt plaintext with AES-256-GCM.
        Returns: IV (12 bytes) + ciphertext + auth tag (16 bytes)
        """
        if key_ref not in self._key_store:
            raise ValueError(f"Key reference '{key_ref}' not found in key store")

        key = self._key_store[key_ref]
        iv = os.urandom(IV_SIZE)
        aesgcm = AESGCM(key)

        ciphertext = aesgcm.encrypt(iv, plaintext.encode("utf-8"), None)

        # Format: IV + ciphertext (includes auth tag appended by AESGCM)
        return iv + ciphertext

    def decrypt(self, ciphertext: bytes, key_ref: str) -> str:
        """
        Decrypt AES-256-GCM ciphertext.
        Input: IV (12 bytes) + ciphertext + auth tag
        """
        if key_ref not in self._key_store:
            raise ValueError(f"Key reference '{key_ref}' not found in key store")

        key = self._key_store[key_ref]
        iv = ciphertext[:IV_SIZE]
        encrypted_data = ciphertext[IV_SIZE:]

        aesgcm = AESGCM(key)
        plaintext = aesgcm.decrypt(iv, encrypted_data, None)

        return plaintext.decode("utf-8")

    def generate_key(self, purpose: str) -> str:
        """
        Generate a new AES-256 key and return its reference.
        Key is persisted to SQLite encrypted with master key.
        Purpose: 'session' | 'user_mek' | 'org_master' | 'memory'
        """
        key_ref = f"key-{purpose}-{uuid4().hex[:12]}"
        key = AESGCM.generate_key(bit_length=256)

        self._key_store[key_ref] = key
        self._key_purposes[key_ref] = purpose

        # Schedule persistence (will be awaited by the caller's event loop)
        import asyncio
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(self.persist_key(key_ref, key, purpose))
        except RuntimeError:
            # No event loop — skip persistence (test context)
            pass

        logger.debug(f"Generated key: {key_ref} (purpose={purpose})")
        return key_ref

    def destroy_key(self, key_ref: str) -> bool:
        """
        Cryptographic annihilation: destroy the key.
        All data encrypted with this key becomes permanently irrecoverable.
        """
        if key_ref not in self._key_store:
            return False

        # Overwrite key in memory before deletion (defense-in-depth)
        key_len = len(self._key_store[key_ref])
        self._key_store[key_ref] = os.urandom(key_len)  # overwrite with random
        del self._key_store[key_ref]
        del self._key_purposes[key_ref]

        # Delete from SQLite
        import asyncio
        try:
            loop = asyncio.get_running_loop()
            loop.create_task(self._delete_persisted_key(key_ref))
        except RuntimeError:
            pass

        logger.info(f"Key destroyed (annihilated): {key_ref}")
        return True

    async def _delete_persisted_key(self, key_ref: str):
        """Delete a key from SQLite persistence."""
        try:
            from soagila.db.sqlite_init import get_db
            db = await get_db()
            await db.execute("DELETE FROM encryption_keys WHERE key_ref = ?", (key_ref,))
            await db.commit()
        except Exception as e:
            logger.warning(f"Failed to delete persisted key {key_ref}: {e}")

    def key_exists(self, key_ref: str) -> bool:
        """Check if a key reference exists."""
        return key_ref in self._key_store

    def get_key_purpose(self, key_ref: str) -> str | None:
        """Get the purpose of a key."""
        return self._key_purposes.get(key_ref)

    def get_stats(self) -> dict:
        """Get key store statistics."""
        purpose_counts: dict[str, int] = {}
        for purpose in self._key_purposes.values():
            purpose_counts[purpose] = purpose_counts.get(purpose, 0) + 1

        return {
            "total_keys": len(self._key_store),
            "by_purpose": purpose_counts,
        }


# Singleton instance
_encryption_engine: SimulatedHEEncryption | None = None


def get_encryption_engine() -> SimulatedHEEncryption:
    """Get the singleton encryption engine."""
    global _encryption_engine
    if _encryption_engine is None:
        _encryption_engine = SimulatedHEEncryption()
    return _encryption_engine
