#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""Dimension F: Working Memory Management.

Per-agent session buffers with LRU+salience eviction, session rotation,
and cross-agent shared context. Public: get_session, rotate_session, post_shared.
"""
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


def _fingerprint(text: str) -> str:
    return hashlib.md5(text[:60].lower().encode()).hexdigest()[:12]


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)
        age_s = (datetime.now(timezone.utc) - ts).total_seconds()
        return max(0.1, 1.0 - (age_s / 14400) * 0.9)  # 4h window
    except Exception:
        return 0.5


def _evict(session: dict):
    items = session.get("items", [])
    cap = session.get("max_items", MAX_ITEMS)
    if len(items) <= cap:
        return
    for it in items:
        it["_es"] = it.get("salience", 1.0) * _recency_weight(
            it.get("last_accessed", it.get("first_accessed", "")))
    items.sort(key=lambda x: x.get("_es", 0), reverse=True)
    session["items"] = items[:cap]
    for it in session["items"]:
        it.pop("_es", None)


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()
        # Reinforce items accessed >= threshold
        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])
        # Summary
        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,
        }
        # Append to episodic index
        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
        # Archive session file
        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)}
