#!/usr/bin/env python3
"""
ump_decay.py — Apply per-kind confidence decay to UMP memory records.

Decay rates (per day since last access):
  identity     λ=0.0001  (very slow, never decays much)
  semantic     λ=0.001   (slow, facts decay over months)
  procedural   λ=0.005   (medium, skills fade if unused)
  working      λ=0.05    (fast, session context fades fast)
  episodic     λ=0.01    (events fade over weeks)
  note         λ=0.003   (medium-slow, session notes)

Formula: confidence_new = confidence_old * exp(-λ * days_since_reference)
  where days_since_reference =
    days since r.time.modified if present
    else days since r.time.created

Floor: 0.05 (never zero).

Status transitions (after applying decay):
  candidate  ->  active    if confidence >= 0.5
  active     ->  archived  if confidence < 0.2
  archived   stays        (no resurrection)
  tombstoned stays        (skip in retrieval, but update confidence for log)

Add fields on first run if missing:
  r.time.modified (default to r.time.created if missing)

Usage:
  ump_decay.py --dry-run                 # show what would change, don't write
  ump_decay.py --apply                   # actually modify memory.ump.json
  ump_decay.py --apply --archive-summary # also print archive candidates

Backups:
  Before --apply, copy memory.ump.json -> memory.ump.json.bak.YYYY-MM-DDTHHMMSSZ

Pure function:
  apply_decay(records, dry_run=True) -> (records, report)
  Mutates a copy of records in-place. Does NOT touch the disk.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import shutil
import sys
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

# ---- Decay configuration --------------------------------------------------

DECAY_RATES: dict[str, float] = {
    "identity":   0.0001,
    "semantic":   0.001,
    "procedural": 0.005,
    "working":    0.05,
    "episodic":   0.01,
    "note":       0.003,
}

CONF_FLOOR = 0.05
ARCHIVE_THRESHOLD = 0.20
PROMOTE_THRESHOLD = 0.50

DEFAULT_PATH = Path(
    "/root/.openclaw/agents/main/workspace/state/ump-local/memory.ump.json"
)


# ---- Time parsing ---------------------------------------------------------

def parse_iso(ts: Any) -> datetime | None:
    """Parse an ISO 8601 timestamp; return None on failure or missing."""
    if not ts or not isinstance(ts, str):
        return None
    s = ts.replace("Z", "+00:00")
    try:
        dt = datetime.fromisoformat(s)
    except ValueError:
        return None
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return dt


# ---- Field accessors (safe) ----------------------------------------------

def _get_time(rec: dict) -> dict:
    t = rec.get("time")
    return t if isinstance(t, dict) else {}


def _get_lifecycle(rec: dict) -> dict:
    lc = rec.get("lifecycle")
    return lc if isinstance(lc, dict) else {}


def get_kind(rec: dict) -> str:
    k = rec.get("kind")
    return k if isinstance(k, str) else "unknown"


def get_status(rec: dict) -> str:
    s = _get_lifecycle(rec).get("status")
    return s if isinstance(s, str) else "active"


def get_confidence(rec: dict) -> float:
    c = _get_lifecycle(rec).get("confidence")
    if isinstance(c, (int, float)):
        return float(c)
    return 1.0


def get_created_at(rec: dict) -> datetime | None:
    return parse_iso(_get_time(rec).get("created"))


def get_modified_at(rec: dict) -> datetime | None:
    return parse_iso(_get_time(rec).get("modified"))


# ---- Normalization --------------------------------------------------------

def normalize_record(rec: dict) -> dict:
    """
    Ensure required fields exist; defaults added in-place.
      - time.modified = time.created if missing
      - lifecycle.confidence = 1.0 if missing
      - lifecycle.status = 'active' if missing
    """
    t = rec.setdefault("time", {})
    if "created" in t and ("modified" not in t or t["modified"] is None):
        t["modified"] = t["created"]
    lc = rec.setdefault("lifecycle", {})
    if "confidence" not in lc or lc["confidence"] is None:
        lc["confidence"] = 1.0
    if "status" not in lc or lc["status"] is None:
        lc["status"] = "active"
    return rec


# ---- Decay math -----------------------------------------------------------

def compute_new_confidence(old: float, lam: float, days_since: float) -> float:
    raw = old * math.exp(-lam * days_since)
    if raw < CONF_FLOOR:
        return CONF_FLOOR
    return raw


# ---- Pure batch function (testable) ---------------------------------------

def apply_decay(
    records: list[dict],
    dry_run: bool = True,
) -> tuple[list[dict], dict]:
    """
    Apply decay + status transitions to a list of records. Pure function.

    Args:
      records: list of memory records (dicts). Will be COPIED at the top level
               so the caller's list is not mutated, but record dicts themselves
               are mutated in place (this is intentional — preserves nested
               references and matches the schema's mutability expectations).
      dry_run: if True, don't mutate any records; just compute what would change.

    Returns:
      (records, report) where report is a dict:
        {
          'total': int,
          'before': Counter[status],
          'after': Counter[status],
          'changes': [change_dict, ...],   # one per record
          'archive_candidates': [...],     # active -> archived transitions
          'promotion_candidates': [...],   # candidate -> active transitions
          'skipped_tombstoned': int,
          'edge_cases': [str, ...],        # weird records, parsing issues
        }
    """
    now = datetime.now(timezone.utc)
    changes: list[dict] = []
    edge_cases: list[str] = []
    archive_candidates: list[dict] = []
    promotion_candidates: list[dict] = []

    before_status: Counter[str] = Counter()
    after_status: Counter[str] = Counter()
    skipped_tombstoned = 0

    # Copy records (shallow) so dry_run really is non-mutating.
    if dry_run:
        records = [dict(r) for r in records]

    for rec in records:
        if not isinstance(rec, dict):
            edge_cases.append(f"non-dict record: {type(rec).__name__}")
            continue

        # Normalize first (adds defaults, mutates rec).
        try:
            normalize_record(rec)
        except Exception as e:
            edge_cases.append(f"normalize failed for {rec.get('id','?')}: {e!r}")
            continue

        urn = rec.get("id", "<no-id>")
        kind = get_kind(rec)
        old_status = get_status(rec)
        old_conf = get_confidence(rec)

        before_status[old_status] += 1

        change: dict[str, Any] = {
            "urn": urn,
            "kind": kind,
            "old_status": old_status,
            "new_status": old_status,
            "old_confidence": old_conf,
            "new_confidence": old_conf,
            "days_since": 0.0,
            "lambda": DECAY_RATES.get(kind, 0.0),
            "status_changed": False,
            "skipped_reason": None,
        }

        # Skip tombstoned entirely (no decay, no transitions).
        if old_status == "tombstoned":
            change["skipped_reason"] = "tombstoned"
            skipped_tombstoned += 1
            after_status[old_status] += 1
            changes.append(change)
            continue

        # Unknown kind -> no decay rate -> skip.
        if kind not in DECAY_RATES:
            change["skipped_reason"] = f"unknown_kind:{kind}"
            edge_cases.append(f"unknown kind {kind!r} for {urn}")
            after_status[old_status] += 1
            changes.append(change)
            continue

        # Reference time: prefer time.modified, fall back to time.created.
        ref_time = get_modified_at(rec) or get_created_at(rec)
        if ref_time is None:
            change["skipped_reason"] = "no_time_reference"
            edge_cases.append(f"no time reference for {urn}")
            after_status[old_status] += 1
            changes.append(change)
            continue

        days = max((now - ref_time).total_seconds() / 86400.0, 0.0)
        change["days_since"] = days

        lam = DECAY_RATES[kind]
        new_conf = compute_new_confidence(old=old_conf, lam=lam, days_since=days)
        change["new_confidence"] = new_conf

        # Mutate record (or skip if dry_run on the math side, but dry_run already
        # copied records above; we mutate the copy).
        rec["lifecycle"]["confidence"] = new_conf

        # Status transitions.
        new_status = old_status
        if old_status == "candidate" and new_conf >= PROMOTE_THRESHOLD:
            new_status = "active"
            change["new_status"] = "active"
            change["status_changed"] = True
        elif old_status == "active" and new_conf < ARCHIVE_THRESHOLD:
            new_status = "archived"
            change["new_status"] = "archived"
            change["status_changed"] = True

        if new_status != old_status:
            rec["lifecycle"]["status"] = new_status

        after_status[new_status] += 1
        changes.append(change)

        # Track top candidates.
        if old_status == "active" and new_status == "archived":
            archive_candidates.append({
                "urn": urn,
                "kind": kind,
                "old_confidence": old_conf,
                "new_confidence": new_conf,
                "days_since": days,
                "lambda": lam,
                "reason": f"λ={lam} × {days:.1f}d drops conf {old_conf:.3f}→{new_conf:.3f}",
            })
        elif old_status == "candidate" and new_status == "active":
            promotion_candidates.append({
                "urn": urn,
                "kind": kind,
                "old_confidence": old_conf,
                "new_confidence": new_conf,
                "days_since": days,
                "lambda": lam,
            })

    report = {
        "total": len(records),
        "before": before_status,
        "after": after_status,
        "changes": changes,
        "archive_candidates": archive_candidates,
        "promotion_candidates": promotion_candidates,
        "skipped_tombstoned": skipped_tombstoned,
        "edge_cases": edge_cases,
    }
    return records, report


# ---- I/O ------------------------------------------------------------------

def read_records(path: Path) -> list[dict]:
    """Read records from memory.ump.json (JSON array). Robust to mid-write."""
    try:
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
    except json.JSONDecodeError as e:
        raise ValueError(f"{path}: invalid JSON: {e}") from e
    if not isinstance(data, list):
        raise ValueError(f"{path}: top-level JSON is not an array (got {type(data).__name__})")
    return data


def write_records(path: Path, records: list[dict]) -> None:
    """Atomic write: tmp file in same dir, fsync, rename. .tmp removed by os.replace."""
    tmp = path.with_suffix(path.suffix + ".tmp")
    try:
        with open(tmp, "w", encoding="utf-8") as f:
            json.dump(records, f, ensure_ascii=False, indent=2)
            f.flush()
            os.fsync(f.fileno())
        os.replace(tmp, path)
    except Exception:
        # Clean up the half-written tmp on failure.
        if tmp.exists():
            try:
                tmp.unlink()
            except OSError:
                pass
        raise


def backup_file(path: Path) -> Path:
    """Copy file to a timestamped backup alongside the original."""
    ts = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H%M%SZ")
    backup = path.with_suffix(path.suffix + f".bak.{ts}")
    shutil.copy2(path, backup)
    return backup


# ---- Reporting ------------------------------------------------------------

def bucket_for(c: float) -> str:
    if c < 0.20:
        return "[0.00-0.20)"
    if c < 0.40:
        return "[0.20-0.40)"
    if c < 0.60:
        return "[0.40-0.60)"
    if c < 0.80:
        return "[0.60-0.80)"
    return "[0.80-1.00]"


def print_report(report: dict, *, archive_summary: bool = False) -> None:
    total = report["total"]
    before = report["before"]
    after = report["after"]
    changes = report["changes"]

    print("\n=== Decay Report ===")
    print(f"Total records:    {total}")
    print(f"Skipped tombstone: {report['skipped_tombstoned']}")

    print("\nStatus distribution (before -> after):")
    statuses = ["active", "candidate", "archived", "tombstoned", "unknown"]
    for s in statuses:
        b = before.get(s, 0)
        a = after.get(s, 0)
        delta = a - b
        sign = "+" if delta > 0 else ""
        print(f"  {s:11s}: {b:4d} -> {a:4d}  ({sign}{delta:+d})")

    # New-confidence histogram (only for records that had decay applied).
    decayed = [c for c in changes if c["skipped_reason"] is None]
    buckets = Counter(bucket_for(c["new_confidence"]) for c in decayed)
    print("\nNew-confidence histogram (decayed records only):")
    for b in ["[0.00-0.20)", "[0.20-0.40)", "[0.40-0.60)", "[0.60-0.80)", "[0.80-1.00]"]:
        print(f"  {b}: {buckets.get(b, 0)}")

    # Top 10 archive candidates.
    archives = sorted(
        report["archive_candidates"],
        key=lambda c: (c["old_confidence"] - c["new_confidence"]),
        reverse=True,
    )
    print(f"\nTop 10 active -> archived ({len(archives)} total):")
    for c in archives[:10]:
        print(f"  {c['urn'][:60]:60s} kind={c['kind']:9s} "
              f"conf {c['old_confidence']:.3f}->{c['new_confidence']:.3f} "
              f"days={c['days_since']:.0f}")

    # Top 10 promotion candidates.
    promotions = sorted(
        report["promotion_candidates"],
        key=lambda c: c["new_confidence"],
        reverse=True,
    )
    print(f"\nTop 10 candidate -> active ({len(promotions)} total):")
    for c in promotions[:10]:
        print(f"  {c['urn'][:60]:60s} kind={c['kind']:9s} "
              f"conf {c['old_confidence']:.3f}->{c['new_confidence']:.3f} "
              f"days={c['days_since']:.0f}")

    if archive_summary:
        print("\n=== Archive Candidates (with reason) ===")
        for c in archives[:10]:
            print(f"  {c['urn']}")
            print(f"    {c['reason']}")

    # Edge cases.
    if report["edge_cases"]:
        print(f"\nEdge cases ({len(report['edge_cases'])}):")
        for e in report["edge_cases"][:20]:
            print(f"  - {e}")
        if len(report["edge_cases"]) > 20:
            print(f"  ... ({len(report['edge_cases']) - 20} more)")


# ---- Main -----------------------------------------------------------------

def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description="Apply UMP memory confidence decay.")
    g = p.add_mutually_exclusive_group(required=True)
    g.add_argument("--dry-run", action="store_true",
                   help="Show what would change, don't write.")
    g.add_argument("--apply", action="store_true",
                   help="Actually modify memory.ump.json (with backup).")
    p.add_argument("--path", type=Path, default=DEFAULT_PATH,
                   help="Path to memory.ump.json.")
    p.add_argument("--archive-summary", action="store_true",
                   help="Print archive candidates with reasons.")
    args = p.parse_args(argv)

    if not args.path.exists():
        print(f"ERROR: memory file not found: {args.path}", file=sys.stderr)
        return 1

    print(f"Reading: {args.path}")
    try:
        records = read_records(args.path)
    except ValueError as e:
        print(f"ERROR parsing memory file: {e}", file=sys.stderr)
        return 1
    print(f"Loaded {len(records)} records.")

    # Parse errors: count records that aren't dicts before passing through.
    parse_errors = sum(1 for r in records if not isinstance(r, dict))
    if parse_errors:
        print(f"WARN: {parse_errors} non-dict records will be skipped.", file=sys.stderr)

    new_records, report = apply_decay(records, dry_run=True)
    print_report(report, archive_summary=args.archive_summary)

    if args.apply:
        backup = backup_file(args.path)
        print(f"\nBackup created: {backup}")
        write_records(args.path, new_records)
        print(f"Wrote {len(new_records)} records to {args.path}")
    else:
        print("\n(dry-run: no changes written)")

    if parse_errors:
        print(f"\nERROR: {parse_errors} records failed to parse.", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())