replica-omnisciente/scripts/data_sharing/collector.py

"""
Collector — gathers agent data (tool calls, thinking, metadata) from the
local replica's brain, filtered by the consent record.

Only data categories explicitly consented to in [data_sharing.categories]
are included in the collected payload. Everything else is skipped.
"""

from __future__ import annotations

import json
import platform
import shutil
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from .consent import ConsentRecord, CATEGORIES, redact


def _safe_read(path: Path, max_bytes: int = 512_000) -> str | None:
    """Read a file safely, returning None on error."""
    try:
        if not path.exists() or not path.is_file():
            return None
        if path.stat().st_size > max_bytes:
            return None
        return path.read_text(encoding="utf-8", errors="replace")
    except Exception:
        return None


def _safe_read_json(path: Path) -> dict | list | None:
    """Read and parse JSON safely."""
    raw = _safe_read(path)
    if raw is None:
        return None
    try:
        return json.loads(raw)
    except Exception:
        return None


def _collect_tool_calls(brain_dir: Path, consent: ConsentRecord) -> list[dict]:
    """Extract tool call records from session files."""
    if not consent.allows("tool_calls"):
        return []

    results: list[dict] = []
    for session_dir in sorted(brain_dir.glob("session-*")):
        meta = _safe_read_json(session_dir / "metadata.json")
        if meta and isinstance(meta, dict):
            session_id = meta.get("session_id", session_dir.name)
            workspace = meta.get("workspace_path", "")
        else:
            session_id = session_dir.name
            workspace = ""

        # Session JSONL — tool calls are typically embedded in messages
        session_jsonl = session_dir / "session.jsonl"
        session_raw = _safe_read(session_jsonl, max_bytes=2_000_000) or ""
        if session_raw:
            for line in session_raw.splitlines():
                if not line.strip():
                    continue
                try:
                    entry = json.loads(line)
                except Exception:
                    continue
                # Look for tool_call entries
                if isinstance(entry, dict):
                    role = entry.get("role", "")
                    if role == "tool" or "tool_call" in entry or "toolCallId" in entry:
                        record = {
                            "session_id": session_id,
                            "workspace": workspace,
                            "tool_name": entry.get("tool_name") or entry.get("name", ""),
                            "tool_call_id": entry.get("toolCallId", ""),
                            "arguments": entry.get("arguments") or entry.get("content", ""),
                            "result": entry.get("result", ""),
                            "timestamp": entry.get("timestamp", ""),
                        }
                        if consent.redact_secrets:
                            record["arguments"] = redact(str(record["arguments"]))[:10_000]
                            record["result"] = redact(str(record["result"]))[:10_000]
                        results.append(record)
    return results


def _collect_thinking(brain_dir: Path, consent: ConsentRecord) -> list[dict]:
    """Extract reasoning/thinking traces from session files."""
    if not consent.allows("thinking"):
        return []

    results: list[dict] = []
    for session_dir in sorted(brain_dir.glob("session-*")):
        session_jsonl = session_dir / "session.jsonl"
        session_raw = _safe_read(session_jsonl, max_bytes=2_000_000) or ""
        if session_raw:
            for line in session_raw.splitlines():
                if not line.strip():
                    continue
                try:
                    entry = json.loads(line)
                except Exception:
                    continue
                if isinstance(entry, dict):
                    # Look for thinking/reasoning entries
                    if entry.get("role") == "thinking" or "reasoning" in entry or "thinking" in entry:
                        results.append({
                            "session_id": session_dir.name,
                            "thinking": (entry.get("content") or entry.get("thinking") or "")[:50_000],
                            "timestamp": entry.get("timestamp", ""),
                        })
    return results


def _collect_chat_messages(brain_dir: Path, consent: ConsentRecord) -> list[dict]:
    """Extract user/assistant chat messages."""
    if not consent.allows("chat_messages"):
        return []

    results: list[dict] = []
    for session_dir in sorted(brain_dir.glob("session-*")):
        session_jsonl = session_dir / "session.jsonl"
        raw = _safe_read(session_jsonl, max_bytes=2_000_000)
        if raw:
            for line in raw.splitlines():
                if not line.strip():
                    continue
                try:
                    entry = json.loads(line)
                except Exception:
                    continue
                if isinstance(entry, dict) and entry.get("role") in ("user", "assistant"):
                    content = str(entry.get("content", ""))
                    if consent.redact_secrets:
                        content = redact(content)
                    results.append({
                        "session_id": session_dir.name,
                        "role": entry["role"],
                        "content": content[:50_000],
                        "timestamp": entry.get("timestamp", ""),
                    })

        # Also check chat_history.json
        chat_hist = _safe_read_json(session_dir / ".system_generated" / "chat_history.json")
        if chat_hist and isinstance(chat_hist, list):
            for msg in chat_hist:
                if isinstance(msg, dict) and msg.get("role") in ("user", "assistant"):
                    content = str(msg.get("content", ""))
                    if consent.redact_secrets:
                        content = redact(content)
                    results.append({
                        "session_id": session_dir.name,
                        "role": msg["role"],
                        "content": content[:50_000],
                    })
    return results


def _collect_session_meta(brain_dir: Path, consent: ConsentRecord) -> list[dict]:
    """Extract session metadata (ids, timestamps, workspaces)."""
    if not consent.allows("session_meta"):
        return []

    results: list[dict] = []
    for session_dir in sorted(brain_dir.glob("session-*")):
        meta = _safe_read_json(session_dir / "metadata.json")
        if meta and isinstance(meta, dict):
            results.append({
                "session_id": meta.get("session_id", session_dir.name),
                "source": meta.get("source", ""),
                "workspace": meta.get("workspace_path", ""),
                "started_at": meta.get("started_at", ""),
                "ended_at": meta.get("ended_at", ""),
                "message_count": meta.get("message_count", 0),
            })
        else:
            results.append({
                "session_id": session_dir.name,
                "source": "",
                "workspace": "",
            })
    return results


def _collect_agent_metadata(brain_dir: Path, consent: ConsentRecord) -> dict[str, Any]:
    """Extract agent/heteronym metadata and model usage."""
    if not consent.allows("agent_metadata"):
        return {}

    # Aggregate metadata from all sessions
    models_used: dict[str, int] = {}
    heteronyms_used: dict[str, int] = {}
    total_sessions = 0

    for session_dir in sorted(brain_dir.glob("session-*")):
        meta = _safe_read_json(session_dir / "metadata.json")
        if meta and isinstance(meta, dict):
            total_sessions += 1
            model = meta.get("model", "unknown")
            models_used[model] = models_used.get(model, 0) + 1
            heteronym = meta.get("heteronym") or meta.get("agent_name", "default")
            heteronyms_used[heteronym] = heteronyms_used.get(heteronym, 0) + 1

    return {
        "total_sessions": total_sessions,
        "models_used": models_used,
        "heteronyms_used": heteronyms_used,
    }


def _collect_error_traces(brain_dir: Path, consent: ConsentRecord) -> list[dict]:
    """Extract error traces and exceptions from session logs."""
    if not consent.allows("error_traces"):
        return []

    results: list[dict] = []
    for session_dir in sorted(brain_dir.glob("session-*")):
        logs_dir = session_dir / ".system_generated" / "logs"
        if logs_dir.exists():
            for log_file in logs_dir.glob("*.txt"):
                content = _safe_read(log_file)
                if content and ("error" in content.lower() or "traceback" in content.lower() or "exception" in content.lower()):
                    if consent.redact_secrets:
                        content = redact(content)
                    results.append({
                        "session_id": session_dir.name,
                        "log_file": log_file.name,
                        "content": content[:50_000],
                    })
    return results


def _collect_file_changes(repo_root: Path, consent: ConsentRecord) -> list[dict]:
    """Collect recent git diffs."""
    if not consent.allows("file_changes"):
        return []

    if not shutil.which("git"):
        return []

    results: list[dict] = []
    try:
        # Get last 20 commits
        r = subprocess.run(
            ["git", "log", "--oneline", "-20", "--format=%H|%s|%ai"],
            cwd=repo_root, capture_output=True, text=True, timeout=10
        )
        if r.returncode == 0:
            for line in r.stdout.strip().splitlines():
                parts = line.split("|", 2)
                if len(parts) == 3:
                    commit_hash, subject, date = parts
                    # Get diff stat
                    r2 = subprocess.run(
                        ["git", "show", "--stat", "--format=", commit_hash],
                        cwd=repo_root, capture_output=True, text=True, timeout=10
                    )
                    diff_stat = r2.stdout.strip() if r2.returncode == 0 else ""
                    if consent.redact_secrets:
                        diff_stat = redact(diff_stat)
                    results.append({
                        "commit": commit_hash,
                        "subject": subject,
                        "date": date,
                        "diff_stat": diff_stat[:5_000],
                    })
    except Exception:
        pass
    return results


def _collect_environment() -> dict[str, Any]:
    """Collect basic environment telemetry."""
    return {
        "os": platform.system(),
        "os_version": platform.version(),
        "python_version": platform.python_version(),
        "machine": platform.machine(),
        "processor": platform.processor()[:100],
        "hostname": platform.node(),
        "collected_at": datetime.now(timezone.utc).isoformat(),
    }


def collect(
    replica_root: Path,
    consent: ConsentRecord,
) -> dict[str, Any]:
    """
    Collect all consented data from the replica.

    Returns a payload dict ready for transmission. Only categories that are
    explicitly enabled in the consent record are included.

    Args:
        replica_root: Path to the replica-omnisciente root.
        consent: Resolved consent record.

    Returns:
        Payload dict with metadata + per-category data.
    """
    brain_dir = replica_root / ".aurelio" / "brain"

    payload: dict[str, Any] = {
        "schema_version": 1,
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "consent": {
            "categories": consent.granted_categories(),
            "retention_days": consent.retention_days,
            "redact_secrets": consent.redact_secrets,
        },
    }

    # Always include environment if consented
    if consent.allows("environment"):
        payload["environment"] = _collect_environment()

    # Per-category collection
    if consent.allows("tool_calls"):
        payload["tool_calls"] = _collect_tool_calls(brain_dir, consent)

    if consent.allows("thinking"):
        payload["thinking"] = _collect_thinking(brain_dir, consent)

    if consent.allows("chat_messages"):
        payload["chat_messages"] = _collect_chat_messages(brain_dir, consent)

    if consent.allows("session_meta"):
        payload["session_meta"] = _collect_session_meta(brain_dir, consent)

    agent_meta = _collect_agent_metadata(brain_dir, consent)
    if agent_meta:
        payload["agent_metadata"] = agent_meta

    if consent.allows("error_traces"):
        payload["error_traces"] = _collect_error_traces(brain_dir, consent)

    if consent.allows("file_changes"):
        payload["file_changes"] = _collect_file_changes(replica_root, consent)

    # Summary counts
    payload["_summary"] = {
        cat: (len(payload[cat]) if isinstance(payload.get(cat), list) else
              (1 if payload.get(cat) else 0))
        for cat in CATEGORIES if cat in payload
    }

    return payload