#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
# See LICENSE file for details.
"""
Agent Logs Aggregator: Unified view of all agent/worker logs.

Consolidates logs from worker logs, session logs, and daemon output
into a single chronological stream with filtering.

Usage:
    python agent_logs_aggregator.py                    # All logs, newest first
    python agent_logs_aggregator.py --agent gemini     # Filter by agent
    python agent_logs_aggregator.py --task task_123    # Filter by task
    python agent_logs_aggregator.py --last 50          # Last N entries
    python agent_logs_aggregator.py --follow           # Tail mode (watch for new)

Dependencies: Python stdlib only (Law 4)
"""

import io
import json
import os
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

if os.name == 'nt' and __name__ == '__main__':
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
    sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8')

SCRIPT_DIR = Path(__file__).resolve().parent
ROOT_DIR = SCRIPT_DIR.parent.parent
STATE_DIR = ROOT_DIR / "state"
WORKER_LOGS_DIR = STATE_DIR / "worker_logs"
SESSIONS_FILE = ROOT_DIR / "Hippocampus" / "Evidence" / "Sessions" / "agent_sessions.json"
BRAIN_QUERY_LOG = STATE_DIR / "brain_query_log.json"


def _load_worker_logs() -> list[dict]:
    """Load all worker log files."""
    entries = []
    if not WORKER_LOGS_DIR.exists():
        return entries
    for log_file in WORKER_LOGS_DIR.glob("*.json"):
        try:
            data = json.loads(log_file.read_text(encoding="utf-8"))
            if isinstance(data, list):
                for item in data:
                    item.setdefault("source", log_file.stem)
                    item.setdefault("type", "worker")
                    entries.append(item)
            elif isinstance(data, dict):
                data.setdefault("source", log_file.stem)
                data.setdefault("type", "worker")
                entries.append(data)
        except (json.JSONDecodeError, Exception):
            continue
    return entries


def _load_session_logs() -> list[dict]:
    """Load agent session logs."""
    entries = []
    if not SESSIONS_FILE.exists():
        return entries
    try:
        data = json.loads(SESSIONS_FILE.read_text(encoding="utf-8"))
        if isinstance(data, list):
            for item in data:
                item.setdefault("source", item.get("agent", "unknown"))
                item.setdefault("type", "session")
                entries.append(item)
    except (json.JSONDecodeError, Exception):
        pass
    return entries


def _load_brain_query_logs() -> list[dict]:
    """Load brain query logs."""
    entries = []
    if not BRAIN_QUERY_LOG.exists():
        return entries
    try:
        data = json.loads(BRAIN_QUERY_LOG.read_text(encoding="utf-8"))
        if isinstance(data, list):
            for item in data:
                item.setdefault("source", "brain_query")
                item.setdefault("type", "brain_query")
                entries.append(item)
    except (json.JSONDecodeError, Exception):
        pass
    return entries


def _load_captured_logs() -> list[dict]:
    """Load captured daemon logs."""
    entries = []
    for log_name in ["pulse_captured_claude.log", "pulse_captured_gemini.log",
                     "pulse_captured_codex.log"]:
        log_file = STATE_DIR / log_name
        if not log_file.exists():
            continue
        try:
            agent = log_name.replace("pulse_captured_", "").replace(".log", "")
            for line in log_file.read_text(encoding="utf-8", errors="replace").splitlines()[-100:]:
                line = line.strip()
                if not line:
                    continue
                # Parse [TIMESTAMP] ROLE: CONTENT
                if line.startswith("[") and "]" in line:
                    ts_end = line.index("]")
                    ts = line[1:ts_end]
                    rest = line[ts_end+2:]
                    entries.append({
                        "timestamp": ts,
                        "source": agent,
                        "type": "captured",
                        "message": rest[:200],
                    })
        except Exception:
            continue
    return entries


def aggregate(agent_filter: str = "", task_filter: str = "",
              last_n: int = 100) -> list[dict]:
    """
    Aggregate all log sources into a single sorted list.

    Args:
        agent_filter: Only show logs from this agent
        task_filter: Only show logs for this task ID
        last_n: Maximum number of entries to return

    Returns:
        List of log entries, newest first
    """
    all_entries = []
    all_entries.extend(_load_worker_logs())
    all_entries.extend(_load_session_logs())
    all_entries.extend(_load_brain_query_logs())
    all_entries.extend(_load_captured_logs())

    # Filter
    if agent_filter:
        agent_lower = agent_filter.lower()
        all_entries = [e for e in all_entries
                       if agent_lower in (e.get("source", "") or "").lower()
                       or agent_lower in (e.get("agent", "") or "").lower()]

    if task_filter:
        all_entries = [e for e in all_entries
                       if task_filter in (e.get("task_id", "") or "")
                       or task_filter in (e.get("task", "") or "")]

    # Sort by timestamp (newest first)
    def _ts_key(entry):
        ts = entry.get("timestamp", "")
        if not ts:
            return ""
        return ts

    all_entries.sort(key=_ts_key, reverse=True)
    return all_entries[:last_n]


def format_entry(entry: dict) -> str:
    """Format a single log entry for display."""
    ts = entry.get("timestamp", "?")[:19]
    source = entry.get("source", "?")
    etype = entry.get("type", "?")

    if etype == "session":
        learnings = entry.get("learnings", [])
        task_id = entry.get("task_id", "")
        msg = f"Session done"
        if task_id:
            msg += f" (task: {task_id})"
        if learnings:
            msg += f" | {len(learnings)} learnings"
        return f"{ts} [{source}] {msg}"

    elif etype == "brain_query":
        query = entry.get("query", "")[:60]
        results = entry.get("results", 0)
        return f"{ts} [brain_query] \"{query}\" → {results} results"

    elif etype == "worker":
        msg = entry.get("message", "") or entry.get("status", "")
        return f"{ts} [{source}] {msg[:100]}"

    elif etype == "captured":
        msg = entry.get("message", "")[:100]
        return f"{ts} [{source}] {msg}"

    else:
        msg = str(entry)[:120]
        return f"{ts} [{source}] {msg}"


def follow_mode(agent_filter: str = "", interval: float = 2.0):
    """Watch for new log entries (tail -f mode)."""
    seen = set()
    print("[AGGREGATOR] Following logs (Ctrl+C to stop)...")
    try:
        while True:
            entries = aggregate(agent_filter=agent_filter, last_n=20)
            for entry in reversed(entries):
                key = f"{entry.get('timestamp', '')}{entry.get('source', '')}"
                if key not in seen:
                    seen.add(key)
                    print(format_entry(entry))
            time.sleep(interval)
    except KeyboardInterrupt:
        print("\n[AGGREGATOR] Stopped.")


if __name__ == "__main__":
    agent = ""
    task = ""
    last_n = 50
    follow = False

    args = sys.argv[1:]
    i = 0
    while i < len(args):
        if args[i] == "--agent" and i + 1 < len(args):
            agent = args[i + 1]
            i += 2
        elif args[i] == "--task" and i + 1 < len(args):
            task = args[i + 1]
            i += 2
        elif args[i] == "--last" and i + 1 < len(args):
            last_n = int(args[i + 1])
            i += 2
        elif args[i] == "--follow":
            follow = True
            i += 1
        else:
            i += 1

    if follow:
        follow_mode(agent_filter=agent)
    else:
        entries = aggregate(agent_filter=agent, task_filter=task, last_n=last_n)
        if not entries:
            print("[AGGREGATOR] No log entries found.")
        else:
            print(f"[AGGREGATOR] Showing {len(entries)} entries:")
            print()
            for entry in entries:
                print(format_entry(entry))
