"""
Encrypted wrapper around EpisodeStore for SEC-01 compliance.

This module provides transparent encryption for episode memory content
while preserving metadata (tags, embedding) for querying.

Requirements: SEC-01 (encryption for agent memories)
"""

from __future__ import annotations

import json
import base64
from typing import Optional, List

from gsd_rlm.memory.hmem import Episode, EpisodeType, EpisodeStore
from .encryption import MemoryEncryptor, EncryptedData


class EncryptedEpisodeStore:
    """
    EpisodeStore wrapper with AES-256-GCM encryption (SEC-01).

    Transparently encrypts episode memory content (context, action, outcome)
    while preserving metadata (tags, embedding) for querying.

    This allows:
    - Storing sensitive agent memories in encrypted form
    - Querying by tags without decryption
    - Semantic search via preserved embeddings
    - Transparent retrieval with automatic decryption

    Example:
        >>> store = EncryptedEpisodeStore("memory/hmem.db")
        >>> episode = Episode(
        ...     episode_id="e1",
        ...     agent_id="agent1",
        ...     session_id="s1",
        ...     episode_type=EpisodeType.TASK_EXECUTION,
        ...     context="sensitive data",
        ... )
        >>> store.store(episode)  # Encrypts and stores
        >>> retrieved = store.get("e1")  # Decrypts and returns
        >>> assert retrieved.context == "sensitive data"
    """

    ENCRYPTED_TAG = "encrypted"

    def __init__(
        self,
        db_path: str = "memory/hmem.db",
        encryptor: Optional[MemoryEncryptor] = None,
    ):
        """
        Initialize the encrypted episode store.

        Args:
            db_path: Path to SQLite database file.
            encryptor: MemoryEncryptor instance. If None, creates a new one.
        """
        self._store = EpisodeStore(db_path)
        self._encryptor = encryptor or MemoryEncryptor()

    def store(self, episode: Episode) -> None:
        """
        Store episode with encrypted memory content.

        The context, action, and outcome fields are encrypted and stored
        in the context field as base64. Tags and embedding are preserved
        unencrypted for querying.

        Args:
            episode: Episode to encrypt and store.
        """
        # Encrypt the sensitive memory content
        memory_content = json.dumps(
            {
                "context": episode.context,
                "action": episode.action,
                "outcome": episode.outcome,
            }
        ).encode("utf-8")

        encrypted = self._encryptor.encrypt(memory_content)

        # Create encrypted episode representation
        encrypted_episode = Episode(
            episode_id=episode.episode_id,
            agent_id=episode.agent_id,
            session_id=episode.session_id,
            episode_type=episode.episode_type,
            # Store encrypted data as base64 in context field
            context=self._encode_encrypted(encrypted),
            action="[ENCRYPTED]",
            outcome="[ENCRYPTED]",
            success=episode.success,
            timestamp=episode.timestamp,
            tokens_used=episode.tokens_used,
            duration_ms=episode.duration_ms,
            # Embedding and tags are NOT encrypted (metadata)
            embedding=episode.embedding,
            tags=episode.tags + [self.ENCRYPTED_TAG],
            trace_id=episode.trace_id,
        )

        self._store.store(encrypted_episode)

    def get(self, episode_id: str) -> Optional[Episode]:
        """
        Retrieve and decrypt episode.

        If the episode is not encrypted (no 'encrypted' tag), returns it as-is.
        This allows mixing encrypted and unencrypted episodes in the same store.

        Args:
            episode_id: Unique episode identifier.

        Returns:
            Decrypted Episode if found, None otherwise.

        Raises:
            cryptography.exceptions.InvalidTag: If ciphertext is tampered.
        """
        encrypted_episode = self._store.get(episode_id)
        if encrypted_episode is None:
            return None

        # Check if encrypted
        if self.ENCRYPTED_TAG not in encrypted_episode.tags:
            return encrypted_episode  # Not encrypted, return as-is

        # Decrypt memory content
        encrypted = self._decode_encrypted(encrypted_episode.context)
        memory_content = self._encryptor.decrypt(encrypted)
        memory_dict = json.loads(memory_content.decode("utf-8"))

        # Return decrypted episode
        return Episode(
            episode_id=encrypted_episode.episode_id,
            agent_id=encrypted_episode.agent_id,
            session_id=encrypted_episode.session_id,
            episode_type=encrypted_episode.episode_type,
            context=memory_dict["context"],
            action=memory_dict["action"],
            outcome=memory_dict["outcome"],
            success=encrypted_episode.success,
            timestamp=encrypted_episode.timestamp,
            tokens_used=encrypted_episode.tokens_used,
            duration_ms=encrypted_episode.duration_ms,
            embedding=encrypted_episode.embedding,
            tags=[t for t in encrypted_episode.tags if t != self.ENCRYPTED_TAG],
            trace_id=encrypted_episode.trace_id,
        )

    def get_episodes_for_session(self, session_id: str) -> List[Episode]:
        """
        Get all episodes for a session, decrypted.

        Args:
            session_id: Session identifier.

        Returns:
            List of decrypted episodes for the session, oldest first.
        """
        encrypted = self._store.get_episodes_for_session(session_id)
        return [self._decrypt_if_needed(ep) for ep in encrypted]

    def get_episodes_by_agent(self, agent_id: str, limit: int = 50) -> List[Episode]:
        """
        Get recent episodes for an agent, decrypted.

        Args:
            agent_id: Agent identifier.
            limit: Maximum number of episodes to return.

        Returns:
            List of decrypted episodes for the agent, most recent first.
        """
        encrypted = self._store.get_episodes_by_agent(agent_id, limit)
        return [self._decrypt_if_needed(ep) for ep in encrypted]

    def delete(self, episode_id: str) -> bool:
        """
        Delete an episode.

        Args:
            episode_id: Episode to delete.

        Returns:
            True if episode was deleted, False if not found.
        """
        return self._store.delete(episode_id)

    def count(self) -> int:
        """
        Get total episode count.

        Returns:
            Total number of episodes in the store.
        """
        return self._store.count()

    def _encode_encrypted(self, encrypted: EncryptedData) -> str:
        """
        Encode encrypted data as base64 string.

        Combines IV and ciphertext into a single base64 string
        for storage in the context field.

        Args:
            encrypted: EncryptedData to encode.

        Returns:
            Base64-encoded string containing IV + ciphertext.
        """
        return base64.b64encode(encrypted.iv + encrypted.ciphertext).decode("utf-8")

    def _decode_encrypted(self, encoded: str) -> EncryptedData:
        """
        Decode base64 string to EncryptedData.

        Splits the decoded bytes into IV (first 12 bytes) and
        ciphertext (remaining bytes).

        Args:
            encoded: Base64-encoded string from storage.

        Returns:
            EncryptedData with IV and ciphertext.
        """
        data = base64.b64decode(encoded.encode("utf-8"))
        return EncryptedData(
            iv=data[:12],  # First 12 bytes are IV
            ciphertext=data[12:],  # Rest is ciphertext + tag
        )

    def _decrypt_if_needed(self, episode: Episode) -> Episode:
        """
        Decrypt episode if it has encrypted tag.

        Used for batch retrieval methods to avoid redundant
        database lookups for already-fetched episodes.

        Args:
            episode: Episode to potentially decrypt.

        Returns:
            Decrypted episode if encrypted, original otherwise.
        """
        if self.ENCRYPTED_TAG not in episode.tags:
            return episode
        # Re-use get() logic for decryption
        return self.get(episode.episode_id) or episode
