"""
NovyxIntegrityMemory - Memory with cryptographic audit trail.

This module provides memory classes that use Novyx Core's built-in audit trail
for verified, tamper-proof memory operations. Every operation is logged with
SHA-256 hashing and can be verified for integrity.

This makes Novyx the ONLY LangChain memory provider with cryptographic audit trails.

In v2.0, the separate Integrity API has been merged into Core. All audit features
are now built-in to Novyx Core and available on all tiers (with retention limits).

Example:
    from novyx_langchain import NovyxIntegrityMemory

    # Audit trail is now built-in to Core
    memory = NovyxIntegrityMemory(
        api_key="nram_tenant_xxx",
        session_id="user-123"
    )

    # Use with LangChain - all writes are audited
    chain = ConversationChain(llm=llm, memory=memory)

    # Verify audit integrity
    verification = memory.verify_audit()
    print(f"Audit valid: {verification['valid']}")
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

from pydantic import Field

from novyx_langchain.memory import NovyxMemory

# LangChain imports
try:
    from langchain_core.messages import BaseMessage
except ImportError:
    from langchain.schema import BaseMessage

logger = logging.getLogger(__name__)


# =============================================================================
# Integrity-Protected Memory Class
# =============================================================================

class NovyxIntegrityMemory(NovyxMemory):
    """
    Full-featured memory with cryptographic audit trail.

    This is the ONLY LangChain memory provider with tamper-proof audit trails.
    Every memory operation is logged with SHA-256 hashing for verification.

    In v2.0, audit features are built directly into Novyx Core:
    - SHA-256 content hashing
    - Immutable audit trail
    - Tamper detection
    - Forensic rollback capability (Pro+)
    - Audit export (Pro+)

    Example:
        from novyx_langchain import NovyxIntegrityMemory
        from langchain.chains import ConversationChain

        memory = NovyxIntegrityMemory(
            api_key="nram_tenant_xxx",
            session_id="user-123"
        )

        chain = ConversationChain(llm=llm, memory=memory)

        # All writes are now cryptographically audited!
        response = chain.invoke({"input": "Remember my name is Alice"})

        # Verify audit integrity
        verification = memory.verify_audit()
        if verification['valid']:
            print("✅ Audit trail integrity verified!")

    Security Features:
        - SHA-256 content hashing on all writes
        - Immutable audit trail with chain verification
        - Tamper detection
        - Rollback capability (Pro+)
        - Cryptographic verification
    """

    # Additional integrity settings
    verify_on_read: bool = Field(False, description="Verify audit on every read")
    audit_checks: bool = Field(True, description="Enable automatic audit checks")

    def __init__(self, **kwargs):
        """
        Initialize integrity-protected memory.

        In v2.0, audit is built into Core - no separate integrity credentials needed.

        Args:
            api_key: Novyx API key
            session_id: Session identifier
            verify_on_read: Verify audit integrity on reads (default: False)
            audit_checks: Enable automatic audit verification (default: True)
            **kwargs: Additional NovyxMemory arguments
        """
        super().__init__(**kwargs)

        logger.info(
            f"NovyxIntegrityMemory initialized with audit trail. "
            f"Verify on read: {self.verify_on_read}, Audit checks: {self.audit_checks}"
        )

        # Perform initial audit check if enabled
        if self.audit_checks:
            try:
                status = self.verify_audit()
                if not status.get("valid", True):
                    logger.warning(
                        f"Audit integrity warning: {status.get('errors', [])}"
                    )
            except Exception as e:
                logger.warning(f"Could not verify audit on init: {e}")

    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Load memory with optional audit verification."""
        result = super().load_memory_variables(inputs)

        if self.verify_on_read and self.audit_checks:
            # Verify audit integrity before returning
            try:
                status = self.verify_audit()
                if not status.get("valid", True):
                    logger.warning(
                        f"Audit integrity warning on read: "
                        f"{status.get('errors', ['Unknown'])}"
                    )
            except Exception as e:
                logger.warning(f"Could not verify audit on read: {e}")

        return result

    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """
        Save context with audit trail.

        Every write is automatically logged in the immutable audit trail with SHA-256 hashing.
        """
        super().save_context(inputs, outputs)

        # Audit is automatic in v2 - no additional action needed
        logger.debug("Memory saved with automatic audit trail")

    # =========================================================================
    # Audit & Verification Methods
    # =========================================================================

    def verify_audit(self) -> Dict[str, Any]:
        """
        Verify audit trail integrity.

        Checks the cryptographic chain of audit entries to detect tampering.

        Returns:
            Dict with verification result:
                - valid: bool (True if no tampering detected)
                - errors: List[str] (any integrity issues found)
                - total_entries: int
                - chain_head: str (latest entry hash)

        Example:
            >>> memory = NovyxIntegrityMemory(api_key="nram_xxx", session_id="user-123")
            >>> verification = memory.verify_audit()
            >>> if verification['valid']:
            >>>     print(f"✅ Verified {verification['total_entries']} audit entries")
            >>> else:
            >>>     print(f"⚠️ Integrity issues: {verification['errors']}")
        """
        try:
            return self._client._client.audit_verify()
        except Exception as e:
            logger.error(f"Audit verification failed: {e}")
            return {
                "valid": False,
                "errors": [str(e)],
                "total_entries": 0,
                "chain_head": None,
            }

    def get_audit_summary(self, limit: int = 100) -> Dict[str, Any]:
        """
        Get summary of audit trail operations.

        Args:
            limit: Maximum audit entries to analyze

        Returns:
            Dict with audit summary:
                - total_operations: int
                - by_type: Dict[str, int] (counts per operation type)
                - first_operation: str (ISO timestamp)
                - last_operation: str (ISO timestamp)
                - integrity_valid: bool

        Example:
            >>> memory = NovyxIntegrityMemory(api_key="nram_xxx", session_id="user-123")
            >>> summary = memory.get_audit_summary()
            >>> print(f"Total operations: {summary['total_operations']}")
            >>> print(f"Creates: {summary['by_type']['CREATE']}")
            >>> print(f"Deletes: {summary['by_type']['DELETE']}")
        """
        audit_entries = self.audit(limit=limit)

        by_type: Dict[str, int] = {}
        timestamps = []

        for entry in audit_entries:
            op = entry.get("operation", "UNKNOWN")
            by_type[op] = by_type.get(op, 0) + 1

            ts = entry.get("timestamp")
            if ts:
                timestamps.append(ts)

        # Sort timestamps
        timestamps.sort()

        # Verify integrity
        verification = self.verify_audit()

        return {
            "total_operations": len(audit_entries),
            "by_type": by_type,
            "first_operation": timestamps[0] if timestamps else None,
            "last_operation": timestamps[-1] if timestamps else None,
            "integrity_valid": verification.get("valid", False),
        }

    def export_audit(self, format: str = "csv", output_file: Optional[str] = None) -> Optional[bytes]:
        """
        Export audit trail (Pro+ only).

        Args:
            format: Export format ("csv", "json", or "jsonl")
            output_file: Optional file path to write export

        Returns:
            Raw bytes if output_file is None, otherwise None (writes to file)

        Raises:
            NovyxForbiddenError: If not on Pro+ tier

        Example:
            >>> memory = NovyxIntegrityMemory(api_key="nram_xxx", session_id="user-123")
            >>> # Export to file
            >>> memory.export_audit(format="csv", output_file="audit.csv")
            >>> # Get raw bytes
            >>> data = memory.export_audit(format="json")
        """
        data = self.audit_export(format=format)

        if output_file:
            with open(output_file, "wb") as f:
                f.write(data)
            logger.info(f"Audit exported to {output_file}")
            return None

        return data

    def rollback_with_verification(self, target: str, dry_run: bool = False) -> Dict[str, Any]:
        """
        Rollback with automatic audit verification (Pro+ only).

        Performs a rollback and then verifies audit integrity.

        Args:
            target: Timestamp or relative time to rollback to
            dry_run: Preview only

        Returns:
            Dict with rollback result and verification status

        Example:
            >>> memory = NovyxIntegrityMemory(api_key="nram_xxx", session_id="user-123")
            >>> result = memory.rollback_with_verification("2 hours ago", dry_run=True)
            >>> if result['safe'] and result['integrity_valid']:
            >>>     result = memory.rollback_with_verification("2 hours ago")
        """
        # Execute rollback
        result = self.rollback(target=target, dry_run=dry_run)

        # Verify integrity after rollback
        if not dry_run:
            verification = self.verify_audit()
            result["integrity_valid"] = verification.get("valid", False)
            result["integrity_errors"] = verification.get("errors", [])

        return result


# =============================================================================
# Factory Functions
# =============================================================================

def create_integrity_memory(
    api_key: str,
    session_id: str,
    **kwargs,
) -> NovyxIntegrityMemory:
    """
    Factory function to create integrity-protected memory.

    In v2.0, audit features are built into Core - no separate credentials needed.

    Args:
        api_key: Novyx API key
        session_id: Session identifier
        **kwargs: Additional arguments for NovyxIntegrityMemory

    Returns:
        Configured NovyxIntegrityMemory instance

    Example:
        memory = create_integrity_memory(
            api_key="nram_xxx",
            session_id="user-123",
            k=10,
            return_messages=True,
            verify_on_read=True,
        )
    """
    return NovyxIntegrityMemory(
        api_key=api_key,
        session_id=session_id,
        **kwargs,
    )
