"""Local audit log - full event history written to disk.

This stays on the user's machine. We never see it.
Full content hashes, timestamps, everything needed for forensics.
"""
import base64
import hashlib
import hmac as hmac_mod
import json
import os
import threading
from datetime import datetime, timezone
from pathlib import Path

_HMAC_PREFIX = "obol_hmac:"
_GCM_PREFIX = "obol_gcm:"


def _get_or_create_archive_key() -> str:
    """Get or create a persistent archive key for local log encryption.

    Returns the base64url-encoded archive key string, or empty string on failure.
    The key is stored at ~/.obol/archive_key with mode 0600.
    """
    archive_key_path = Path(os.path.expanduser("~/.obol/archive_key"))
    try:
        if archive_key_path.exists():
            return archive_key_path.read_text().strip()
    except OSError:
        return ""

    # Generate new 32-byte random key, base64url-encode it
    try:
        raw_key = os.urandom(32)
        encoded_key = base64.urlsafe_b64encode(raw_key).decode()
        archive_key_path.parent.mkdir(parents=True, exist_ok=True)
        # Write with restrictive permissions
        fd = os.open(str(archive_key_path), os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
        with os.fdopen(fd, "w") as f:
            f.write(encoded_key)
        return encoded_key
    except OSError:
        return ""


def _derive_archive_gcm_key(archive_key_str: str) -> bytes:
    """Derive AES-256-GCM key from the persistent archive key via HKDF."""
    from cryptography.hazmat.primitives.hashes import SHA256
    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
    raw_key = base64.urlsafe_b64decode(archive_key_str.encode())
    return HKDF(
        algorithm=SHA256(), length=32,
        salt=b"obol-archive-aes256-v1", info=b"local-log-encryption",
    ).derive(raw_key)


def _derive_apikey_gcm_key(api_key: str) -> bytes:
    """Derive AES-256-GCM key from the API key via HKDF."""
    from cryptography.hazmat.primitives.hashes import SHA256
    from cryptography.hazmat.primitives.kdf.hkdf import HKDF
    return HKDF(
        algorithm=SHA256(), length=32,
        salt=b"obol-apikey-aes256-v1", info=b"local-log-encryption",
    ).derive(api_key.encode())


def _encrypt_line(line: str, api_key: str, archive_key: str = "") -> str:
    """Encrypt a JSONL line. Prefers archive key, falls back to api_key.

    Primary: AES-256-GCM (post-quantum safe). Fallback: HMAC+base64 (no cryptography lib).
    """
    if not api_key and not archive_key:
        return line
    try:
        from cryptography.hazmat.primitives.ciphers.aead import AESGCM
        if archive_key:
            key = _derive_archive_gcm_key(archive_key)
        else:
            key = _derive_apikey_gcm_key(api_key)
        aesgcm = AESGCM(key)
        nonce = os.urandom(12)
        ct = aesgcm.encrypt(nonce, line.encode(), None)
        encoded = base64.urlsafe_b64encode(nonce + ct).decode()
        return _GCM_PREFIX + encoded
    except ImportError:
        # Fallback: HMAC-signed base64 (integrity + obfuscation, not true encryption)
        if archive_key:
            raw_key = base64.urlsafe_b64decode(archive_key.encode())
            key_bytes = hashlib.sha256(raw_key).digest()
        else:
            key_bytes = hashlib.sha256(api_key.encode()).digest()
        encoded = base64.b64encode(line.encode()).decode()
        sig = hmac_mod.new(key_bytes, line.encode(), hashlib.sha256).hexdigest()[:16]
        return _HMAC_PREFIX + sig + ":" + encoded


def _decrypt_line(line: str, api_key: str, archive_key: str = "") -> str:
    """Decrypt a JSONL line. Handles AES-256-GCM, HMAC, and legacy plaintext."""
    if not line:
        return line
    if not api_key and not archive_key:
        return line
    # AES-256-GCM encrypted lines (current format)
    if line.startswith(_GCM_PREFIX):
        try:
            from cryptography.hazmat.primitives.ciphers.aead import AESGCM
            raw = base64.urlsafe_b64decode(line[len(_GCM_PREFIX):])
            nonce = raw[:12]
            ct_with_tag = raw[12:]
            # Try archive key first, then api_key
            if archive_key:
                try:
                    key = _derive_archive_gcm_key(archive_key)
                    return AESGCM(key).decrypt(nonce, ct_with_tag, None).decode()
                except Exception:
                    pass
            if api_key:
                try:
                    key = _derive_apikey_gcm_key(api_key)
                    return AESGCM(key).decrypt(nonce, ct_with_tag, None).decode()
                except Exception:
                    pass
        except (ImportError, Exception):
            pass
        return line  # Can't decrypt
    if line.startswith(_HMAC_PREFIX):
        try:
            rest = line[len(_HMAC_PREFIX):]
            sig, encoded = rest.split(":", 1)
            content = base64.b64decode(encoded).decode()
            # Try archive key HMAC first, then api_key HMAC
            if archive_key:
                raw_key = base64.urlsafe_b64decode(archive_key.encode())
                key_bytes = hashlib.sha256(raw_key).digest()
                expected_sig = hmac_mod.new(key_bytes, content.encode(), hashlib.sha256).hexdigest()[:16]
                if hmac_mod.compare_digest(sig, expected_sig):
                    return content
            if api_key:
                key_bytes = hashlib.sha256(api_key.encode()).digest()
                expected_sig = hmac_mod.new(key_bytes, content.encode(), hashlib.sha256).hexdigest()[:16]
                if hmac_mod.compare_digest(sig, expected_sig):
                    return content
            return line  # Tampered or wrong key, return raw
        except Exception:
            return line
    # Legacy plaintext
    return line


class LocalAuditLog:
    """Append-only local audit log for complete event history."""

    def __init__(self, agent_id: int, session_id: str, log_dir: str = None, api_key: str = None):
        self.agent_id = agent_id
        self.session_id = session_id
        self._api_key = api_key
        self._archive_key = _get_or_create_archive_key()

        # Default: ~/.obol/audit/
        self.log_dir = Path(log_dir or os.path.expanduser("~/.obol/audit"))
        self.log_dir.mkdir(parents=True, exist_ok=True)
        # Restrict directory permissions to owner only
        try:
            os.chmod(self.log_dir, 0o700)
        except OSError:
            pass

        # One file per session
        self.log_file = self.log_dir / f"{agent_id}_{session_id}.jsonl"
        self._file = None
        self._lock = threading.Lock()

    def start(self):
        """Open log file for writing."""
        # Use restrictive permissions (0o600) to prevent other users from reading
        fd = os.open(self.log_file, os.O_WRONLY | os.O_CREAT | os.O_APPEND, 0o600)
        self._file = os.fdopen(fd, "a", encoding="utf-8")
        # Write session header
        self._write_entry({
            "type": "session_start",
            "agent_id": self.agent_id,
            "session_id": self.session_id,
        })

    def stop(self):
        """Close log file."""
        if self._file:
            self._write_entry({
                "type": "session_end",
                "agent_id": self.agent_id,
                "session_id": self.session_id,
            })
            self._file.close()
            self._file = None

    def write(self, event: dict):
        """Write event to local log (append-only)."""
        entry = {
            "type": "event",
            **event,
        }
        self._write_entry(entry)

    def _write_entry(self, entry: dict):
        """Write a single JSONL entry with timestamp, encrypted if key available."""
        if not self._file:
            return

        entry["timestamp"] = datetime.now(timezone.utc).isoformat()

        with self._lock:
            try:
                line = json.dumps(entry)
                if self._archive_key or self._api_key:
                    line = _encrypt_line(line, self._api_key, self._archive_key)
                self._file.write(line + "\n")
                self._file.flush()
            except Exception:
                pass  # Never let logging break the agent
