phase-memory/src/phase_memory/management.py

"""Federated store discovery and activity reporting."""

from __future__ import annotations

import json
import os
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Mapping

from .adapters import LOCAL_STORE_SCHEMA, FileBackedMemoryGraphStore, JsonlAuditSink, JsonlMemoryEventLog
from .models import Diagnostic
from .ops_warden import OPS_WARDEN_PROFILE_ID, OPS_WARDEN_RUNTIME_SCHEMA, OpsWardenMemoryStore, default_memory_store_path
from .utils import parse_iso_datetime, stable_digest, utc_now_iso

STORE_REGISTRY_SCHEMA = "phase_memory.management.store_registry.v1"
STORE_LIST_SCHEMA = "phase_memory.management.store_list.v1"
FEDERATED_REPORT_SCHEMA = "phase_memory.management.federated_report.v1"

STORE_KIND_OPS_WARDEN = "ops_warden_coordination"
STORE_KIND_LOCAL_GRAPH = "local_graph"

DEFAULT_OPS_WARDEN_STORE_ID = "ops-warden-default"
TOP_ROUTE_IDS = 10
DETAIL_EPISODE_LIMIT = 20


@dataclass(frozen=True)
class MemoryStoreDescriptor:
    store_id: str
    path: Path
    store_kind: str
    profile_id: str = ""
    label: str = ""
    source: str = "discovered"


def default_registry_path(environ: Mapping[str, str] | None = None) -> Path:
    environ = environ or os.environ
    override = str(environ.get("PHASE_MEMORY_REGISTRY") or "").strip()
    if override:
        return Path(override).expanduser()
    xdg_data = str(environ.get("XDG_DATA_HOME") or "").strip()
    base = Path(xdg_data).expanduser() if xdg_data else Path.home() / ".local" / "share"
    return base / "phase-memory" / "stores.json"


def classify_store(path: str | Path) -> str | None:
    root = Path(path).expanduser().resolve()
    local_metadata = root / "phase-memory.json"
    if local_metadata.is_file():
        try:
            data = json.loads(local_metadata.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            return None
        if str(data.get("schema_version") or "") == LOCAL_STORE_SCHEMA:
            return STORE_KIND_LOCAL_GRAPH

    ops_metadata = root / "metadata.json"
    if ops_metadata.is_file():
        try:
            data = json.loads(ops_metadata.read_text(encoding="utf-8"))
        except json.JSONDecodeError:
            return None
        if str(data.get("schema_version") or "") == OPS_WARDEN_RUNTIME_SCHEMA:
            return STORE_KIND_OPS_WARDEN
    return None


def load_store_registry(
    path: str | Path | None = None,
    environ: Mapping[str, str] | None = None,
) -> dict[str, Any]:
    registry_path = Path(path).expanduser() if path is not None else default_registry_path(environ)
    diagnostics: list[dict[str, Any]] = []
    if not registry_path.is_file():
        return {
            "schema_version": STORE_REGISTRY_SCHEMA,
            "valid": True,
            "registry_path": str(registry_path),
            "entries": [],
            "diagnostics": diagnostics,
        }
    try:
        payload = json.loads(registry_path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        return {
            "schema_version": STORE_REGISTRY_SCHEMA,
            "valid": False,
            "registry_path": str(registry_path),
            "entries": [],
            "diagnostics": [
                Diagnostic(
                    "warn",
                    "corrupt_store_registry",
                    "Store registry file is not valid JSON.",
                    str(registry_path),
                    {"error": str(exc)},
                ).to_dict()
            ],
        }
    if not isinstance(payload, dict):
        diagnostics.append(
            Diagnostic(
                "warn",
                "invalid_store_registry_shape",
                "Store registry must be a JSON object.",
                str(registry_path),
            ).to_dict()
        )
        return {
            "schema_version": STORE_REGISTRY_SCHEMA,
            "valid": False,
            "registry_path": str(registry_path),
            "entries": [],
            "diagnostics": diagnostics,
        }

    schema = str(payload.get("schema_version") or "")
    if schema and schema != STORE_REGISTRY_SCHEMA:
        diagnostics.append(
            Diagnostic(
                "warn",
                "unknown_store_registry_schema",
                "Store registry declares an unknown schema version.",
                str(registry_path),
                {"schema_version": schema},
            ).to_dict()
        )

    entries: list[dict[str, Any]] = []
    for item in payload.get("stores") or ():
        if not isinstance(item, dict):
            continue
        store_path = str(item.get("path") or "").strip()
        store_kind = str(item.get("store_kind") or "").strip()
        store_id = str(item.get("store_id") or "").strip()
        if not store_path or not store_kind or not store_id:
            diagnostics.append(
                Diagnostic(
                    "warn",
                    "incomplete_registry_entry",
                    "Registry entry must include store_id, path, and store_kind.",
                    store_id or store_path or "unknown",
                ).to_dict()
            )
            continue
        entries.append(item)
    return {
        "schema_version": STORE_REGISTRY_SCHEMA,
        "valid": not any(item.get("severity") == "error" for item in diagnostics),
        "registry_path": str(registry_path),
        "entries": entries,
        "diagnostics": diagnostics,
    }


def _store_id_for_path(path: Path, environ: Mapping[str, str] | None = None) -> str:
    resolved = path.expanduser().resolve()
    default_ops = default_memory_store_path(environ).expanduser().resolve()
    if resolved == default_ops:
        return DEFAULT_OPS_WARDEN_STORE_ID
    return f"store:{stable_digest(str(resolved))}"


def _descriptor_from_path(
    path: Path,
    *,
    store_kind: str,
    source: str,
    profile_id: str = "",
    label: str = "",
    store_id: str = "",
    environ: Mapping[str, str] | None = None,
) -> MemoryStoreDescriptor:
    resolved = path.expanduser().resolve()
    return MemoryStoreDescriptor(
        store_id=store_id or _store_id_for_path(resolved, environ),
        path=resolved,
        store_kind=store_kind,
        profile_id=profile_id,
        label=label,
        source=source,
    )


def discover_memory_stores(
    environ: Mapping[str, str] | None = None,
    *,
    registry_path: str | Path | None = None,
    extra_paths: tuple[str | Path, ...] = (),
) -> tuple[list[MemoryStoreDescriptor], list[dict[str, Any]]]:
    environ = environ or os.environ
    diagnostics: list[dict[str, Any]] = []
    by_path: dict[str, MemoryStoreDescriptor] = {}

    def add_descriptor(descriptor: MemoryStoreDescriptor) -> None:
        key = str(descriptor.path)
        existing = by_path.get(key)
        if existing is None:
            by_path[key] = descriptor
            return
        if descriptor.source == "registry" and existing.source != "registry":
            by_path[key] = descriptor

    default_ops = default_memory_store_path(environ)
    kind = classify_store(default_ops)
    if kind:
        profile_id = OPS_WARDEN_PROFILE_ID if kind == STORE_KIND_OPS_WARDEN else ""
        add_descriptor(
            _descriptor_from_path(
                default_ops,
                store_kind=kind,
                source="default",
                profile_id=profile_id,
                label="ops-warden coordination store",
                store_id=DEFAULT_OPS_WARDEN_STORE_ID,
                environ=environ,
            )
        )

    registry = load_store_registry(registry_path, environ=environ)
    diagnostics.extend(registry.get("diagnostics", ()))
    registry_file = Path(str(registry.get("registry_path") or "")).expanduser()
    for entry in registry.get("entries", ()):
        entry_path = Path(str(entry.get("path") or "")).expanduser()
        if not entry_path.is_absolute() and registry_file.is_file():
            entry_path = (registry_file.parent / entry_path).resolve()
        if not entry_path.is_dir():
            diagnostics.append(
                Diagnostic(
                    "warn",
                    "registry_store_missing",
                    "Registry entry path does not exist.",
                    str(entry.get("store_id") or entry_path),
                    {"path": str(entry_path)},
                ).to_dict()
            )
            continue
        declared_kind = str(entry.get("store_kind") or "")
        actual_kind = classify_store(entry_path)
        if actual_kind and actual_kind != declared_kind:
            diagnostics.append(
                Diagnostic(
                    "warn",
                    "registry_store_kind_mismatch",
                    "Registry store_kind does not match on-disk classification.",
                    str(entry.get("store_id") or entry_path),
                    {"declared": declared_kind, "actual": actual_kind},
                ).to_dict()
            )
        store_kind = actual_kind or declared_kind
        if not store_kind:
            continue
        add_descriptor(
            _descriptor_from_path(
                entry_path,
                store_kind=store_kind,
                source="registry",
                profile_id=str(entry.get("profile_id") or ""),
                label=str(entry.get("label") or ""),
                store_id=str(entry.get("store_id") or ""),
                environ=environ,
            )
        )

    env_paths = str(environ.get("PHASE_MEMORY_STORE_PATHS") or "").strip()
    path_candidates = [item.strip() for item in env_paths.split(":") if item.strip()]
    path_candidates.extend(str(item) for item in extra_paths)
    for candidate in path_candidates:
        candidate_path = Path(candidate).expanduser()
        if not candidate_path.is_dir():
            continue
        store_kind = classify_store(candidate_path)
        if not store_kind:
            continue
        add_descriptor(
            _descriptor_from_path(
                candidate_path,
                store_kind=store_kind,
                source="env",
                environ=environ,
            )
        )

    stores = sorted(by_path.values(), key=lambda item: (item.store_kind, item.store_id))
    return stores, diagnostics


def resolve_store_reference(
    reference: str,
    stores: list[MemoryStoreDescriptor],
    environ: Mapping[str, str] | None = None,
) -> MemoryStoreDescriptor | None:
    ref = str(reference or "").strip()
    if not ref:
        return None
    for store in stores:
        if ref == store.store_id or ref == str(store.path):
            return store
    candidate = Path(ref).expanduser()
    if candidate.is_dir():
        store_kind = classify_store(candidate)
        if store_kind:
            return _descriptor_from_path(candidate, store_kind=store_kind, source="explicit", environ=environ)
    for store in stores:
        if candidate.resolve() == store.path:
            return store
    return None


def compute_report_window(
    *,
    days: int,
    window_end: datetime | None = None,
) -> tuple[datetime, datetime]:
    if days < 1:
        raise ValueError("days must be at least 1")
    end = window_end or datetime.now(timezone.utc)
    if end.tzinfo is None:
        end = end.replace(tzinfo=timezone.utc)
    else:
        end = end.astimezone(timezone.utc)
    start = end - timedelta(days=days)
    return start, end


def _event_in_window(timestamp: str | None, *, window_start: datetime, window_end: datetime) -> bool | None:
    parsed = parse_iso_datetime(timestamp)
    if parsed is None:
        return None
    return window_start <= parsed <= window_end


def _read_jsonl(path: Path) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    if not path.is_file():
        return [], []
    events: list[dict[str, Any]] = []
    diagnostics: list[dict[str, Any]] = []
    for line_number, raw in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
        if not raw.strip():
            continue
        try:
            data = json.loads(raw)
        except json.JSONDecodeError as exc:
            diagnostics.append(
                Diagnostic(
                    "warn",
                    "malformed_jsonl_line",
                    "JSONL line is not valid JSON.",
                    f"{path}:line:{line_number}",
                    {"error": str(exc)},
                ).to_dict()
            )
            continue
        if isinstance(data, dict):
            events.append(data)
    return events, diagnostics


def _increment(counter: dict[str, int], key: str) -> None:
    normalized = key or "unknown"
    counter[normalized] = counter.get(normalized, 0) + 1


def _top_counter(counter: dict[str, int], *, limit: int) -> dict[str, int]:
    return dict(sorted(counter.items(), key=lambda item: (-item[1], item[0]))[:limit])


def aggregate_ops_warden_activity(
    store: MemoryStoreDescriptor,
    *,
    window_start: datetime,
    window_end: datetime,
) -> dict[str, Any]:
    events_path = store.path / "events.jsonl"
    raw_events, line_diagnostics = _read_jsonl(events_path)
    diagnostics = list(line_diagnostics)
    by_session_kind: dict[str, int] = {}
    by_command: dict[str, int] = {}
    by_outcome: dict[str, int] = {}
    by_route_id: dict[str, int] = {}
    episode_timeline: list[dict[str, Any]] = []
    first_activity_at = ""
    last_activity_at = ""
    in_window = 0
    missing_timestamp = 0

    for event in raw_events:
        timestamp = str(event.get("recorded_at") or "")
        in_range = _event_in_window(timestamp, window_start=window_start, window_end=window_end)
        if in_range is None:
            missing_timestamp += 1
            diagnostics.append(
                Diagnostic(
                    "warn",
                    "episode_missing_timestamp",
                    "Episode lacks a parseable recorded_at timestamp.",
                    str(event.get("event_id") or "unknown"),
                ).to_dict()
            )
            continue
        if not in_range:
            continue
        in_window += 1
        _increment(by_session_kind, str(event.get("session_kind") or ""))
        _increment(by_command, str(event.get("command") or ""))
        _increment(by_outcome, str(event.get("outcome") or ""))
        _increment(by_route_id, str(event.get("route_id") or ""))
        if not first_activity_at or timestamp < first_activity_at:
            first_activity_at = timestamp
        if not last_activity_at or timestamp > last_activity_at:
            last_activity_at = timestamp
        episode_timeline.append(
            {
                "event_id": str(event.get("event_id") or ""),
                "recorded_at": timestamp,
                "session_kind": str(event.get("session_kind") or ""),
                "command": str(event.get("command") or ""),
                "route_id": str(event.get("route_id") or ""),
                "outcome": str(event.get("outcome") or ""),
            }
        )

    episode_timeline.sort(key=lambda item: item.get("recorded_at", ""), reverse=True)
    metadata_path = store.path / "metadata.json"
    schema_version = ""
    if metadata_path.is_file():
        try:
            schema_version = str(json.loads(metadata_path.read_text(encoding="utf-8")).get("schema_version") or "")
        except json.JSONDecodeError:
            pass

    return {
        "store_id": store.store_id,
        "store_kind": store.store_kind,
        "path": str(store.path),
        "profile_id": store.profile_id or OPS_WARDEN_PROFILE_ID,
        "schema_version": schema_version,
        "episode_count": in_window,
        "audit_event_count": 0,
        "missing_timestamp_count": missing_timestamp,
        "first_activity_at": first_activity_at,
        "last_activity_at": last_activity_at,
        "by_session_kind": dict(sorted(by_session_kind.items())),
        "by_command": dict(sorted(by_command.items())),
        "by_outcome": dict(sorted(by_outcome.items())),
        "by_route_id": _top_counter(by_route_id, limit=TOP_ROUTE_IDS),
        "by_operation": {},
        "episode_timeline": episode_timeline[:DETAIL_EPISODE_LIMIT],
        "repair_diagnostics": [],
        "diagnostics": diagnostics,
    }


def aggregate_local_graph_activity(
    store: MemoryStoreDescriptor,
    *,
    window_start: datetime,
    window_end: datetime,
) -> dict[str, Any]:
    events_path = store.path / "events.jsonl"
    audit_path = store.path / "audit.jsonl"
    raw_events, event_line_diagnostics = _read_jsonl(events_path)
    raw_audit, audit_line_diagnostics = _read_jsonl(audit_path)
    diagnostics = list(event_line_diagnostics)
    diagnostics.extend(audit_line_diagnostics)

    by_session_kind: dict[str, int] = {}
    by_command: dict[str, int] = {}
    by_outcome: dict[str, int] = {}
    by_operation: dict[str, int] = {}
    episode_timeline: list[dict[str, Any]] = []
    first_activity_at = ""
    last_activity_at = ""
    episode_count = 0
    audit_count = 0
    missing_timestamp = 0

    for event in raw_events:
        timestamp = str(event.get("timestamp") or event.get("recorded_at") or "")
        in_range = _event_in_window(timestamp, window_start=window_start, window_end=window_end)
        if in_range is None:
            missing_timestamp += 1
            diagnostics.append(
                Diagnostic(
                    "warn",
                    "episode_missing_timestamp",
                    "Event lacks a parseable timestamp.",
                    str(event.get("event_id") or event.get("id") or "unknown"),
                ).to_dict()
            )
            continue
        if not in_range:
            continue
        episode_count += 1
        kind = str(event.get("kind") or event.get("session_kind") or "")
        _increment(by_session_kind, kind)
        _increment(by_command, str(event.get("command") or kind))
        _increment(by_outcome, str(event.get("outcome") or event.get("state") or ""))
        if not first_activity_at or timestamp < first_activity_at:
            first_activity_at = timestamp
        if not last_activity_at or timestamp > last_activity_at:
            last_activity_at = timestamp
        episode_timeline.append(
            {
                "event_id": str(event.get("event_id") or event.get("id") or ""),
                "recorded_at": timestamp,
                "session_kind": kind,
                "command": str(event.get("command") or kind),
                "route_id": "",
                "outcome": str(event.get("outcome") or event.get("state") or ""),
            }
        )

    for event in raw_audit:
        timestamp = str(event.get("timestamp") or event.get("recorded_at") or "")
        in_range = _event_in_window(timestamp, window_start=window_start, window_end=window_end)
        if in_range is None:
            missing_timestamp += 1
            diagnostics.append(
                Diagnostic(
                    "warn",
                    "audit_missing_timestamp",
                    "Audit event lacks a parseable timestamp.",
                    str(event.get("operation_id") or event.get("operation") or "unknown"),
                ).to_dict()
            )
            continue
        if not in_range:
            continue
        audit_count += 1
        _increment(by_operation, str(event.get("operation") or ""))
        if not first_activity_at or timestamp < first_activity_at:
            first_activity_at = timestamp
        if not last_activity_at or timestamp > last_activity_at:
            last_activity_at = timestamp

    episode_timeline.sort(key=lambda item: item.get("recorded_at", ""), reverse=True)
    graph_store = FileBackedMemoryGraphStore(store.path)
    repair = graph_store.repair_diagnostics()
    repair_diagnostics = [
        item.to_dict()
        for item in repair
        if item.severity in {"error", "warn"}
    ]
    schema_version = ""
    try:
        schema_version = str(graph_store.metadata().get("schema_version") or "")
    except (json.JSONDecodeError, OSError):
        pass

    return {
        "store_id": store.store_id,
        "store_kind": store.store_kind,
        "path": str(store.path),
        "profile_id": store.profile_id,
        "schema_version": schema_version,
        "episode_count": episode_count,
        "audit_event_count": audit_count,
        "missing_timestamp_count": missing_timestamp,
        "first_activity_at": first_activity_at,
        "last_activity_at": last_activity_at,
        "by_session_kind": dict(sorted(by_session_kind.items())),
        "by_command": dict(sorted(by_command.items())),
        "by_outcome": dict(sorted(by_outcome.items())),
        "by_route_id": {},
        "by_operation": dict(sorted(by_operation.items())),
        "episode_timeline": episode_timeline[:DETAIL_EPISODE_LIMIT],
        "repair_diagnostics": repair_diagnostics,
        "diagnostics": diagnostics,
    }


def aggregate_store_activity(
    store: MemoryStoreDescriptor,
    *,
    window_start: datetime,
    window_end: datetime,
) -> dict[str, Any]:
    if store.store_kind == STORE_KIND_OPS_WARDEN:
        return aggregate_ops_warden_activity(store, window_start=window_start, window_end=window_end)
    return aggregate_local_graph_activity(store, window_start=window_start, window_end=window_end)


def build_store_list(
    environ: Mapping[str, str] | None = None,
    *,
    registry_path: str | Path | None = None,
    extra_paths: tuple[str | Path, ...] = (),
) -> dict[str, Any]:
    stores, diagnostics = discover_memory_stores(environ, registry_path=registry_path, extra_paths=extra_paths)
    return {
        "schema_version": STORE_LIST_SCHEMA,
        "valid": not any(item.get("severity") == "error" for item in diagnostics),
        "store_count": len(stores),
        "stores": [
            {
                "store_id": store.store_id,
                "path": str(store.path),
                "store_kind": store.store_kind,
                "profile_id": store.profile_id,
                "label": store.label,
                "source": store.source,
            }
            for store in stores
        ],
        "diagnostics": diagnostics,
    }


def build_federated_report(
    *,
    days: int = 7,
    focus_store_id: str | None = None,
    focus_store_path: str | Path | None = None,
    environ: Mapping[str, str] | None = None,
    registry_path: str | Path | None = None,
    extra_paths: tuple[str | Path, ...] = (),
    window_end: datetime | None = None,
) -> dict[str, Any]:
    environ = environ or os.environ
    stores, discovery_diagnostics = discover_memory_stores(
        environ,
        registry_path=registry_path,
        extra_paths=extra_paths,
    )
    diagnostics = list(discovery_diagnostics)

    focus_reference = str(focus_store_id or focus_store_path or "").strip()
    focus_store: MemoryStoreDescriptor | None = None
    if focus_reference:
        focus_store = resolve_store_reference(focus_reference, stores, environ=environ)
        if focus_store is None:
            diagnostics.append(
                Diagnostic(
                    "error",
                    "unknown_store_reference",
                    "Requested store was not discovered and could not be classified.",
                    focus_reference,
                ).to_dict()
            )
            return {
                "schema_version": FEDERATED_REPORT_SCHEMA,
                "valid": False,
                "window_days": days,
                "window_start": "",
                "window_end": "",
                "focus_store_id": focus_reference,
                "store_count": 0,
                "aggregate": {},
                "stores": [],
                "store_detail": None,
                "diagnostics": diagnostics,
            }
        stores = [focus_store]

    window_start, window_end_dt = compute_report_window(days=days, window_end=window_end)
    summaries: list[dict[str, Any]] = []
    aggregate = {
        "episode_count": 0,
        "audit_event_count": 0,
        "active_store_count": 0,
        "by_outcome": {},
        "by_session_kind": {},
        "by_operation": {},
    }
    by_outcome: dict[str, int] = {}
    by_session_kind: dict[str, int] = {}
    by_operation: dict[str, int] = {}

    store_detail: dict[str, Any] | None = None
    for store in stores:
        summary = aggregate_store_activity(store, window_start=window_start, window_end=window_end_dt)
        diagnostics.extend(summary.get("diagnostics", ()))
        public_summary = {
            key: value
            for key, value in summary.items()
            if key not in {"episode_timeline", "repair_diagnostics", "diagnostics"}
        }
        summaries.append(public_summary)
        aggregate["episode_count"] += int(summary.get("episode_count") or 0)
        aggregate["audit_event_count"] += int(summary.get("audit_event_count") or 0)
        if int(summary.get("episode_count") or 0) or int(summary.get("audit_event_count") or 0):
            aggregate["active_store_count"] += 1
        for key, value in (summary.get("by_outcome") or {}).items():
            by_outcome[key] = by_outcome.get(key, 0) + int(value)
        for key, value in (summary.get("by_session_kind") or {}).items():
            by_session_kind[key] = by_session_kind.get(key, 0) + int(value)
        for key, value in (summary.get("by_operation") or {}).items():
            by_operation[key] = by_operation.get(key, 0) + int(value)
        if focus_store is not None and store.store_id == focus_store.store_id:
            store_detail = summary

    aggregate["by_outcome"] = dict(sorted(by_outcome.items()))
    aggregate["by_session_kind"] = dict(sorted(by_session_kind.items()))
    aggregate["by_operation"] = dict(sorted(by_operation.items()))

    return {
        "schema_version": FEDERATED_REPORT_SCHEMA,
        "valid": not any(item.get("severity") == "error" for item in diagnostics),
        "window_days": days,
        "window_start": window_start.replace(microsecond=0).isoformat(),
        "window_end": window_end_dt.replace(microsecond=0).isoformat(),
        "focus_store_id": focus_store.store_id if focus_store else "",
        "store_count": len(summaries),
        "aggregate": aggregate,
        "stores": summaries,
        "store_detail": store_detail,
        "diagnostics": diagnostics,
        "generated_at": utc_now_iso(),
    }