"""
Phase 52: Long-term Agent Memory with Semantic Organization.

Gives each agent a persistent memory store organized by category and
importance. Supports keyword-based recall, importance scoring, and
automatic consolidation of stale/low-value memories.

Storage:
  state/broker/agent_memory/{agent_id}/
    memories.json    — all memories for this agent
"""

import json
import logging
import time
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

from brynq.server import BROKER_DIR, _write_json_sync, _read_jsonl_sync, _append_jsonl_sync

logger = logging.getLogger(__name__)

# ── Paths ──────────────────────────────────────────────────────────────────

AGENT_MEMORY_DIR = BROKER_DIR / "agent_memory"
AGENT_MEMORY_DIR.mkdir(parents=True, exist_ok=True)


# ── Internal helpers ──────────────────────────────────────────────────────

def _agent_dir(agent_id: str) -> Path:
    d = AGENT_MEMORY_DIR / agent_id
    d.mkdir(parents=True, exist_ok=True)
    return d


def _memories_file(agent_id: str) -> Path:
    return _agent_dir(agent_id) / "memories.json"


def _read_memories(agent_id: str) -> List[dict]:
    path = _memories_file(agent_id)
    if not path.exists():
        return []
    try:
        return json.loads(path.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return []


def _save_memories(agent_id: str, memories: List[dict]) -> None:
    _write_json_sync(_memories_file(agent_id), memories)


def _keyword_score(text: str, query: str) -> int:
    """Simple keyword overlap score between text and query."""
    text_lower = text.lower()
    query_words = query.lower().split()
    return sum(1 for w in query_words if w in text_lower)


# ── Public API ────────────────────────────────────────────────────────────

def remember(
    agent_id: str,
    key: str,
    value: Any,
    category: str = "general",
    importance: int = 5,
) -> dict:
    """Store a memory for an agent.

    Args:
        agent_id: Agent identifier.
        key: Unique key for this memory.
        value: The value to store (any JSON-serializable data).
        category: Category tag (e.g. "general", "task", "preference").
        importance: 1-10 scale of importance.

    Returns:
        The stored memory record.
    """
    importance = max(1, min(10, importance))
    memories = _read_memories(agent_id)

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

    # Update existing memory if key already exists
    for m in memories:
        if m["key"] == key:
            m["value"] = value
            m["category"] = category
            m["importance"] = importance
            m["updated_at"] = now
            m["access_count"] = m.get("access_count", 0)
            _save_memories(agent_id, memories)
            logger.info("Updated memory '%s' for agent %s", key, agent_id)
            return m

    memory = {
        "memory_id": uuid.uuid4().hex[:12],
        "key": key,
        "value": value,
        "category": category,
        "importance": importance,
        "created_at": now,
        "updated_at": now,
        "access_count": 0,
    }

    memories.append(memory)
    _save_memories(agent_id, memories)
    logger.info("Stored memory '%s' for agent %s", key, agent_id)
    return memory


def recall(
    agent_id: str,
    key: Optional[str] = None,
    category: Optional[str] = None,
    query: Optional[str] = None,
) -> List[dict]:
    """Retrieve memories by key, category, or text search.

    Args:
        agent_id: Agent identifier.
        key: Exact key match.
        category: Filter by category.
        query: Text search across key and value.

    Returns:
        List of matching memory records.
    """
    memories = _read_memories(agent_id)
    results = []

    for m in memories:
        if key is not None and m["key"] != key:
            continue
        if category is not None and m.get("category") != category:
            continue
        if query is not None:
            searchable = f"{m['key']} {json.dumps(m['value'])}".lower()
            if not any(w in searchable for w in query.lower().split()):
                continue
        results.append(m)

    # Update access counts
    result_ids = {r["memory_id"] for r in results}
    for m in memories:
        if m["memory_id"] in result_ids:
            m["access_count"] = m.get("access_count", 0) + 1
    _save_memories(agent_id, memories)

    return results


def forget(agent_id: str, key: str) -> bool:
    """Delete a memory by key.

    Returns:
        True if memory was found and deleted, False otherwise.
    """
    memories = _read_memories(agent_id)
    before = len(memories)
    memories = [m for m in memories if m["key"] != key]
    if len(memories) == before:
        return False
    _save_memories(agent_id, memories)
    logger.info("Forgot memory '%s' for agent %s", key, agent_id)
    return True


def get_context(
    agent_id: str,
    task_description: str,
    limit: int = 10,
) -> List[dict]:
    """Retrieve the most relevant memories for a task description.

    Uses keyword matching + importance weighting to rank memories.

    Args:
        agent_id: Agent identifier.
        task_description: Description of the task to find context for.
        limit: Maximum number of memories to return.

    Returns:
        List of memories sorted by relevance (most relevant first).
    """
    memories = _read_memories(agent_id)
    if not memories:
        return []

    scored = []
    for m in memories:
        searchable = f"{m['key']} {json.dumps(m['value'])} {m.get('category', '')}"
        kw_score = _keyword_score(searchable, task_description)
        importance = m.get("importance", 5)
        # Combined relevance: keyword hits * 2 + importance
        relevance = kw_score * 2 + importance
        scored.append((relevance, m))

    scored.sort(key=lambda x: x[0], reverse=True)

    # Only return memories with some relevance (keyword match > 0) or high importance
    results = []
    for relevance, m in scored[:limit]:
        results.append(m)

    return results


def consolidate(agent_id: str) -> dict:
    """Merge similar memories and prune low-importance old ones.

    - Memories with importance <= 2 and older than 30 days are pruned.
    - Memories with the same category and similar keys are flagged.

    Returns:
        Dict with pruned_count and remaining_count.
    """
    memories = _read_memories(agent_id)
    now = datetime.now(timezone.utc)

    pruned = 0
    kept = []

    for m in memories:
        importance = m.get("importance", 5)
        created = m.get("created_at", "")

        should_prune = False
        if importance <= 2 and created:
            try:
                ct = datetime.fromisoformat(created)
                age_days = (now - ct).total_seconds() / 86400
                if age_days > 30:
                    should_prune = True
            except (ValueError, TypeError):
                pass

        if should_prune:
            pruned += 1
        else:
            kept.append(m)

    _save_memories(agent_id, kept)
    logger.info("Consolidated agent %s: pruned %d, kept %d", agent_id, pruned, len(kept))
    return {"pruned_count": pruned, "remaining_count": len(kept)}


def get_memory_stats(agent_id: str) -> dict:
    """Return statistics about an agent's memories.

    Returns:
        Dict with total, by_category, oldest, newest timestamps.
    """
    memories = _read_memories(agent_id)

    if not memories:
        return {
            "total": 0,
            "by_category": {},
            "oldest": None,
            "newest": None,
        }

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

    for m in memories:
        cat = m.get("category", "general")
        by_category[cat] = by_category.get(cat, 0) + 1
        ts = m.get("created_at")
        if ts:
            timestamps.append(ts)

    timestamps.sort()

    return {
        "total": len(memories),
        "by_category": by_category,
        "oldest": timestamps[0] if timestamps else None,
        "newest": timestamps[-1] if timestamps else None,
    }
