"""
agentpaas.services.memory_service — Persistent memory service

Three-tier memory model:
  Working  — current session context (already exists in Context)
  Episodic — recent conversation summaries (auto-compressed, 7 days)
  Semantic — permanent knowledge (user preferences, project structure)

Storage: SQLite for dev, Redis for prod.
"""
from __future__ import annotations

import json
import os
import time
import hashlib
from dataclasses import dataclass, field, asdict
from typing import Any, Dict, List, Optional

# ════════════════════════════════════════════════════════════
# Memory Entry
# ════════════════════════════════════════════════════════════

@dataclass
class MemoryEntry:
    """A single memory record."""
    key: str
    content: str
    tier: str = "semantic"      # working / episodic / semantic
    tags: List[str] = field(default_factory=list)
    source: str = ""            # agent_id or "user"
    created_at: float = 0
    expires_at: float = 0       # 0 = never expires
    access_count: int = 0
    relevance_score: float = 0  # for search ranking

    def __post_init__(self):
        if self.created_at == 0:
            self.created_at = time.time()
        if not self.key:
            self.key = hashlib.md5(self.content.encode()).hexdigest()[:12]

    def is_expired(self) -> bool:
        return self.expires_at > 0 and time.time() > self.expires_at

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, data: dict) -> MemoryEntry:
        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})


# ════════════════════════════════════════════════════════════
# Memory Store (SQLite backend)
# ════════════════════════════════════════════════════════════

class MemoryStore:
    """Persistent memory store using SQLite."""

    def __init__(self, db_path: str = ""):
        self.db_path = db_path or os.path.expanduser("~/.lambdagent/memory.db")
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        self._init_db()

    def _init_db(self):
        import sqlite3
        conn = sqlite3.connect(self.db_path)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS memories (
                key TEXT PRIMARY KEY,
                content TEXT NOT NULL,
                tier TEXT DEFAULT 'semantic',
                tags TEXT DEFAULT '[]',
                source TEXT DEFAULT '',
                created_at REAL,
                expires_at REAL DEFAULT 0,
                access_count INTEGER DEFAULT 0,
                relevance_score REAL DEFAULT 0
            )
        """)
        conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_tier ON memories(tier)
        """)
        conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_source ON memories(source)
        """)
        conn.commit()
        conn.close()

    def _conn(self):
        import sqlite3
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn

    def store(self, entry: MemoryEntry) -> str:
        """Store or update a memory entry."""
        conn = self._conn()
        conn.execute("""
            INSERT OR REPLACE INTO memories
            (key, content, tier, tags, source, created_at, expires_at, access_count, relevance_score)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (entry.key, entry.content, entry.tier, json.dumps(entry.tags),
              entry.source, entry.created_at, entry.expires_at,
              entry.access_count, entry.relevance_score))
        conn.commit()
        conn.close()
        return entry.key

    def recall(self, query: str, tier: str = "", source: str = "",
               top_k: int = 5) -> List[MemoryEntry]:
        """Recall memories by relevance to query."""
        conn = self._conn()

        # Build WHERE clause
        conditions = ["expires_at = 0 OR expires_at > ?"]
        params: list = [time.time()]
        if tier:
            conditions.append("tier = ?")
            params.append(tier)
        if source:
            conditions.append("source = ?")
            params.append(source)

        where = " AND ".join(conditions)
        rows = conn.execute(
            f"SELECT * FROM memories WHERE {where} ORDER BY created_at DESC LIMIT 100",
            params
        ).fetchall()
        conn.close()

        # Score by simple keyword matching (TF-IDF would be better but this is fast)
        query_words = set(query.lower().split())
        results = []
        for row in rows:
            entry = MemoryEntry(
                key=row["key"], content=row["content"], tier=row["tier"],
                tags=json.loads(row["tags"]), source=row["source"],
                created_at=row["created_at"], expires_at=row["expires_at"],
                access_count=row["access_count"],
            )
            # Relevance scoring
            content_words = set(entry.content.lower().split())
            tag_words = set(t.lower() for t in entry.tags)
            overlap = len(query_words & (content_words | tag_words))
            entry.relevance_score = overlap / max(1, len(query_words))
            results.append(entry)

        # Sort by relevance, then by recency
        results.sort(key=lambda e: (-e.relevance_score, -e.created_at))

        # Update access count for returned entries
        if results[:top_k]:
            conn = self._conn()
            for entry in results[:top_k]:
                conn.execute("UPDATE memories SET access_count = access_count + 1 WHERE key = ?",
                             (entry.key,))
            conn.commit()
            conn.close()

        return results[:top_k]

    def get(self, key: str) -> Optional[MemoryEntry]:
        """Get a specific memory by key."""
        conn = self._conn()
        row = conn.execute("SELECT * FROM memories WHERE key = ?", (key,)).fetchone()
        conn.close()
        if not row:
            return None
        return MemoryEntry(
            key=row["key"], content=row["content"], tier=row["tier"],
            tags=json.loads(row["tags"]), source=row["source"],
            created_at=row["created_at"], expires_at=row["expires_at"],
            access_count=row["access_count"],
        )

    def forget(self, key: str) -> bool:
        """Delete a memory entry."""
        conn = self._conn()
        conn.execute("DELETE FROM memories WHERE key = ?", (key,))
        affected = conn.total_changes
        conn.commit()
        conn.close()
        return affected > 0

    def list_all(self, tier: str = "", source: str = "", limit: int = 50) -> List[MemoryEntry]:
        """List memories."""
        conn = self._conn()
        conditions = ["1=1"]
        params: list = []
        if tier:
            conditions.append("tier = ?")
            params.append(tier)
        if source:
            conditions.append("source = ?")
            params.append(source)
        where = " AND ".join(conditions)
        rows = conn.execute(
            f"SELECT * FROM memories WHERE {where} ORDER BY created_at DESC LIMIT ?",
            params + [limit]
        ).fetchall()
        conn.close()
        return [
            MemoryEntry(key=r["key"], content=r["content"], tier=r["tier"],
                        tags=json.loads(r["tags"]), source=r["source"],
                        created_at=r["created_at"], access_count=r["access_count"])
            for r in rows
        ]

    def cleanup_expired(self) -> int:
        """Remove expired episodic memories."""
        conn = self._conn()
        conn.execute("DELETE FROM memories WHERE expires_at > 0 AND expires_at < ?",
                     (time.time(),))
        affected = conn.total_changes
        conn.commit()
        conn.close()
        return affected

    def stats(self) -> dict:
        """Memory statistics."""
        conn = self._conn()
        total = conn.execute("SELECT COUNT(*) FROM memories").fetchone()[0]
        by_tier = {}
        for row in conn.execute("SELECT tier, COUNT(*) as cnt FROM memories GROUP BY tier"):
            by_tier[row["tier"]] = row["cnt"]
        conn.close()
        return {"total": total, "by_tier": by_tier, "db_path": self.db_path}


# ════════════════════════════════════════════════════════════
# Agent-side tool functions
# ════════════════════════════════════════════════════════════

_default_store = None

def _get_store() -> MemoryStore:
    global _default_store
    if _default_store is None:
        _default_store = MemoryStore()
    return _default_store


def memory_store(input_val: Any) -> str:
    """Store a memory. Input: {content, tags?, tier?, key?}"""
    params = _parse(input_val)
    content = params.get("content", "")
    if not content:
        return "[ERROR] content is required"

    entry = MemoryEntry(
        key=params.get("key", ""),
        content=content,
        tier=params.get("tier", "semantic"),
        tags=params.get("tags", []),
        source=params.get("source", "agent"),
    )
    ttl = params.get("ttl", 0)
    if ttl > 0:
        entry.expires_at = time.time() + ttl
        entry.tier = "episodic"

    key = _get_store().store(entry)
    return f"[OK] Memory stored: key={key}, tier={entry.tier}"


def memory_recall(input_val: Any) -> str:
    """Recall memories by query. Input: {query, top_k?, tier?}"""
    params = _parse(input_val)
    query = params.get("query", "")
    if not query:
        return "[ERROR] query is required"

    results = _get_store().recall(
        query,
        tier=params.get("tier", ""),
        source=params.get("source", ""),
        top_k=params.get("top_k", 5),
    )

    if not results:
        return f"[NO_MATCH] No memories matching '{query}'"

    lines = []
    for entry in results:
        tags = f" [{', '.join(entry.tags)}]" if entry.tags else ""
        lines.append(f"[{entry.tier}] (score={entry.relevance_score:.2f}{tags})")
        lines.append(f"  {entry.content[:300]}")
        lines.append("")
    return "\n".join(lines)


def memory_list(input_val: Any) -> str:
    """List all memories. Input: {tier?, limit?}"""
    params = _parse(input_val) if input_val else {}
    entries = _get_store().list_all(
        tier=params.get("tier", ""),
        limit=params.get("limit", 20),
    )
    if not entries:
        return "[EMPTY] No memories stored"

    lines = []
    for e in entries:
        age = _format_age(time.time() - e.created_at)
        tags = f" [{', '.join(e.tags)}]" if e.tags else ""
        lines.append(f"  [{e.tier}] {e.key}: {e.content[:80]}...{tags} ({age})")
    stats = _get_store().stats()
    lines.insert(0, f"Memories: {stats['total']} total ({stats['by_tier']})")
    return "\n".join(lines)


def memory_forget(input_val: Any) -> str:
    """Forget a memory. Input: {key}"""
    params = _parse(input_val)
    key = params.get("key", "")
    if not key:
        return "[ERROR] key is required"
    if _get_store().forget(key):
        return f"[OK] Memory forgotten: {key}"
    return f"[ERROR] Memory not found: {key}"


def _format_age(seconds: float) -> str:
    if seconds < 60:
        return "just now"
    if seconds < 3600:
        return f"{int(seconds/60)}m ago"
    if seconds < 86400:
        return f"{int(seconds/3600)}h ago"
    return f"{int(seconds/86400)}d ago"


def _parse(input_val: Any) -> dict:
    if isinstance(input_val, dict):
        return input_val
    if isinstance(input_val, str):
        try:
            return json.loads(input_val)
        except (json.JSONDecodeError, ValueError):
            return {"query": input_val} if len(input_val) < 200 else {"content": input_val}
    return {}
