#!/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.

Split from the original combined file for Law 7 compliance.
Cross-agent shared broadcast (post_shared, get_shared, _load_shared,
_save_shared, _evict_ttl) lives in working_memory_shared.py and is
re-exported here so all existing callers keep working unchanged.

LAW 7 NOTE: 331 lines — justified keep. Single responsibility: per-agent
session lifecycle (16 tightly-coupled functions: load/save/evict/chunk/
fingerprint/phonological-loop/visuospatial-sketch/rotate). Splitting further
would require passing session state across module boundaries on every call,
adding latency to the hot search path.
"""
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"
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.memory.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 _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}


# ---------------------------------------------------------------------------
# Re-export shared broadcast surface so all existing callers keep working
# with `from pulse.brain.memory.working_memory import post_shared, _load_shared, ...`
# ---------------------------------------------------------------------------
from pulse.brain.memory.working_memory_shared import (  # noqa: E402, F401
    SHARED_CONTEXT_FILE,
    post_shared,
    get_shared,
    _load_shared,
    _save_shared,
    _evict_ttl,
)
