"""
NovyxTracedMemory - Memory with automatic trace audit logging.

This module provides memory classes that automatically create trace sessions
and log every memory operation as a trace step. This is the "compliance mode"
for LangChain agents that need full audit trails of all operations.

Trace audit is a Pro+ feature that provides:
- RSA-signed trace sessions
- Real-time policy enforcement
- Circuit breaker for violations
- Full audit trail of agent actions

This makes Novyx the ONLY LangChain memory provider with compliance-grade tracing.

Example:
    from novyx_langchain import NovyxTracedMemory

    # Automatically creates trace sessions (Pro+ only)
    memory = NovyxTracedMemory(
        api_key="nram_tenant_xxx",
        session_id="user-123",
        agent_id="my-agent"
    )

    # Use with LangChain - all operations are automatically traced
    chain = ConversationChain(llm=llm, memory=memory)

    # Every memory read/write is logged as a trace step
    response = chain.invoke({"input": "Hello!"})

    # Trace is auto-completed when memory is destroyed or manually
    memory.complete_trace()
"""

from __future__ import annotations

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

from pydantic import Field

from novyx_langchain.memory import NovyxMemory

# Novyx SDK imports
try:
    from novyx import NovyxForbiddenError
    NOVYX_SDK_AVAILABLE = True
except ImportError:
    NOVYX_SDK_AVAILABLE = False
    NovyxForbiddenError = None

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

logger = logging.getLogger(__name__)


class NovyxTracedMemory(NovyxMemory):
    """
    Memory with automatic trace audit logging (Pro+ only).

    This is the ONLY LangChain memory provider with compliance-grade trace audit.

    Every memory operation is automatically logged as a trace step:
    - Memory writes (save_context) → trace step
    - Memory reads (load_memory_variables) → trace step
    - Memory searches → trace step
    - Custom memories → trace step

    Traces are RSA-signed for tamper-proof verification and can enforce
    real-time policies (e.g., circuit breaker for violations).

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

        memory = NovyxTracedMemory(
            api_key="nram_tenant_xxx",
            session_id="user-123",
            agent_id="my-agent"
        )

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

        # All operations are automatically traced
        response = chain.invoke({"input": "Remember my name is Alice"})

        # Complete trace (auto-completed on cleanup)
        result = memory.complete_trace()
        print(f"Trace signature: {result['signature']}")

    Compliance Features:
        - RSA-signed trace sessions
        - Immutable audit trail
        - Real-time policy enforcement
        - Circuit breaker for violations
        - Tamper-proof verification
        - Pro+ only (graceful degradation to regular memory on Free tier)
    """

    # Trace settings
    agent_id: str = Field("langchain-agent", description="Agent identifier for tracing")
    auto_trace: bool = Field(True, description="Automatically create trace sessions")
    trace_reads: bool = Field(False, description="Trace read operations (can be verbose)")
    trace_on_free_tier: bool = Field(False, description="Allow on Free tier (will error)")

    # Internal state
    _trace_id: Optional[str] = None
    _trace_active: bool = False
    _trace_unavailable: bool = False
    _step_count: int = 0

    def __init__(self, **kwargs):
        """
        Initialize traced memory.

        Automatically creates a trace session if Pro+ tier.
        On Free tier, gracefully degrades to regular NovyxMemory unless trace_on_free_tier=True.

        Args:
            api_key: Novyx API key
            session_id: Session identifier
            agent_id: Agent identifier for tracing (default: "langchain-agent")
            auto_trace: Automatically create trace sessions (default: True)
            trace_reads: Trace read operations (default: False, can be verbose)
            trace_on_free_tier: Attempt trace on Free tier (will raise error, default: False)
            **kwargs: Additional NovyxMemory arguments
        """
        super().__init__(**kwargs)

        if self.auto_trace:
            self._create_trace()

    def _create_trace(self) -> None:
        """Create a new trace session."""
        if self._trace_active:
            return  # Already have active trace

        try:
            result = self._client._client.trace_create(
                agent_id=self.agent_id,
                session_id=self.session_id,
                metadata={
                    "source": "novyx-langchain",
                    "memory_class": "NovyxTracedMemory",
                },
            )

            self._trace_id = result.get("trace_id")
            self._trace_active = True
            self._step_count = 0

            logger.info(
                f"Trace session created: {self._trace_id} for agent {self.agent_id}"
            )

        except NovyxForbiddenError:
            # Not on Pro+ tier
            if self.trace_on_free_tier:
                logger.error("Trace audit requires Pro+ tier")
                raise
            else:
                logger.warning(
                    "Trace audit unavailable (Pro+ only). "
                    "Falling back to regular memory without tracing."
                )
                self._trace_unavailable = True
                self._trace_active = False

        except Exception as e:
            logger.error(f"Failed to create trace session: {e}")
            self._trace_unavailable = True
            self._trace_active = False

    def _add_trace_step(
        self,
        step_type: str,
        name: str,
        content: Optional[str] = None,
        attributes: Optional[Dict] = None,
    ) -> None:
        """Add a step to the active trace."""
        if not self._trace_active or self._trace_unavailable:
            return

        try:
            result = self._client._client.trace_step(
                trace_id=self._trace_id,
                step_type=step_type,
                content=content or name,
                metadata=attributes,
            )

            self._step_count += 1

            # Check for circuit breaker
            if result.get("circuit_breaker_triggered"):
                logger.critical(
                    f"⚠️ CIRCUIT BREAKER TRIGGERED: {result.get('reason', 'Policy violation')}"
                )
                logger.critical(f"Action: {result.get('action', 'block')}")

        except Exception as e:
            logger.error(f"Failed to add trace step: {e}")

    def load_memory_variables(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """Load memory with automatic trace step."""
        # Trace read operation if enabled
        if self.trace_reads:
            input_text = inputs.get(self.input_key, "")
            self._add_trace_step(
                step_type="observation",
                name="load_memory",
                content=f"Loading memory for input: {input_text[:100]}",
                attributes={
                    "operation": "load",
                    "session_id": self.session_id,
                },
            )

        result = super().load_memory_variables(inputs)

        if self.trace_reads:
            memory_count = len(result.get(self.memory_key, []))
            self._add_trace_step(
                step_type="observation",
                name="memory_loaded",
                content=f"Loaded {memory_count} memory items",
                attributes={
                    "memory_count": memory_count,
                },
            )

        return result

    def save_context(self, inputs: Dict[str, Any], outputs: Dict[str, str]) -> None:
        """Save context with automatic trace step."""
        input_text = inputs.get(self.input_key, "")
        output_text = outputs.get(self.output_key, "")

        # Trace the memory write
        self._add_trace_step(
            step_type="action",
            name="save_memory",
            content=f"Saving conversation turn",
            attributes={
                "operation": "save",
                "session_id": self.session_id,
                "input_length": len(input_text),
                "output_length": len(output_text),
            },
        )

        # Perform save
        super().save_context(inputs, outputs)

        # Trace completion
        self._add_trace_step(
            step_type="observation",
            name="memory_saved",
            content="Conversation turn saved successfully",
            attributes={
                "success": True,
            },
        )

    def store_memory(
        self,
        content: str,
        tags: Optional[List[str]] = None,
        importance: int = 5,
    ) -> str:
        """Store custom memory with automatic trace step."""
        # Trace the memory write
        self._add_trace_step(
            step_type="action",
            name="store_custom_memory",
            content=f"Storing custom memory: {content[:100]}",
            attributes={
                "operation": "store",
                "tags": tags or [],
                "importance": importance,
            },
        )

        # Perform store
        memory_id = super().store_memory(content, tags, importance)

        # Trace completion
        self._add_trace_step(
            step_type="observation",
            name="memory_stored",
            content=f"Custom memory stored with ID: {memory_id}",
            attributes={
                "memory_id": memory_id,
                "success": True,
            },
        )

        return memory_id

    def search(self, query: str, limit: int = 5) -> List:
        """Search memories with automatic trace step."""
        # Trace the search
        self._add_trace_step(
            step_type="action",
            name="search_memory",
            content=f"Searching for: {query}",
            attributes={
                "operation": "search",
                "query": query,
                "limit": limit,
            },
        )

        # Perform search
        results = super().search(query, limit)

        # Trace results
        self._add_trace_step(
            step_type="observation",
            name="search_results",
            content=f"Found {len(results)} matching memories",
            attributes={
                "result_count": len(results),
            },
        )

        return results

    def clear(self) -> None:
        """Clear memory with automatic trace step."""
        # Trace the clear
        self._add_trace_step(
            step_type="action",
            name="clear_memory",
            content=f"Clearing all memory for session {self.session_id}",
            attributes={
                "operation": "clear",
                "session_id": self.session_id,
            },
        )

        # Perform clear
        super().clear()

        # Trace completion
        self._add_trace_step(
            step_type="observation",
            name="memory_cleared",
            content="Memory cleared successfully",
            attributes={
                "success": True,
            },
        )

    # =========================================================================
    # Trace Management Methods
    # =========================================================================

    def complete_trace(self) -> Dict[str, Any]:
        """
        Complete the trace session with RSA signature.

        This finalizes the trace and generates a cryptographic signature
        for tamper-proof verification.

        Returns:
            Dict with trace completion result:
                - trace_id: str
                - status: str ("completed")
                - signature: str (RSA signature)
                - public_key_id: str
                - signed_at: str (ISO timestamp)
                - integrity_chain_valid: bool
                - total_steps: int

        Example:
            >>> memory = NovyxTracedMemory(api_key="nram_xxx", session_id="user-123")
            >>> # ... use memory ...
            >>> result = memory.complete_trace()
            >>> print(f"Trace signed with {result['total_steps']} steps")
            >>> print(f"Signature: {result['signature'][:32]}...")
        """
        if not self._trace_active:
            logger.warning("No active trace to complete")
            return {
                "trace_id": None,
                "status": "no_trace",
                "signature": None,
                "total_steps": 0,
            }

        try:
            result = self._client._client.trace_complete(trace_id=self._trace_id)

            self._trace_active = False
            logger.info(
                f"Trace completed: {self._trace_id} with {self._step_count} steps"
            )

            return result

        except Exception as e:
            logger.error(f"Failed to complete trace: {e}")
            raise

    def verify_trace(self) -> Dict[str, Any]:
        """
        Verify trace integrity.

        Checks the cryptographic chain of trace steps to detect tampering.

        Returns:
            Dict with verification result:
                - trace_id: str
                - valid: bool
                - errors: List[str]
                - steps_verified: int

        Example:
            >>> memory = NovyxTracedMemory(api_key="nram_xxx", session_id="user-123")
            >>> verification = memory.verify_trace()
            >>> if verification['valid']:
            >>>     print(f"✅ Verified {verification['steps_verified']} trace steps")
        """
        if not self._trace_id:
            return {
                "trace_id": None,
                "valid": False,
                "errors": ["No trace ID available"],
                "steps_verified": 0,
            }

        try:
            return self._client._client.trace_verify(trace_id=self._trace_id)
        except Exception as e:
            logger.error(f"Trace verification failed: {e}")
            return {
                "trace_id": self._trace_id,
                "valid": False,
                "errors": [str(e)],
                "steps_verified": 0,
            }

    def get_trace_status(self) -> Dict[str, Any]:
        """
        Get current trace status.

        Returns:
            Dict with trace status:
                - trace_id: Optional[str]
                - active: bool
                - step_count: int
                - agent_id: str
                - session_id: str
                - unavailable: bool (if trace not available on current tier)

        Example:
            >>> memory = NovyxTracedMemory(api_key="nram_xxx", session_id="user-123")
            >>> status = memory.get_trace_status()
            >>> print(f"Trace active: {status['active']}")
            >>> print(f"Steps logged: {status['step_count']}")
        """
        return {
            "trace_id": self._trace_id,
            "active": self._trace_active,
            "step_count": self._step_count,
            "agent_id": self.agent_id,
            "session_id": self.session_id,
            "unavailable": self._trace_unavailable,
        }

    def __del__(self):
        """Auto-complete trace on cleanup."""
        try:
            if self._trace_active:
                logger.info("Auto-completing trace on cleanup")
                self.complete_trace()
        except Exception:
            pass  # Ignore errors during cleanup


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

def create_traced_memory(
    api_key: str,
    session_id: str,
    agent_id: str = "langchain-agent",
    **kwargs,
) -> NovyxTracedMemory:
    """
    Factory function to create traced memory.

    Args:
        api_key: Novyx API key (must be Pro+ tier)
        session_id: Session identifier
        agent_id: Agent identifier for tracing
        **kwargs: Additional arguments for NovyxTracedMemory

    Returns:
        Configured NovyxTracedMemory instance

    Example:
        memory = create_traced_memory(
            api_key="nram_xxx",
            session_id="user-123",
            agent_id="my-agent",
            k=10,
            return_messages=True,
            trace_reads=True,
        )
    """
    return NovyxTracedMemory(
        api_key=api_key,
        session_id=session_id,
        agent_id=agent_id,
        **kwargs,
    )
