import hashlib, json
from datetime import datetime, timezone
from uuid import uuid4
from typing import Any, Dict, List, Optional
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding, rsa
from pydantic import BaseModel, Field
 
 
class MemoryAttestation(BaseModel):
    """Signed proof of memory retrieval state — EU AI Act Art.13."""
    retrieval_id: str = Field(default_factory=lambda: str(uuid4()))
    timestamp_utc: str = Field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )
    session_id: str
    agent_id: str
    memory_chunk_hashes: List[str]
    combined_memory_hash: str
    query_context_hash: str
    signature: Optional[str] = None
 
    def compute_hash(self) -> str:
        data = (
            f"{self.retrieval_id}{self.timestamp_utc}"
            f"{self.session_id}{self.agent_id}"
            f"{self.combined_memory_hash}{self.query_context_hash}"
        )
        return hashlib.sha256(data.encode()).hexdigest()
 
    def sign(self, private_key: rsa.RSAPrivateKey) -> "MemoryAttestation":
        hash_value = self.compute_hash()
        sig_bytes = private_key.sign(
            hash_value.encode(),
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )
        self.signature = sig_bytes.hex()
        return self
 
 
class AttestationProxy:
    """
    Drop-in proxy for any memory client.
    Compatible with: Mem0, LangGraph BaseStore, Zep, custom clients.
    """
    def __init__(self, memory_client, private_key, redis_client=None):
        self.memory_client = memory_client
        self.private_key = private_key
        self.redis = redis_client
        self.local_audit_log: List[MemoryAttestation] = []
 
    def retrieve(self, session_id, agent_id, query, **kwargs):
        # 1. Pass-through retrieval
        results = self.memory_client.retrieve(query, **kwargs)
 
        # 2. Hash each chunk (content-addressed)
        chunk_hashes = [
            hashlib.sha256(
                json.dumps(c, sort_keys=True).encode()
            ).hexdigest()
            for c in results.get("results", [])
        ]
 
        # 3. Combined retrieval set hash
        combined_hash = hashlib.sha256(
            "".join(sorted(chunk_hashes)).encode()
        ).hexdigest()
 
        # 4. Query context hash
        query_hash = hashlib.sha256(
            f"{session_id}{agent_id}{query}".encode()
        ).hexdigest()
 
        # 5. Build + sign attestation
        attestation = MemoryAttestation(
            session_id=session_id, agent_id=agent_id,
            memory_chunk_hashes=chunk_hashes,
            combined_memory_hash=combined_hash,
            query_context_hash=query_hash
        ).sign(self.private_key)
 
        # 6. Store — 90-day TTL (GDPR Art.17 compatible)
        if self.redis:
            self.redis.setex(
                f"attestation:{attestation.retrieval_id}",
                86400 * 90,
                attestation.model_dump_json()
            )
        self.local_audit_log.append(attestation)
 
        # 7. Return enriched response
        return {
            "memory": results,
            "attestation_id": attestation.retrieval_id,
            "attestation_hash": attestation.compute_hash(),
            "timestamp": attestation.timestamp_utc,
            "signature": attestation.signature
        }
 
 
# USAGE — drop in place of your existing memory client
if __name__ == "__main__":
    from mem0 import Memory
    import redis as redis_lib
 
    # Production: load from KMS/HSM — never generate at runtime
    private_key = rsa.generate_private_key(
        public_exponent=65537, key_size=2048
    )
    proxy = AttestationProxy(
        memory_client=Memory(),
        private_key=private_key,
        redis_client=redis_lib.Redis(host="localhost", port=6379)
    )
    result = proxy.retrieve(
        session_id="user_123_session_456",
        agent_id="fraud_detector_v2",
        query="customer recent transaction preferences"
    )

version: '3.8'
 
services:
  postgres:
    image: pgvector/pgvector:pg16
    environment:
      POSTGRES_DB: agent_memory
      POSTGRES_USER: agent
      POSTGRES_PASSWORD: ${PG_PASSWORD:-ranksquire2026}
    volumes:
      - pgdata:/var/lib/postgresql/data
      - ./init.sql:/docker-entrypoint-initdb.d/init.sql
    ports: ["5432:5432"]
    healthcheck:
      test: ["CMD-SHELL", "pg_isready -U agent -d agent_memory"]
      interval: 10s
      timeout: 5s
      retries: 5
 
  qdrant:
    image: qdrant/qdrant:v1.10.1
    ports: ["6333:6333", "6334:6334"]
    volumes: [qdrant_storage:/qdrant/storage]
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:6333/readyz"]
      interval: 10s
      timeout: 5s
      retries: 5
 
  redis-audit:
    image: redis:7.2-alpine
    command: redis-server --appendonly yes --maxmemory 2gb
             --maxmemory-policy allkeys-lru
    ports: ["6379:6379"]
    volumes: [redis_audit:/data]
 
  attestation-proxy:
    build: {context: ./attestation-proxy, dockerfile: Dockerfile}
    environment:
      PG_CONNECTION: postgresql://agent:${PG_PASSWORD}@postgres:5432/agent_memory
      QDRANT_URL: http://qdrant:6333
      REDIS_URL: redis://redis-audit:6379
      PRIVATE_KEY_B64: ${PRIVATE_KEY_B64}   # Use KMS/HSM in production
    ports: ["8003:8003"]
    depends_on:
      postgres: {condition: service_healthy}
      qdrant: {condition: service_healthy}
 
  langfuse:
    image: langfuse/langfuse:2.84.0
    environment:
      DATABASE_URL: postgresql://agent:${PG_PASSWORD}@postgres:5432/langfuse
      NEXTAUTH_SECRET: ${NEXTAUTH_SECRET:-ranksquire-obs-secret}
      NEXTAUTH_URL: http://localhost:3000
    ports: ["3000:3000"]
    depends_on:
      postgres: {condition: service_healthy}
 
volumes:
  pgdata:
  qdrant_storage:
  redis_audit:

# BEFORE (fails with any subgraph):
from langgraph.checkpoint.postgres import PostgresSaver
memory = PostgresSaver.from_conn_string("postgresql://...")
graph = builder.compile(checkpointer=memory)  # TypeError on subgraph
 
# AFTER — Option A: Remove checkpointer (trade: no auto-recovery)
graph = builder.compile()
 
# AFTER — Option B: MemorySaver (trade: in-memory only, not persistent)
from langgraph.checkpoint import MemorySaver
graph = builder.compile(checkpointer=MemorySaver())
 
# AFTER — Option C: App-layer checkpointing (RECOMMENDED)
# Handle state persistence in your PostgreSQL agent_memory table
# See: github.com/ranksquire/memory-benchmark/patterns/manual_checkpoint.py

# BEFORE (misses voice-transcribed queries):
CACHEABLE_QUERY_TYPES = [QueryType.RETRIEVAL]
 
# AFTER (apply to voice agent paths only):
CACHEABLE_QUERY_TYPES = [QueryType.RETRIEVAL, QueryType.GENERAL]
 
# Monitoring: track cache_hit_rate by query_type in Langfuse
# Alert threshold: cache_hit_rate < 70% for QueryType.RETRIEVAL

# Cron: 0 2 * * * python prune_memory.py --threshold 0.6
# DRY RUN first: python prune_memory.py --dry-run
 
import psycopg2
 
def prune_low_value_memories(
    conn_string: str,
    confidence_threshold: float = 0.6,
    max_age_days: int = 90,
    dry_run: bool = False
) -> dict:
    with psycopg2.connect(conn_string) as conn:
        with conn.cursor() as cur:
            cur.execute("""
                SELECT COUNT(*) FROM agent_memory
                WHERE confidence < %s
                   OR (valid_to IS NULL
                       AND created_at < NOW() - INTERVAL '%s days'
                       AND importance_score < 0.7)
            """, (confidence_threshold, max_age_days))
            to_remove = cur.fetchone()[0]
 
            if not dry_run:
                # Soft delete — preserves EU AI Act audit trail
                cur.execute("""
                    UPDATE agent_memory SET valid_to = NOW()
                    WHERE confidence < %s
                       OR (valid_to IS NULL
                           AND created_at < NOW() - INTERVAL '%s days'
                           AND importance_score < 0.7)
                """, (confidence_threshold, max_age_days))
                conn.commit()
 
    return {"removed": to_remove, "dry_run": dry_run}

from sentence_transformers import SentenceTransformer
import numpy as np
 
encoder = SentenceTransformer('all-MiniLM-L6-v2')
 
def resolve_entity(
    candidate: str,
    existing_entities: list,
    threshold: float = 0.85
) -> str:
    """
    Returns matching entity if similarity > threshold, else candidate.
    resolve_entity("Alex Smith", ["Alexander Smith"]) → "Alexander Smith"
    resolve_entity("Bob Jones", ["Alex Smith"])       → "Bob Jones"
    """
    if not existing_entities:
        return candidate
    candidate_emb = encoder.encode(candidate)
    existing_embs = encoder.encode(existing_entities)
    sims = np.dot(existing_embs, candidate_emb) / (
        np.linalg.norm(existing_embs, axis=1) *
        np.linalg.norm(candidate_emb)
    )
    max_idx = np.argmax(sims)
    if sims[max_idx] > threshold:
        return existing_entities[max_idx]
    return candidate

# WRONG — shared collection, optional filter = PII exposure risk
collection.query(query_embedding=emb, where={"tenant_id": tid})
 
# CORRECT — collection-per-tenant (isolation at storage layer)
class TenantIsolatedMemory:
    def __init__(self, qdrant_client):
        self.client = qdrant_client
        self._collections = {}
 
    def _get_collection(self, tenant_id: str) -> str:
        name = f"memory_{tenant_id}"
        if name not in self._collections:
            self.client.create_collection(
                collection_name=name,
                vectors_config={"size": 1536, "distance": "Cosine"},
                optimizers_config={"default_segment_number": 2}
            )
            self._collections[name] = True
        return name
 
    def retrieve(self, tenant_id: str, query_embedding, limit=5):
        return self.client.search(
            collection_name=self._get_collection(tenant_id),
            query_vector=query_embedding,
            limit=limit
        )

async def dual_write_migration(
    managed_client,
    sovereign_client,
    session_id: str,
    query: str
):
    """
    Write to both. Read from managed (primary).
    Log every diff for reconciliation review.
    """
    # Parallel writes — neither blocks the other
    await asyncio.gather(
        managed_client.add(query, session_id=session_id),
        sovereign_client.add(query, session_id=session_id)
    )
 
    # Primary read: managed during Phase 1
    managed_result = await managed_client.retrieve(query)
 
    # Validation: sovereign must match managed
    sovereign_result = await sovereign_client.retrieve(query)
    if managed_result != sovereign_result:
        logger.warning(
            f"Diff: session={session_id} query_len={len(query)}"
        )
 
    return managed_result  # Serve managed until Phase 2 trigger