#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""Dimension F: Working Memory — LRU+salience per-agent session buffers.

Modularized: wm_buffer.py (broadcast/IOR/phonological), wm_io.py (session I/O).
All existing callers keep working via re-exports below.
"""
import shutil
from datetime import datetime
from pulse.core.brain_utils import STATE_DIR, now_iso, atomic_update_json
from pulse.brain.memory.wm_buffer import (  # noqa: F401
    _active_agent, _RECENTLY_BROADCAST, _PRECONSCIOUS_BUFFER,
    _PHONOLOGICAL_BUFFER, _VISUOSPATIAL_SKETCH,
    _fingerprint, _mark_broadcast, _get_ior_penalty, _add_phonological, _update_sketch,
)
from pulse.brain.memory.wm_io import (  # noqa: F401
    WORKING_MEMORY_DIR, COMPLETED_DIR, MAX_ITEMS,
    _session_path, _new_session, _load_session, _save_session,
)

MAX_SHARED = 10
REINFORCE_THRESHOLD = 3


def _recency_weight(iso_ts: str) -> float:
    from datetime import timezone
    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))
    try:
        agent = _active_agent.get()
        buf = _PRECONSCIOUS_BUFFER.setdefault(agent, [])
        buf = [b for b in buf if _now - b.get("_buf_ts", 0) < 300]
        _PRECONSCIOUS_BUFFER[agent] = buf
        _sfps = {it.get("fingerprint"): it for it in unlocked}
        for _bi in list(buf):
            _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:
                buf.remove(_bi)
                unlocked.insert(0, _promote)
    except Exception:
        pass
    kept = locked + unlocked[:keep_unlocked]
    try:
        for _ev in unlocked[keep_unlocked:]:
            if 0.3 <= _ev.get("salience", 0) <= 0.6:
                _ev["_buf_ts"] = _now
                buf.append(_ev)
        _PRECONSCIOUS_BUFFER[_active_agent.get()] = buf[-20:]
    except Exception:
        pass
    for i, it in enumerate(kept):
        it["slot_type"] = "focal" if i < 4 else ("activated" if i < 9 else "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."""
    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


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
        _RECENTLY_BROADCAST.pop(agent, None)
        _PRECONSCIOUS_BUFFER.pop(agent, None)
        _PHONOLOGICAL_BUFFER.pop(agent, None)
        _VISUOSPATIAL_SKETCH.pop(agent, None)
        _save_session(agent, _new_session(agent))
        return summary
    except Exception as e:
        return {"error": str(e), "agent": agent}


# Re-export shared broadcast surface for backward compatibility
from pulse.brain.memory.working_memory_shared import (  # noqa: E402, F401
    SHARED_CONTEXT_FILE, post_shared, get_shared, _load_shared, _save_shared, _evict_ttl,
)
