#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""Dimension F: Working Memory — LRU+salience per-agent session buffers."""
import hashlib
import json
import shutil
from datetime import datetime, timezone, timedelta
from pathlib import Path

from pulse.core.brain_utils import STATE_DIR, now_iso, atomic_update_json
from pulse.core.brain_io import _acquire

WORKING_MEMORY_DIR = STATE_DIR / "working_memory"
SHARED_CONTEXT_FILE = WORKING_MEMORY_DIR / "shared_context.json"
COMPLETED_DIR = STATE_DIR / "completed"
MAX_ITEMS = 20
MAX_SHARED = 10
REINFORCE_THRESHOLD = 3
_RECENTLY_BROADCAST: dict = {}  # {item_id: timestamp}
_PRECONSCIOUS_BUFFER: list = []  # Item 56: salience 0.3-0.6 near-threshold; 300s TTL
_PHONOLOGICAL_BUFFER: dict = {}; _PHONOLOGICAL_TTL = 300  # Item 66: 5-min phonological loop
_VISUOSPATIAL_SKETCH: list = []; _SKETCH_MAX = 50  # Item 67: spatial sketch
def _mark_broadcast(item_ids: list):
    import time
    now = time.time()
    for iid in item_ids:
        _RECENTLY_BROADCAST[iid] = now
    cutoff = now - 30
    for k in list(_RECENTLY_BROADCAST):
        if _RECENTLY_BROADCAST[k] < cutoff:
            del _RECENTLY_BROADCAST[k]
def _get_ior_penalty(item_id: str) -> float:
    import time
    ts = _RECENTLY_BROADCAST.get(item_id)
    if ts and (time.time() - ts) < 30:
        return 0.15
    return 0.0
def _fingerprint(text: str) -> str:
    """Aligned with brain_search._item_key (80 chars)."""
    return hashlib.md5(text[:80].lower().strip().encode("utf-8")).hexdigest()[:12]
def _add_phonological(text: str):
    """Item 66: phonological loop (5min TTL)."""
    import time
    try:
        fp = _fingerprint(text); ts = time.time()
        _PHONOLOGICAL_BUFFER[fp] = {"text": text[:200], "ts": ts}
        cutoff = ts - _PHONOLOGICAL_TTL
        for k in list(_PHONOLOGICAL_BUFFER):
            if _PHONOLOGICAL_BUFFER[k]["ts"] < cutoff: del _PHONOLOGICAL_BUFFER[k]
    except Exception: pass
def _update_sketch(node_ids: list):
    """Item 67: visuospatial sketchpad (recent graph nodes)."""
    try:
        for nid in node_ids:
            if nid not in _VISUOSPATIAL_SKETCH: _VISUOSPATIAL_SKETCH.append(nid)
        while len(_VISUOSPATIAL_SKETCH) > _SKETCH_MAX: _VISUOSPATIAL_SKETCH.pop(0)
    except Exception: pass
def _session_path(agent: str) -> Path:
    WORKING_MEMORY_DIR.mkdir(parents=True, exist_ok=True)
    return WORKING_MEMORY_DIR / f"{agent}_session.json"
def _new_session(agent, task_id=None):
    now = now_iso()
    return {"session_id": f"{agent}_{now}", "agent": agent, "task_id": task_id,
            "started_at": now, "items": [], "max_items": MAX_ITEMS}
def _load_session(agent: str) -> dict:
    p = _session_path(agent)
    if p.exists():
        try:
            data = json.loads(p.read_text(encoding="utf-8"))
            if isinstance(data, dict) and data.get("agent") == agent:
                return data
        except (json.JSONDecodeError, OSError):
            pass
    return _new_session(agent)
def _save_session(agent: str, data: dict):
    p = _session_path(agent)
    p.parent.mkdir(parents=True, exist_ok=True)
    try:
        with _acquire(p):
            p.write_text(json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
    except Exception:
        pass
def _recency_weight(iso_ts: str) -> float:
    try:
        ts = datetime.fromisoformat(iso_ts)
        if ts.tzinfo is None: ts = ts.replace(tzinfo=timezone.utc)
        return max(0.1, 1.0 - ((datetime.now(timezone.utc) - ts).total_seconds() / 14400) * 0.9)
    except Exception:
        return 0.5
def _chunk_similar(items, threshold=0.8):
    """P3-4: chunk similar items to free working memory slots."""
    try:
        from pulse.core.brain_analysis import text_similarity
    except Exception:
        return items
    chunked, used = [], set()
    for i, item in enumerate(items):
        if i in used: continue
        chunk = [item]
        for j, other in enumerate(items[i + 1:], i + 1):
            if j not in used and text_similarity(item.get("text", ""), other.get("text", "")) > threshold:
                chunk.append(other); used.add(j)
        best = max(chunk, key=lambda x: x.get("salience", 0))
        if len(chunk) > 1: best["_chunked_count"] = len(chunk)
        chunked.append(best); used.add(i)
    return chunked
def _evict(session: dict):
    import time as _t
    _now = _t.time()
    items = session.get("items", [])
    cap = session.get("max_items", MAX_ITEMS)
    locked = [it for it in items if it.get("priority_locked")]
    unlocked = [it for it in items if not it.get("priority_locked")]
    try: unlocked = _chunk_similar(unlocked)
    except Exception: pass
    for it in unlocked:
        it["_es"] = float(it.get("salience", 1.0)) * _recency_weight(
            it.get("last_accessed", it.get("first_accessed", ""))) * (
            1.0 - _get_ior_penalty(it.get("fingerprint", "")))
    unlocked.sort(key=lambda x: x.get("_es", 0), reverse=True)
    keep_unlocked = max(0, cap - len(locked))
    # Item 56: promote pre-conscious items that rose above 0.6 OR re-appeared in session
    try:
        global _PRECONSCIOUS_BUFFER
        _PRECONSCIOUS_BUFFER = [b for b in _PRECONSCIOUS_BUFFER if _now - b.get('_buf_ts', 0) < 300]
        _sfps = {it.get('fingerprint'): it for it in unlocked}
        for _bi in list(_PRECONSCIOUS_BUFFER):
            _fp = _bi.get('fingerprint')
            _promote = (_sfps.get(_fp) if _fp and _fp in _sfps else None) or (_bi if _bi.get('salience', 0) > 0.6 else None)
            if _promote: _PRECONSCIOUS_BUFFER.remove(_bi); unlocked.insert(0, _promote)
    except Exception: pass
    kept = locked + unlocked[:keep_unlocked]
    try:  # Item 56: buffer near-threshold evicted items
        for _ev in unlocked[keep_unlocked:]:
            if 0.3 <= _ev.get("salience", 0) <= 0.6:
                _ev["_buf_ts"] = _now; _PRECONSCIOUS_BUFFER.append(_ev)
        _PRECONSCIOUS_BUFFER = _PRECONSCIOUS_BUFFER[-20:]
    except Exception: pass
    for i, it in enumerate(kept):
        if i < 4:
            it["slot_type"] = "focal"
        elif i < 9:
            it["slot_type"] = "activated"
        else:
            it["slot_type"] = "background"
        it.pop("_es", None)
    _evicted_count = len(items) - len(kept)
    if _evicted_count > 0:
        try:
            from pulse.brain.working_memory_p36 import increment_refresh_count
            increment_refresh_count(_evicted_count)
        except Exception: pass
    session["items"] = kept
def get_active_contexts(agent: str) -> list:
    """Returns focal items (Cowan 2001: 4±1) plus priority-locked items to feed Global Workspace."""
    session = _load_session(agent)
    items = session.get("items", [])
    return [it for it in items if it.get("slot_type") == "focal" or it.get("priority_locked")][:4]
def _load_shared() -> dict:
    WORKING_MEMORY_DIR.mkdir(parents=True, exist_ok=True)
    if SHARED_CONTEXT_FILE.exists():
        try:
            data = json.loads(SHARED_CONTEXT_FILE.read_text(encoding="utf-8"))
            if isinstance(data, dict):
                return data
        except (json.JSONDecodeError, OSError):
            pass
    return {"shared_items": [], "max_shared": MAX_SHARED}
def _save_shared(data: dict):
    try:
        with _acquire(SHARED_CONTEXT_FILE):
            SHARED_CONTEXT_FILE.write_text(
                json.dumps(data, indent=2, ensure_ascii=False), encoding="utf-8")
    except Exception:
        pass
def _evict_ttl(data: dict):
    now = datetime.now(timezone.utc)
    kept = []
    for item in data.get("shared_items", []):
        try:
            posted = datetime.fromisoformat(item["posted_at"])
            if posted.tzinfo is None:
                posted = posted.replace(tzinfo=timezone.utc)
            if (now - posted) <= timedelta(minutes=item.get("ttl_minutes", 60)):
                kept.append(item)
        except Exception:
            continue
    data["shared_items"] = kept
def _add_to_session(agent: str, items: list):
    if not items:
        return
    try:
        session = _load_session(agent)
        now = now_iso()
        existing = {it.get("fingerprint"): it for it in session.get("items", [])}
        for hit in items:
            text = hit.get("text", "") or hit.get("title", "")
            if not text:
                continue
            fp = _fingerprint(text)
            if fp in existing:
                existing[fp]["accessed_count"] = existing[fp].get("accessed_count", 1) + 1
                existing[fp]["last_accessed"] = now
                existing[fp]["salience"] = max(
                    existing[fp].get("salience", 1.0), hit.get("salience", 1.0))
            else:
                existing[fp] = {
                    "text": text[:200], "source": hit.get("source", "unknown"),
                    "fingerprint": fp, "salience": hit.get("salience", 1.0),
                    "accessed_count": 1, "first_accessed": now, "last_accessed": now,
                }
        session["items"] = list(existing.values())
        _evict(session)
        _save_session(agent, session)
    except Exception:
        pass
# === PUBLIC API ===
def get_session(agent: str, task_id: str = None) -> dict:
    """Get or create the current working memory session for an agent."""
    try:
        session = _load_session(agent)
        if task_id and session.get("task_id") != task_id:
            session = _new_session(agent, task_id)
            _save_session(agent, session)
        return session
    except Exception:
        return _new_session(agent, task_id)
def rotate_session(agent: str) -> dict:
    """Close current session: reinforce LTM for frequent items, archive."""
    try:
        session = _load_session(agent)
        items = session.get("items", [])
        now = now_iso()
        reinforced = []
        for it in items:
            if it.get("accessed_count", 0) >= REINFORCE_THRESHOLD:
                try:
                    from pulse.core.uncertainty import bayesian_update
                    from pulse.core.temporal import reinforce
                    ec = it.get("evidence_count", 1)
                    new_conf, new_ec = bayesian_update(
                        it.get("confidence", 0.5), ec, observation_positive=True)
                    it["confidence"], it["evidence_count"] = new_conf, new_ec
                    reinforce(it)
                except Exception:
                    pass
                reinforced.append(it["text"][:60])
        sorted_items = sorted(items, key=lambda x: x.get("accessed_count", 0), reverse=True)
        top_5 = [it["text"][:80] for it in sorted_items[:5]]
        started = session.get("started_at", now)
        try: dur_min = int((datetime.fromisoformat(now) - datetime.fromisoformat(started)).total_seconds() / 60)
        except Exception: dur_min = 0
        summary = {
            "session_id": session.get("session_id", ""), "agent": agent,
            "task_id": session.get("task_id"), "started_at": started, "ended_at": now,
            "duration_minutes": dur_min, "total_items": len(items),
            "reinforced_count": len(reinforced), "top_items": top_5,
        }
        try:
            from pulse.core.brain_utils import EPISODIC_INDEX_FILE
            def _updater(data):
                if not isinstance(data, list):
                    data = []
                data.append({
                    "type": "session_summary",
                    "summary": f"Session {agent}: {len(items)} items, {len(reinforced)} reinforced",
                    "agent": agent, "timestamp": now,
                    "topic_keywords": [it.get("source", "") for it in sorted_items[:5]],
                    "entities_mentioned": [], "emotional_intensity": 0,
                    "_fp": _fingerprint(summary["session_id"]),
                })
                return data[-2000:]
            atomic_update_json(EPISODIC_INDEX_FILE, _updater, default=[])
        except Exception:
            pass
        try:
            src = _session_path(agent)
            if src.exists():
                COMPLETED_DIR.mkdir(parents=True, exist_ok=True)
                ts_safe = now.replace(":", "-").replace("+", "_")[:19]
                shutil.move(str(src), str(COMPLETED_DIR / f"{agent}_{ts_safe}.json"))
        except Exception:
            pass
        _save_session(agent, _new_session(agent))
        return summary
    except Exception as e:
        return {"error": str(e), "agent": agent}
def post_shared(agent: str, text: str, priority: str = "normal", ttl_minutes: int = 60) -> dict:
    """Post to cross-agent shared context. TTL-based eviction, max 10 items."""
    try:
        data = _load_shared()
        _evict_ttl(data)
        now = now_iso()
        item = {"text": text[:300], "posted_by": agent, "posted_at": now,
                "priority": priority, "ttl_minutes": ttl_minutes, "acknowledged_by": []}
        data["shared_items"].append(item)
        prio_map = {"high": 3, "normal": 2, "low": 1}
        data["shared_items"].sort(
            key=lambda x: prio_map.get(x.get("priority", "normal"), 2), reverse=True)
        data["shared_items"] = data["shared_items"][:data.get("max_shared", MAX_SHARED)]
        _save_shared(data)
        return {"status": "posted", "item": item}
    except Exception as e:
        return {"status": "error", "error": str(e)}
