"""ARIS Memory — Cross-session, cross-model persistent memory."""

import json
import os
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Optional


@dataclass
class MemoryEntry:
    """A single memory entry."""
    key: str
    value: str
    category: str
    model: str
    created_at: float
    updated_at: float
    access_count: int = 0
    tags: list[str] = field(default_factory=list)


class MemoryStore:
    """5-layer memory: Hot(1h) > Warm(24h) > Cool(7d) > Cold(30d) > Archive."""

    def __init__(self, base_dir: Optional[str] = None) -> None:
        self._base = Path(base_dir or os.path.expanduser("~/.aris/memory"))
        self._base.mkdir(parents=True, exist_ok=True)
        self._store_file = self._base / "memories.json"
        self._memories: dict[str, MemoryEntry] = {}
        self._load()

    def _load(self) -> None:
        if self._store_file.exists():
            try:
                data = json.loads(self._store_file.read_text())
                for key, entry in data.items():
                    self._memories[key] = MemoryEntry(**entry)
            except (json.JSONDecodeError, TypeError):
                self._memories = {}

    def _save(self) -> None:
        self._dirty = True
        self._dirty_count = getattr(self, "_dirty_count", 0) + 1
        # Batch: write every 100 ops or on explicit flush
        if self._dirty_count < 100:
            return
        self.flush()

    def flush(self) -> None:
        """Force write to disk."""
        data = {}
        for key, entry in self._memories.items():
            data[key] = {
                "key": entry.key, "value": entry.value,
                "category": entry.category, "model": entry.model,
                "created_at": entry.created_at, "updated_at": entry.updated_at,
                "access_count": entry.access_count, "tags": entry.tags,
            }
        self._store_file.write_text(json.dumps(data, indent=2))
        self._dirty = False
        self._dirty_count = 0

    def save(self, key: str, value: str, category: str = "context",
             model: str = "claude", tags: Optional[list[str]] = None) -> MemoryEntry:
        now = time.time()
        if key in self._memories:
            entry = self._memories[key]
            entry.value = value
            entry.updated_at = now
            entry.category = category
            entry.model = model
            if tags:
                entry.tags = list(set(entry.tags + tags))
        else:
            entry = MemoryEntry(
                key=key, value=value, category=category,
                model=model, created_at=now, updated_at=now,
                tags=tags or [],
            )
            self._memories[key] = entry
        self._save()
        return entry

    def recall(self, key: str) -> Optional[MemoryEntry]:
        entry = self._memories.get(key)
        if entry:
            entry.access_count += 1
            self._save()
        return entry

    def search(self, query: str, category: Optional[str] = None,
               limit: int = 10) -> list[MemoryEntry]:
        query_lower = query.lower()
        results = []
        for entry in self._memories.values():
            if category and entry.category != category:
                continue
            score = 0
            if query_lower in entry.key.lower():
                score += 3
            if query_lower in entry.value.lower():
                score += 2
            if any(query_lower in tag.lower() for tag in entry.tags):
                score += 1
            if score > 0:
                results.append((score, entry))
        results.sort(key=lambda x: (-x[0], -x[1].updated_at))
        return [entry for _, entry in results[:limit]]

    def list_all(self, category: Optional[str] = None, limit: int = 50) -> list[MemoryEntry]:
        entries = list(self._memories.values())
        if category:
            entries = [e for e in entries if e.category == category]
        entries.sort(key=lambda e: -e.updated_at)
        return entries[:limit]

    def delete(self, key: str) -> bool:
        if key in self._memories:
            del self._memories[key]
            self.flush()  # Always flush on delete
            return True
        return False

    def get_layer(self, entry: MemoryEntry) -> str:
        age = time.time() - entry.updated_at
        if age < 3600:
            return "hot"
        elif age < 86400:
            return "warm"
        elif age < 604800:
            return "cool"
        elif age < 2592000:
            return "cold"
        return "archive"

    def stats(self) -> dict[str, Any]:
        layers = {"hot": 0, "warm": 0, "cool": 0, "cold": 0, "archive": 0}
        categories: dict[str, int] = {}
        models: dict[str, int] = {}
        for entry in self._memories.values():
            layers[self.get_layer(entry)] += 1
            categories[entry.category] = categories.get(entry.category, 0) + 1
            models[entry.model] = models.get(entry.model, 0) + 1
        return {"total": len(self._memories), "layers": layers,
                "categories": categories, "models": models}
