"""Episodic Memory (Tier 2) — SQLite-backed with relationship tables.

Adapted from PostgreSQL + Memgraph to use SQLite via soagila.db.sqlite_init.
Graph relationships stored in memory_relationships and memory_topics tables.
"""

import json
import hashlib
import logging
from datetime import datetime, timezone
from uuid import uuid4

from soagila.db.sqlite_init import get_db
from soagila.db.event_bus import publish_event
from soagila.core.tcmg.memory_models import EpisodicMemory
from soagila.core.tcmg.embedding_service import generate_embedding, cosine_similarity

logger = logging.getLogger("soagila.tcmg.episodic")

CHANNEL = "soagila:metrics"
SIMILARITY_THRESHOLD = 0.7  # For RELATED_TO relationship creation

# Canonical column list for SELECTs that feed ``_row_to_episodic``.
# Changing this order requires updating the row index in _row_to_episodic.
_EPISODIC_COLUMNS = (
    "id, user_id, org_id, encrypted_content, encryption_key_ref, embedding, "
    "pii_tags, classification, topic_tags, confidence_score, relevance_score, "
    "access_count, last_accessed, decay_rate, stability, is_canary, is_decayed, "
    "source_type, created_at, session_id"
)


def _decrypt_content(encrypted_content, key_ref: str | None) -> str:
    """Decrypt an episodic-memory blob using the referenced AES key.

    Falls back gracefully for three corner cases:
        1. Legacy rows with ``encryption_key_ref == 'mvp-plaintext'`` (pre-0.5.0)
        2. Rows whose key has been destroyed (annihilated) — returns a
           placeholder instead of raising so decay/recall don't crash
        3. Keys that aren't in the in-memory store yet (engine not loaded)
    """
    if encrypted_content is None:
        return ""
    raw = bytes(encrypted_content) if isinstance(encrypted_content, (bytes, memoryview)) else \
        str(encrypted_content).encode("utf-8")

    # Legacy rows: plaintext stored directly in the blob
    if key_ref in (None, "", "mvp-plaintext"):
        try:
            return raw.decode("utf-8")
        except UnicodeDecodeError:
            return ""

    try:
        from soagila.core.encryption.simulated_he import get_encryption_engine
        engine = get_encryption_engine()
        if not engine.key_exists(key_ref):
            # Key gone → memory is cryptographically shredded
            return "[encrypted: key destroyed]"
        return engine.decrypt(raw, key_ref)
    except Exception as exc:  # noqa: BLE001
        logger.warning(f"Decrypt failed for key_ref={key_ref}: {exc}")
        # Last resort: treat as legacy plaintext if it decodes cleanly
        try:
            return raw.decode("utf-8")
        except UnicodeDecodeError:
            return "[decrypt failed]"


def _row_to_episodic(row) -> EpisodicMemory:
    """Convert a SQLite row to an EpisodicMemory instance.

    Expected column order:
        id, user_id, org_id, encrypted_content, encryption_key_ref,
        embedding, pii_tags, classification, topic_tags, confidence_score,
        relevance_score, access_count, last_accessed, decay_rate, stability,
        is_canary, is_decayed, source_type, created_at, session_id
    """
    key_ref = row[4] if len(row) > 4 else None
    content = _decrypt_content(row[3], key_ref)
    embedding = json.loads(row[5]) if row[5] else None
    pii_tags = json.loads(row[6]) if row[6] else []
    topic_tags = json.loads(row[8]) if row[8] else []

    return EpisodicMemory(
        id=str(row[0]),
        user_id=str(row[1]),
        org_id=str(row[2]),
        content=content,
        embedding=embedding,
        pii_tags=pii_tags,
        classification=row[7] or "general",
        topic_tags=topic_tags,
        confidence_score=row[9] or 1.0,
        relevance_score=row[10] or 1.0,
        access_count=row[11] or 0,
        last_accessed=row[12],
        decay_rate=row[13] or 0.1,
        stability=row[14] or 1.0,
        is_canary=bool(row[15]),
        is_decayed=bool(row[16]),
        source_type=row[17] or "chat",
        created_at=row[18],
        session_id=row[19] or "",
    )


# =============================================
# STORE
# =============================================

async def store(memory: EpisodicMemory) -> str:
    """
    Store an episodic memory:
    1. Generate embedding (from plaintext — then plaintext is dropped)
    2. Encrypt content with AES-256-GCM
    3. Store only ciphertext + key reference in SQLite
    4. Create topic entries and relationships
    5. Create similarity relationships
    6. Publish event
    """
    # 1. Generate embedding over the plaintext before encrypting. The
    #    embedding itself leaks semantic information but is much harder to
    #    invert than plaintext; trading that off keeps similarity search
    #    working. This mirrors how Pinecone/Weaviate handle vectors.
    memory.embedding = generate_embedding(memory.content)

    # Content hash for dedup (over plaintext)
    content_hash = hashlib.sha256(memory.content.encode()).hexdigest()

    # 2. Encrypt content with a freshly-generated per-memory key so that
    #    destroying a single key annihilates exactly one memory. The key
    #    itself is persisted in SQLite, encrypted with the project master
    #    key (PBKDF2 from machine-id + project_hash + salt).
    from soagila.core.encryption.simulated_he import get_encryption_engine

    engine = get_encryption_engine()
    key_ref = engine.generate_key("memory")
    encrypted_content = engine.encrypt(memory.content, key_ref)
    encryption_key_ref = key_ref

    db = await get_db()

    now = datetime.now(timezone.utc).isoformat()

    await db.execute(
        """
        INSERT INTO episodic_memories (
            id, user_id, org_id, encrypted_content, encryption_key_ref,
            embedding, content_hash, pii_tags, classification, topic_tags,
            confidence_score, relevance_score, access_count, decay_rate,
            stability, is_canary, is_decayed, source_type, session_id,
            created_at, updated_at
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        (
            memory.id,
            memory.user_id,
            memory.org_id,
            encrypted_content,
            encryption_key_ref,
            json.dumps(memory.embedding),
            content_hash,
            json.dumps(memory.pii_tags),
            memory.classification,
            json.dumps(memory.topic_tags),
            memory.confidence_score,
            memory.relevance_score,
            memory.access_count,
            memory.decay_rate,
            memory.stability,
            int(memory.is_canary),
            int(memory.is_decayed),
            memory.source_type,
            memory.session_id,
            memory.created_at.isoformat() if isinstance(memory.created_at, datetime) else str(memory.created_at),
            now,
        ),
    )

    # 3. Create topic entries and memory-topic links
    for tag in memory.topic_tags:
        await db.execute("INSERT OR IGNORE INTO topics (name) VALUES (?)", (tag,))
        await db.execute(
            "INSERT OR IGNORE INTO memory_topics (memory_id, topic_name) VALUES (?, ?)",
            (memory.id, tag),
        )

    # 3b. Create session relationships (FOLLOWED_BY, SAME_SESSION)
    cursor = await db.execute(
        """
        SELECT id FROM episodic_memories
        WHERE session_id = ? AND id != ?
        ORDER BY created_at DESC LIMIT 1
        """,
        (memory.session_id, memory.id),
    )
    prev_row = await cursor.fetchone()
    if prev_row:
        prev_id = str(prev_row[0])
        await db.execute(
            "INSERT OR IGNORE INTO memory_relationships (id, source_id, target_id, relationship_type) VALUES (?, ?, ?, ?)",
            (str(uuid4()), prev_id, memory.id, "FOLLOWED_BY"),
        )
        await db.execute(
            "INSERT OR IGNORE INTO memory_relationships (id, source_id, target_id, relationship_type) VALUES (?, ?, ?, ?)",
            (str(uuid4()), prev_id, memory.id, "SAME_SESSION"),
        )

    # 3c. Create SAME_TOPIC relationships
    if memory.topic_tags:
        placeholders = ",".join("?" for _ in memory.topic_tags)
        cursor = await db.execute(
            f"""
            SELECT DISTINCT mt.memory_id FROM memory_topics mt
            JOIN episodic_memories em ON em.id = mt.memory_id
            WHERE mt.topic_name IN ({placeholders})
              AND mt.memory_id != ?
              AND em.user_id = ?
            """,
            (*memory.topic_tags, memory.id, memory.user_id),
        )
        same_topic_rows = await cursor.fetchall()
        for r in same_topic_rows:
            other_id = str(r[0])
            await db.execute(
                "INSERT OR IGNORE INTO memory_relationships (id, source_id, target_id, relationship_type) VALUES (?, ?, ?, ?)",
                (str(uuid4()), memory.id, other_id, "SAME_TOPIC"),
            )

    await db.commit()

    # 4. Create similarity relationships
    await _create_similarity_relationships(memory)

    # 5. Publish event
    event = json.dumps({
        "type": "memory_stored",
        "tier": "episodic",
        "memory_id": memory.id,
        "user_id": memory.user_id,
        "session_id": memory.session_id,
        "topic": memory.topic_tags[0] if memory.topic_tags else "general",
        "classification": memory.classification,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    })
    await publish_event(CHANNEL, event)

    logger.debug(f"Stored episodic memory {memory.id} for user {memory.user_id}")
    return memory.id


async def _create_similarity_relationships(memory: EpisodicMemory):
    """Find similar memories by embedding and create RELATED_TO edges in SQLite."""
    if not memory.embedding:
        return

    db = await get_db()
    cursor = await db.execute(
        """
        SELECT id, embedding FROM episodic_memories
        WHERE user_id = ? AND id != ? AND is_decayed = 0
          AND embedding IS NOT NULL
        ORDER BY created_at DESC LIMIT 50
        """,
        (memory.user_id, memory.id),
    )
    rows = await cursor.fetchall()

    similar_ids = []
    for row in rows:
        other_id = str(row[0])
        other_embedding = json.loads(row[1]) if isinstance(row[1], str) else row[1]
        if other_embedding:
            sim = cosine_similarity(memory.embedding, other_embedding)
            if sim >= SIMILARITY_THRESHOLD:
                similar_ids.append((other_id, sim))

    # Create RELATED_TO edges in memory_relationships
    for other_id, strength in similar_ids:
        await db.execute(
            """
            INSERT OR IGNORE INTO memory_relationships (id, source_id, target_id, relationship_type, strength)
            VALUES (?, ?, ?, ?, ?)
            """,
            (str(uuid4()), memory.id, other_id, "RELATED_TO", strength),
        )

    if similar_ids:
        await db.commit()
        logger.debug(f"Created {len(similar_ids)} RELATED_TO edges for memory {memory.id}")


# =============================================
# QUERY
# =============================================

async def query(
    user_id: str,
    query_text: str,
    limit: int = 20,
    min_relevance: float = 0.15,
) -> list[EpisodicMemory]:
    """
    Query episodic memories using embedding similarity + topic matching.
    Returns merged, deduplicated, relevance-sorted results.
    """
    query_embedding = generate_embedding(query_text)

    # SQLite: embedding similarity search
    pg_results = await _query_sqlite(user_id, query_embedding, limit, min_relevance)

    # Topic-based search (replaces Memgraph graph search)
    topic_results = await _query_by_topics(user_id, query_text)

    # Merge and deduplicate
    seen_ids = set()
    merged = []
    for mem in pg_results + topic_results:
        if mem.id not in seen_ids:
            seen_ids.add(mem.id)
            # Compute similarity score for ranking
            if mem.embedding:
                mem.metadata["similarity_score"] = cosine_similarity(query_embedding, mem.embedding)
            merged.append(mem)

    # Sort: similarity_score * relevance_score (combined ranking)
    merged.sort(
        key=lambda m: m.metadata.get("similarity_score", 0.5) * m.relevance_score,
        reverse=True,
    )

    results = merged[:limit]

    # Update access counts for returned memories
    for mem in results:
        await update_access(mem.id)

    # Publish query event
    event = json.dumps({
        "type": "memory_queried",
        "tier": "episodic",
        "user_id": user_id,
        "query": query_text,
        "results_count": len(results),
        "timestamp": datetime.now(timezone.utc).isoformat(),
    })
    await publish_event(CHANNEL, event)

    return results


async def _query_sqlite(
    user_id: str,
    query_embedding: list[float],
    limit: int,
    min_relevance: float,
) -> list[EpisodicMemory]:
    """Similarity search via SQLite embedding comparison."""
    db = await get_db()
    cursor = await db.execute(
        f"""
        SELECT {_EPISODIC_COLUMNS}
        FROM episodic_memories
        WHERE user_id = ? AND is_decayed = 0
            AND relevance_score >= ?
        ORDER BY created_at DESC
        LIMIT ?
        """,
        (user_id, min_relevance, limit * 2),
    )
    rows = await cursor.fetchall()

    memories = []
    for row in rows:
        mem = _row_to_episodic(row)

        # Compute similarity for ranking
        if mem.embedding and query_embedding:
            sim = cosine_similarity(query_embedding, mem.embedding)
            mem.metadata["similarity_score"] = sim

        memories.append(mem)

    # Sort by similarity
    memories.sort(key=lambda m: m.metadata.get("similarity_score", 0), reverse=True)
    return memories


async def _query_by_topics(user_id: str, query_text: str) -> list[EpisodicMemory]:
    """Topic-based search via SQLite (replaces Memgraph graph traversal)."""
    query_words = [w.lower() for w in query_text.split() if len(w) > 2]
    if not query_words:
        return []

    db = await get_db()

    # Find memories linked to matching topics
    placeholders = ",".join("?" for _ in query_words)
    em_cols = ", ".join(f"em.{c.strip()}" for c in _EPISODIC_COLUMNS.split(","))
    cursor = await db.execute(
        f"""
        SELECT DISTINCT {em_cols}
        FROM episodic_memories em
        JOIN memory_topics mt ON mt.memory_id = em.id
        WHERE em.user_id = ? AND em.is_decayed = 0
          AND LOWER(mt.topic_name) IN ({placeholders})
        ORDER BY em.created_at DESC
        LIMIT 20
        """,
        (user_id, *query_words),
    )
    rows = await cursor.fetchall()

    return [_row_to_episodic(row) for row in rows]


# =============================================
# GET / ACCESS / CLUSTERS
# =============================================

async def get(memory_id: str) -> EpisodicMemory | None:
    """Get a single episodic memory by ID."""
    db = await get_db()
    cursor = await db.execute(
        f"SELECT {_EPISODIC_COLUMNS} FROM episodic_memories WHERE id = ?",
        (memory_id,),
    )
    row = await cursor.fetchone()

    if not row:
        return None

    return _row_to_episodic(row)


async def get_related(memory_id: str, depth: int = 2) -> list[EpisodicMemory]:
    """Get related memories via recursive CTE (replaces Memgraph graph traversal)."""
    db = await get_db()

    cursor = await db.execute(
        """
        WITH RECURSIVE related_chain(id, depth) AS (
            -- Base: direct relationships from the starting memory
            SELECT CASE
                WHEN source_id = ? THEN target_id
                ELSE source_id
            END AS id, 1 AS depth
            FROM memory_relationships
            WHERE (source_id = ? OR target_id = ?)
              AND relationship_type IN ('RELATED_TO', 'SAME_TOPIC')

            UNION

            -- Recursive: follow relationships up to given depth
            SELECT CASE
                WHEN mr.source_id = rc.id THEN mr.target_id
                ELSE mr.source_id
            END AS id, rc.depth + 1
            FROM memory_relationships mr
            JOIN related_chain rc ON (mr.source_id = rc.id OR mr.target_id = rc.id)
            WHERE rc.depth < ?
              AND mr.relationship_type IN ('RELATED_TO', 'SAME_TOPIC')
        )
        SELECT DISTINCT {em_cols}
        FROM episodic_memories em
        JOIN related_chain rc ON em.id = rc.id
        WHERE em.is_decayed = 0 AND em.id != ?
        LIMIT 20
        """.format(em_cols=", ".join(f"em.{c.strip()}" for c in _EPISODIC_COLUMNS.split(","))),
        (memory_id, memory_id, memory_id, depth, memory_id),
    )
    rows = await cursor.fetchall()

    return [_row_to_episodic(row) for row in rows]


async def update_access(memory_id: str):
    """Update access count and last_accessed timestamp. Reinforces relevance."""
    now = datetime.now(timezone.utc).isoformat()

    db = await get_db()
    await db.execute(
        """
        UPDATE episodic_memories
        SET access_count = access_count + 1,
            last_accessed = ?,
            stability = MIN(stability + 0.1, 3.0),
            updated_at = ?
        WHERE id = ?
        """,
        (now, now, memory_id),
    )
    await db.commit()


async def find_topic_clusters(user_id: str) -> list[dict]:
    """Find clusters of related episodic memories grouped by topic (for promotion engine)."""
    db = await get_db()

    cursor = await db.execute(
        """
        SELECT mt.topic_name, COUNT(*) as cnt
        FROM memory_topics mt
        JOIN episodic_memories em ON em.id = mt.memory_id
        WHERE em.user_id = ? AND em.is_decayed = 0
        GROUP BY mt.topic_name
        HAVING COUNT(*) >= 3
        ORDER BY cnt DESC
        """,
        (user_id,),
    )
    topic_rows = await cursor.fetchall()

    clusters = []
    for topic_row in topic_rows:
        topic_name = topic_row[0]
        count = topic_row[1]

        # Fetch the memories in this cluster
        mem_cursor = await db.execute(
            """
            SELECT em.id, em.encrypted_content, em.encryption_key_ref, em.created_at
            FROM episodic_memories em
            JOIN memory_topics mt ON mt.memory_id = em.id
            WHERE mt.topic_name = ? AND em.user_id = ? AND em.is_decayed = 0
            ORDER BY em.created_at DESC
            """,
            (topic_name, user_id),
        )
        mem_rows = await mem_cursor.fetchall()

        memories = []
        for mr in mem_rows:
            memories.append({
                "id": str(mr[0]),
                "content": _decrypt_content(mr[1], mr[2]),
                "created_at": mr[3],
            })

        clusters.append({
            "topic": topic_name,
            "count": count,
            "memories": memories,
        })

    return clusters


async def get_user_memories(
    user_id: str,
    tier: str = "episodic",
    limit: int = 50,
    min_relevance: float = 0.15,
) -> list[EpisodicMemory]:
    """Get all episodic memories for a user, sorted by relevance."""
    db = await get_db()
    cursor = await db.execute(
        f"""
        SELECT {_EPISODIC_COLUMNS}
        FROM episodic_memories
        WHERE user_id = ? AND is_decayed = 0
            AND relevance_score >= ?
        ORDER BY relevance_score DESC, created_at DESC
        LIMIT ?
        """,
        (user_id, min_relevance, limit),
    )
    rows = await cursor.fetchall()

    return [_row_to_episodic(row) for row in rows]
