#!/usr/bin/env python3
"""drift_cache.py — Drift-aware cache with wash-count-based eviction.

The cache contains the most proven items — those that have survived washing.
Cache priority = wash_count × importance.
Items with high wash_count stay in cache longer.
Items with low wash_count get evicted sooner.

This is the "most proven layer" — what the agent reaches for first
because it's been tested and held.

Research basis:
- AMV-L (arXiv:2603.04443): value-driven eviction across lifecycle tiers
- SmartVector (arXiv:2604.20598): access reinforcement and confidence decay

Run under ~/.venvs/memory/bin/python3.
"""

import json
import math
import os
import sys
import time
from datetime import datetime, timezone
from typing import Dict, List, Optional, Tuple

SCRIPTS_DIR = os.path.expanduser("~/.config/opencode/scripts")
CACHE_DIR = os.path.expanduser("~/.config/opencode/memory/cache")
CACHE_FILE = os.path.join(CACHE_DIR, "drift_cache.json")

# Cache configuration
MAX_CACHE_SIZE = 200        # maximum entries in cache
EVICTION_BATCH = 20         # entries to evict when cache is full
MIN_WASH_FOR_CACHE = 1      # minimum wash_count to enter cache
CACHE_BOOST_ROCK = 2.0      # multiplier for rock-status entries
CACHE_BOOST_ESTABLISHED = 1.5  # multiplier for established entries


class DriftCache:
    """Drift-aware cache with wash-count-based priority."""

    def __init__(self):
        os.makedirs(CACHE_DIR, exist_ok=True)
        self.entries: Dict[int, dict] = {}
        self._load()

    def _load(self):
        """Load cache from disk."""
        if os.path.exists(CACHE_FILE):
            try:
                with open(CACHE_FILE) as f:
                    data = json.load(f)
                    self.entries = {int(k): v for k, v in data.items()}
            except (json.JSONDecodeError, ValueError):
                self.entries = {}

    def _save(self):
        """Save cache to disk."""
        with open(CACHE_FILE, "w") as f:
            json.dump({str(k): v for k, v in self.entries.items()}, f, indent=2)

    def _compute_priority(self, entry: dict) -> float:
        """Compute cache priority: wash_count × importance × status boost."""
        wash = entry.get("wash_count", 0) or 0
        imp = entry.get("importance", 0.5) or 0.5
        status = entry.get("rock_status", "washable")

        # Logarithmic wash scaling (diminishing returns)
        wash_factor = min(1.0, math.log1p(wash) / math.log1p(20))

        # Status boost
        status_boost = 1.0
        if status == "rock":
            status_boost = CACHE_BOOST_ROCK
        elif status == "established":
            status_boost = CACHE_BOOST_ESTABLISHED

        return wash_factor * imp * status_boost

    def add(self, entry_id: int, entry: dict):
        """Add an entry to the cache.

        Entries qualify for cache if:
        - wash_count >= MIN_WASH_FOR_CACHE (proven through use), OR
        - importance >= 0.8 (high importance, worth caching even if untested)
        """
        wash = entry.get("wash_count", 0) or 0
        imp = entry.get("importance", 0.5) or 0.5
        if wash < MIN_WASH_FOR_CACHE and imp < 0.8:
            return  # Not proven enough and not important enough

        priority = self._compute_priority(entry)
        self.entries[entry_id] = {
            "priority": priority,
            "added_at": datetime.now(timezone.utc).isoformat(),
            "last_accessed": datetime.now(timezone.utc).isoformat(),
            "access_count": 0,
            "wash_count": wash,
            "rock_status": entry.get("rock_status", "washable"),
        }

        # Evict if over capacity
        if len(self.entries) > MAX_CACHE_SIZE:
            self._evict()

        self._save()

    def access(self, entry_id: int):
        """Record a cache access (reinforces the entry)."""
        if entry_id in self.entries:
            self.entries[entry_id]["last_accessed"] = datetime.now(timezone.utc).isoformat()
            self.entries[entry_id]["access_count"] = \
                self.entries[entry_id].get("access_count", 0) + 1
            self._save()

    def contains(self, entry_id: int) -> bool:
        """Check if an entry is in the cache."""
        return entry_id in self.entries

    def get_priority(self, entry_id: int) -> float:
        """Get the cache priority of an entry."""
        if entry_id in self.entries:
            return self.entries[entry_id].get("priority", 0.0)
        return 0.0

    def _evict(self):
        """Evict lowest-priority entries."""
        # Sort by priority (ascending), evict lowest
        sorted_entries = sorted(
            self.entries.items(),
            key=lambda x: x[1].get("priority", 0.0)
        )
        to_evict = sorted_entries[:EVICTION_BATCH]
        for entry_id, _ in to_evict:
            del self.entries[entry_id]

    def update_from_store(self, db):
        """Sync cache with memory store — update priorities from current wash counts."""
        for entry_id in list(self.entries.keys()):
            row = db.execute(
                "SELECT wash_count, importance, rock_status FROM entries WHERE id = ?",
                (entry_id,)
            ).fetchone()
            if row:
                entry = {
                    "wash_count": row[0] or 0,
                    "importance": row[1] or 0.5,
                    "rock_status": row[2] or "washable",
                }
                self.entries[entry_id]["priority"] = self._compute_priority(entry)
                self.entries[entry_id]["wash_count"] = entry["wash_count"]
                self.entries[entry_id]["rock_status"] = entry["rock_status"]
            else:
                # Entry deleted from store — remove from cache
                del self.entries[entry_id]
        self._save()

    def get_stats(self) -> dict:
        """Get cache statistics."""
        if not self.entries:
            return {"size": 0, "avg_priority": 0.0, "max_priority": 0.0, "min_priority": 0.0, "rock_count": 0, "established_count": 0, "washable_count": 0}

        priorities = [e.get("priority", 0.0) for e in self.entries.values()]
        rock_count = sum(1 for e in self.entries.values() if e.get("rock_status") == "rock")
        est_count = sum(1 for e in self.entries.values() if e.get("rock_status") == "established")

        return {
            "size": len(self.entries),
            "avg_priority": sum(priorities) / len(priorities),
            "max_priority": max(priorities),
            "min_priority": min(priorities),
            "rock_count": rock_count,
            "established_count": est_count,
            "washable_count": len(self.entries) - rock_count - est_count,
        }


# ═══════════════════════════════════════════════════════════════════════════════
# CLI
# ═══════════════════════════════════════════════════════════════════════════════

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Drift-aware cache")
    sub = parser.add_subparsers(dest="cmd")

    # status
    sub.add_parser("status", help="Show cache statistics")

    # sync
    sub.add_parser("sync", help="Sync cache with memory store")

    # add <entry_id>
    p_add = sub.add_parser("add", help="Add entry to cache")
    p_add.add_argument("entry_id", type=int)

    args = parser.parse_args()

    if not args.cmd:
        parser.print_help()
        sys.exit(1)

    cache = DriftCache()

    if args.cmd == "status":
        stats = cache.get_stats()
        print("=== Drift Cache Status ===")
        print(f"Size: {stats['size']}/{MAX_CACHE_SIZE}")
        print(f"Avg priority: {stats['avg_priority']:.4f}")
        print(f"Rock: {stats['rock_count']}")
        print(f"Established: {stats['established_count']}")
        print(f"Washable: {stats['washable_count']}")

    elif args.cmd == "sync":
        sys.path.insert(0, SCRIPTS_DIR)
        import memory_store
        db = memory_store.connect()
        cache.update_from_store(db)
        db.close()
        stats = cache.get_stats()
        print(f"Synced. Cache size: {stats['size']}")

    elif args.cmd == "add":
        sys.path.insert(0, SCRIPTS_DIR)
        import memory_store
        db = memory_store.connect()
        row = db.execute(
            "SELECT id, text, importance, wash_count, rock_status FROM entries WHERE id = ?",
            (args.entry_id,)
        ).fetchone()
        if row:
            entry = {
                "text": row[1],
                "importance": row[2] or 0.5,
                "wash_count": row[3] or 0,
                "rock_status": row[4] or "washable",
            }
            cache.add(args.entry_id, entry)
            print(f"Added entry {args.entry_id} to cache")
        else:
            print(f"Entry {args.entry_id} not found")
        db.close()


if __name__ == "__main__":
    main()