"""
Phase 13: Swarm Shared Memory

Long-term shared memory for agent swarms. Agents store and retrieve
knowledge using TF-IDF text scoring with proper relevance ranking,
retention policies, and metadata filtering.
"""

import json
import math
import os
import re
import time
import uuid
from collections import Counter
from pathlib import Path
from typing import Optional


MEMORY_DIR = Path("state/broker/swarm_memory")
os.makedirs(MEMORY_DIR, exist_ok=True)

# ── Index for fast search (rebuilt on demand) ─────────────────────────────
_INDEX_FILE = MEMORY_DIR / "_index.json"


def _tokenize(text: str) -> list[str]:
    """Lowercase, split on non-alphanumeric, drop short words."""
    return [w for w in re.split(r"[^a-z0-9]+", text.lower()) if len(w) > 1]


def _rebuild_index() -> dict:
    """Build an inverted index of token -> [memory_id, ...] for fast search."""
    index: dict[str, list[str]] = {}
    doc_count = 0
    for fp in MEMORY_DIR.glob("*.json"):
        if fp.name.startswith("_"):
            continue
        try:
            data = json.loads(fp.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        doc_count += 1
        mid = data.get("id", fp.stem)
        tokens = set(_tokenize(data.get("content", "")))
        # Also index metadata values and tags
        meta = data.get("metadata", {})
        if isinstance(meta, dict):
            for v in meta.values():
                if isinstance(v, str):
                    tokens.update(_tokenize(v))
        for tag in data.get("tags", []):
            tokens.update(_tokenize(tag))
        for token in tokens:
            index.setdefault(token, []).append(mid)

    index_data = {"tokens": index, "doc_count": doc_count, "built_at": time.time()}
    try:
        _INDEX_FILE.write_text(json.dumps(index_data), encoding="utf-8")
    except OSError:
        pass
    return index_data


def _load_index() -> dict:
    """Load index, rebuild if stale or missing."""
    try:
        if _INDEX_FILE.exists():
            data = json.loads(_INDEX_FILE.read_text("utf-8"))
            # Rebuild if older than 60 seconds
            if time.time() - data.get("built_at", 0) < 60:
                return data
    except (json.JSONDecodeError, OSError):
        pass
    return _rebuild_index()


# ── Store ─────────────────────────────────────────────────────────────────


def store_memory(
    content: str,
    author: str,
    metadata: Optional[dict[str, str]] = None,
    tags: Optional[list[str]] = None,
    importance: float = 0.5,
) -> str:
    """Store a piece of knowledge in the swarm's shared memory.

    Args:
        content: The text content to store.
        author: Agent ID or name that created this memory.
        metadata: Optional key-value metadata (e.g., {"topic": "architecture"}).
        tags: Optional list of tags for categorization.
        importance: Float 0-1 indicating how important this memory is (affects search ranking).

    Returns:
        The memory node ID.
    """
    node_id = uuid.uuid4().hex[:12]
    importance = max(0.0, min(1.0, importance))

    node = {
        "id": node_id,
        "content": content,
        "metadata": metadata or {},
        "tags": tags or [],
        "timestamp": time.time(),
        "author": author,
        "importance": importance,
        "access_count": 0,
        "last_accessed": 0.0,
    }

    fp = MEMORY_DIR / f"{node_id}.json"
    fp.write_text(json.dumps(node, indent=2), encoding="utf-8")

    # Invalidate index
    try:
        _INDEX_FILE.unlink(missing_ok=True)
    except OSError:
        pass

    return node_id


# ── Retrieve ──────────────────────────────────────────────────────────────


def retrieve_memory(
    query: str,
    limit: int = 5,
    author: Optional[str] = None,
    tags: Optional[list[str]] = None,
    min_importance: float = 0.0,
) -> list[dict]:
    """Retrieve memories using TF-IDF scoring with filters.

    Scoring:
    - TF: frequency of query words in content
    - IDF: words that appear in fewer documents score higher
    - Importance multiplier: author-declared importance
    - Recency bias: newer memories get a small boost
    """
    query_tokens = _tokenize(query)
    if not query_tokens:
        return []

    index = _load_index()
    token_index = index.get("tokens", {})
    doc_count = max(index.get("doc_count", 1), 1)

    # Collect candidate memory IDs from index
    candidates: Counter = Counter()
    for qt in query_tokens:
        for mid in token_index.get(qt, []):
            candidates[mid] += 1

    if not candidates:
        return []

    results = []
    now = time.time()

    for mid, _ in candidates.most_common(limit * 3):  # over-fetch for filtering
        fp = MEMORY_DIR / f"{mid}.json"
        if not fp.exists():
            continue
        try:
            data = json.loads(fp.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue

        # Apply filters
        if author and data.get("author") != author:
            continue
        if min_importance > 0 and data.get("importance", 0) < min_importance:
            continue
        if tags:
            node_tags = set(t.lower() for t in data.get("tags", []))
            if not any(t.lower() in node_tags for t in tags):
                continue

        # Score
        content_tokens = _tokenize(data.get("content", ""))
        content_counter = Counter(content_tokens)
        total_tokens = max(len(content_tokens), 1)

        score = 0.0
        for qt in query_tokens:
            tf = content_counter.get(qt, 0) / total_tokens
            # IDF: log(total_docs / docs_containing_term)
            df = len(token_index.get(qt, []))
            idf = math.log(1 + doc_count / (1 + df))
            score += tf * idf

        # Importance multiplier (0.5 base + up to 0.5 from importance)
        score *= (0.5 + data.get("importance", 0.5) * 0.5)

        # Recency bias: up to 10% boost, decays over 30 days
        age_hours = (now - data.get("timestamp", 0)) / 3600
        recency = max(0.0, 1.0 - age_hours / 720)
        score *= (1.0 + recency * 0.1)

        if score > 0:
            data["_score"] = round(score, 4)
            results.append(data)

    results.sort(key=lambda x: x.get("_score", 0), reverse=True)

    # Track access on returned results
    for item in results[:limit]:
        _touch_memory(item["id"])

    return results[:limit]


def _touch_memory(memory_id: str) -> None:
    """Increment access count and update last_accessed."""
    fp = MEMORY_DIR / f"{memory_id}.json"
    try:
        data = json.loads(fp.read_text("utf-8"))
        data["access_count"] = data.get("access_count", 0) + 1
        data["last_accessed"] = time.time()
        fp.write_text(json.dumps(data, indent=2), encoding="utf-8")
    except (json.JSONDecodeError, OSError):
        pass


# ── List & Manage ─────────────────────────────────────────────────────────


def list_memories(
    limit: int = 50,
    author: Optional[str] = None,
    sort_by: str = "timestamp",
) -> list[dict]:
    """List memories. sort_by: 'timestamp', 'importance', 'access_count'."""
    results = []
    for fp in MEMORY_DIR.glob("*.json"):
        if fp.name.startswith("_"):
            continue
        try:
            data = json.loads(fp.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue
        if author and data.get("author") != author:
            continue
        results.append(data)

    key = sort_by if sort_by in ("timestamp", "importance", "access_count") else "timestamp"
    results.sort(key=lambda x: x.get(key, 0), reverse=True)
    return results[:limit]


def delete_memory(memory_id: str) -> dict:
    """Delete a specific memory node."""
    fp = MEMORY_DIR / f"{memory_id}.json"
    if not fp.exists():
        return {"ok": False, "error": "Memory not found"}
    fp.unlink()
    try:
        _INDEX_FILE.unlink(missing_ok=True)
    except OSError:
        pass
    return {"ok": True, "deleted": memory_id}


def prune_memories(max_age_days: int = 90, min_importance: float = 0.1) -> dict:
    """Remove old, low-importance memories. Keeps anything above min_importance
    regardless of age, and anything newer than max_age_days regardless of importance."""
    cutoff = time.time() - (max_age_days * 86400)
    removed = 0

    for fp in MEMORY_DIR.glob("*.json"):
        if fp.name.startswith("_"):
            continue
        try:
            data = json.loads(fp.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue

        ts = data.get("timestamp", 0)
        imp = data.get("importance", 0)

        # Keep if recent enough OR important enough
        if ts > cutoff or imp >= min_importance:
            continue

        fp.unlink(missing_ok=True)
        removed += 1

    if removed:
        try:
            _INDEX_FILE.unlink(missing_ok=True)
        except OSError:
            pass

    return {"pruned": removed}


def memory_stats() -> dict:
    """Return statistics about the shared memory store."""
    total = 0
    total_size = 0
    authors: Counter = Counter()
    for fp in MEMORY_DIR.glob("*.json"):
        if fp.name.startswith("_"):
            continue
        total += 1
        total_size += fp.stat().st_size
        try:
            data = json.loads(fp.read_text("utf-8"))
            authors[data.get("author", "unknown")] += 1
        except (json.JSONDecodeError, OSError):
            pass

    return {
        "total_memories": total,
        "storage_bytes": total_size,
        "authors": dict(authors.most_common(20)),
    }
