#!/usr/bin/env python3
"""
collect-data.py — Data collector for JARVIS Command Center
Focuses on JARVIS's world: MacBook activity, JARVIS↔FRIDAY communication,
strategy/writing tasks, and EDITH framework status.

Data sources:
  1. /root/.hermes/shared-memory/QUEUE.json — inter-agent messages (JARVIS side)
  2. /root/.hermes/logs/agent.log — FRIDAY's log showing JARVIS-related activity
  3. /root/.hermes/sessions/ — active session files
  4. System metrics — MacBook is remote, so we show what we can from this server

Output: /root/jarvis-dashboard/data/live/data.json
"""

import json
import os
import re
import subprocess
import time
from collections import defaultdict
from datetime import datetime, date, timezone
from pathlib import Path

# ── Paths ──────────────────────────────────────────────────────────────────
SHARED_QUEUE = "/root/.hermes/shared-memory/QUEUE.json"
AGENT_LOG = "/root/.hermes/logs/agent.log"
SESSIONS_DIR = "/root/.hermes/sessions"
OUTPUT_DIR = "/root/jarvis-dashboard/data/live"
OUTPUT_FILE = os.path.join(OUTPUT_DIR, "data.json")

os.makedirs(OUTPUT_DIR, exist_ok=True)

# ── JARVIS Model Info ──────────────────────────────────────────────────────
JARVIS_MODEL = "Qwen2.5-7B-Instruct Q4_K_M"
JARVIS_PROVIDER = "Local (MacBook)"
JARVIS_LOCATION = "MacBook Pro (192.168.68.112)"
JARVIS_ROLE = "Strategy & Writing Agent"


def parse_shared_queue():
    """Read inter-agent messages from shared queue — JARVIS's perspective."""
    messages = []
    if not os.path.exists(SHARED_QUEUE):
        return messages
    try:
        with open(SHARED_QUEUE) as f:
            queue = json.load(f)
    except (json.JSONDecodeError, IOError):
        return messages

    for direction in ["friday_to_jarvis", "jarvis_to_friday"]:
        sender = "FRIDAY" if direction == "friday_to_jarvis" else "JARVIS"
        receiver = "JARVIS" if direction == "friday_to_jarvis" else "FRIDAY"
        for msg in queue.get(direction, []):
            ts = msg.get("timestamp", 0)
            if isinstance(ts, (int, float)) and ts > 1e9:
                dt = datetime.fromtimestamp(ts, tz=timezone.utc)
                ts_str = dt.strftime("%Y-%m-%d %H:%M:%S UTC")
            else:
                ts_str = str(ts)
            messages.append({
                "timestamp": ts_str,
                "sender": sender,
                "receiver": receiver,
                "msgType": msg.get("type", "info"),
                "content": msg.get("content", "")[:300],
                "read": msg.get("read", True),
            })
    return messages


def parse_friday_log_for_jarvis():
    """Parse FRIDAY's agent.log for JARVIS-related activity."""
    entries = []
    if not os.path.exists(AGENT_LOG):
        return entries

    jarvis_patterns = [
        re.compile(r"(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}).*?(?:jarvis|JARVIS|shared.queue|QUEUE)", re.IGNORECASE),
        re.compile(r"(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}).*?(?:MCP|friday_to_jarvis|jarvis_to_friday)", re.IGNORECASE),
        re.compile(r"(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}).*?(?:shared-memory|inter.agent|EDITH)", re.IGNORECASE),
    ]

    tool_pattern = re.compile(
        r"(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}),\d+.*tool (\w+) completed"
    )
    response_pattern = re.compile(
        r"(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}),\d+.*response ready:.*platform=(\S+).*time=([\d.]+)s"
    )

    with open(AGENT_LOG, "r", errors="replace") as f:
        for line in f:
            # JARVIS-specific mentions
            for pattern in jarvis_patterns:
                if pattern.search(line):
                    ts_match = re.search(r"(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2})", line)
                    if ts_match:
                        entries.append({
                            "type": "jarvisMention",
                            "timestamp": ts_match.group(1),
                            "message": line.strip()[:200],
                        })
                        break

            # Tool calls
            tool_match = tool_pattern.search(line)
            if tool_match:
                entries.append({
                    "type": "toolCall",
                    "timestamp": tool_match.group(1),
                    "tool": tool_match.group(2),
                })

            # Responses sent
            resp_match = response_pattern.search(line)
            if resp_match:
                entries.append({
                    "type": "responseSent",
                    "timestamp": resp_match.group(1),
                    "platform": resp_match.group(2),
                    "responseTime": float(resp_match.group(3)),
                })

    return entries


def get_system_metrics():
    """Collect system health metrics from the server."""
    metrics = {}
    try:
        with open("/proc/uptime") as f:
            uptime_seconds = float(f.readline().split()[0])
        days = int(uptime_seconds // 86400)
        hours = int((uptime_seconds % 86400) // 3600)
        metrics["uptime"] = f"{days}d {hours}h"
        metrics["uptimeSeconds"] = int(uptime_seconds)
    except Exception:
        metrics["uptime"] = "N/A"
        metrics["uptimeSeconds"] = 0

    try:
        mem = subprocess.run(["free", "-m"], capture_output=True, text=True, timeout=5).stdout
        for line in mem.splitlines():
            if line.startswith("Mem:"):
                parts = line.split()
                total, used = int(parts[1]), int(parts[2])
                metrics["memoryTotalMb"] = total
                metrics["memoryUsedMb"] = used
                metrics["memoryPct"] = round(used / total * 100, 1)
                break
    except Exception:
        metrics["memoryPct"] = 0

    try:
        disk = subprocess.run(["df", "-h", "/"], capture_output=True, text=True, timeout=5).stdout
        for line in disk.splitlines()[1:]:
            parts = line.split()
            metrics["diskTotal"] = parts[1]
            metrics["diskUsed"] = parts[2]
            metrics["diskPct"] = int(parts[3].rstrip("%"))
            break
    except Exception:
        metrics["diskPct"] = 0

    try:
        load = os.getloadavg()
        metrics["cpuLoad1m"] = round(load[0], 2)
        metrics["cpuLoad5m"] = round(load[1], 2)
        metrics["cpuLoad15m"] = round(load[2], 2)
        cpu_count = os.cpu_count() or 1
        metrics["cpuPct"] = round(load[0] / cpu_count * 100, 1)
        metrics["cpuCount"] = cpu_count
    except Exception:
        metrics["cpuLoad1m"] = 0
        metrics["cpuPct"] = 0
        metrics["cpuCount"] = 1

    return metrics


def build_jarvis_status(queue_messages, activity_entries, today_str):
    """Build JARVIS's agent status."""
    jarvis = {
        "name": "JARVIS",
        "location": JARVIS_LOCATION,
        "model": JARVIS_MODEL,
        "provider": JARVIS_PROVIDER,
        "role": JARVIS_ROLE,
        "status": "online",
        "lastActivity": None,
        "currentTask": "Standing by",
        "todayMessages": 0,
        "todayToolCalls": 0,
        "accentColor": "#F59E0B",
    }

    # Count today's JARVIS activity from queue
    for msg in queue_messages:
        if msg.get("sender") == "JARVIS" and today_str in msg.get("timestamp", ""):
            jarvis["todayMessages"] += 1
            jarvis["lastActivity"] = msg["timestamp"]

    # Count today's tool calls from FRIDAY's log
    for entry in activity_entries:
        if entry.get("timestamp", "").startswith(today_str):
            if entry["type"] == "toolCall":
                jarvis["todayToolCalls"] += 1
                jarvis["lastActivity"] = entry["timestamp"]

    # Determine current task
    jarvis_mentions = [e for e in activity_entries if e["type"] == "jarvisMention" and e.get("timestamp", "").startswith(today_str)]
    if jarvis_mentions:
        last = jarvis_mentions[-1]
        msg_lower = last.get("message", "").lower()
        if "email" in msg_lower or "send" in msg_lower:
            jarvis["currentTask"] = "Composing communications"
        elif "research" in msg_lower or "search" in msg_lower:
            jarvis["currentTask"] = "Researching"
        elif "strateg" in msg_lower:
            jarvis["currentTask"] = "Strategy work"
        elif "write" in msg_lower or "draft" in msg_lower:
            jarvis["currentTask"] = "Writing"
        elif "queue" in msg_lower:
            jarvis["currentTask"] = "Processing queue"
        else:
            jarvis["currentTask"] = "Active"

    # Check idle
    if jarvis["lastActivity"]:
        try:
            last_dt = datetime.strptime(jarvis["lastActivity"].split(",")[0][:19], "%Y-%m-%d %H:%M:%S")
            idle_minutes = (datetime.now() - last_dt).total_seconds() / 60
            if idle_minutes > 60:
                jarvis["status"] = "idle"
                jarvis["currentTask"] = "Standing by"
        except (ValueError, TypeError):
            pass

    return jarvis


def build_activity_feed(queue_messages, activity_entries):
    """Build combined activity feed from queue and logs."""
    feed = []

    # Queue messages
    for msg in queue_messages:
        feed.append({
            "timestamp": msg["timestamp"],
            "agent": msg["sender"],
            "message": msg.get("content", "")[:150],
            "type": "queueMessage",
        })

    # Activity entries
    for entry in activity_entries:
        if entry["type"] == "jarvisMention":
            feed.append({
                "timestamp": entry["timestamp"],
                "agent": "FRIDAY",
                "message": f"JARVIS mention: {entry.get('message', '')[:100]}",
                "type": "jarvisMention",
            })
        elif entry["type"] == "responseSent":
            feed.append({
                "timestamp": entry["timestamp"],
                "agent": "FRIDAY",
                "message": f"Response sent via {entry.get('platform', 'unknown')}",
                "type": "response",
            })
        elif entry["type"] == "toolCall":
            feed.append({
                "timestamp": entry["timestamp"],
                "agent": "FRIDAY",
                "message": f"Tool: {entry.get('tool', 'unknown')}",
                "type": "toolCall",
            })

    feed.sort(key=lambda x: x.get("timestamp", ""))
    return feed[-100:]


def build_heatmap_from_queue(queue_messages):
    """Build activity heatmap from queue message frequency."""
    days = defaultdict(int)
    for msg in queue_messages:
        ts = msg.get("timestamp", "")
        if " " in ts:
            day = ts.split(" ")[0]
            days[day] += 1

    if not days:
        return []

    max_count = max(days.values())
    result = []
    for day in sorted(days.keys()):
        count = days[day]
        intensity = count / max_count if max_count else 0
        level = 3 if intensity > 0.75 else 2 if intensity > 0.45 else 1 if intensity > 0.15 else 0
        result.append({
            "date": day,
            "messages": count,
            "intensity": level,
        })
    return result[-30:]  # Last 30 days


def generate_fallback_data():
    """Generate fallback data when no real data is available."""
    today = datetime.now().strftime("%Y-%m-%d")
    return {
        "queue": [],
        "activityFeed": [],
        "agent": {
            "name": "JARVIS",
            "location": JARVIS_LOCATION,
            "model": JARVIS_MODEL,
            "provider": JARVIS_PROVIDER,
            "role": JARVIS_ROLE,
            "status": "idle",
            "lastActivity": None,
            "currentTask": "No data available — JARVIS may be offline",
            "todayMessages": 0,
            "todayToolCalls": 0,
            "accentColor": "#F59E0B",
        },
        "system": get_system_metrics(),
        "stats": {
            "totalLogEntries": 0,
            "totalQueueMessages": 0,
            "activeSessions24h": 0,
        },
    }


# ── Main ───────────────────────────────────────────────────────────────────
def main():
    today = date.today().isoformat()

    # 1. Parse shared queue
    queue_messages = parse_shared_queue()

    # 2. Parse FRIDAY's log for JARVIS-related activity
    activity_entries = parse_friday_log_for_jarvis()

    # 3. System metrics
    system = get_system_metrics()

    # 4. JARVIS status
    agent = build_jarvis_status(queue_messages, activity_entries, today)

    # 5. Activity feed
    activity_feed = build_activity_feed(queue_messages, activity_entries)

    # 6. Heatmap
    heatmap = build_heatmap_from_queue(queue_messages)

    # Build output with camelCase keys matching DashboardData type
    output = {
        "generatedAt": datetime.now().isoformat(),
        "agent": agent,
        "queue": queue_messages,
        "activityFeed": activity_feed,
        "heatmap": heatmap,
        "system": system,
        "stats": {
            "totalLogEntries": len(activity_entries),
            "totalQueueMessages": len(queue_messages),
            "activeSessions24h": len([e for e in activity_entries if e.get("timestamp", "").startswith(today)]),
        },
        "modelDistribution": [
            {
                "model": JARVIS_MODEL,
                "provider": JARVIS_PROVIDER,
                "tokenVolume": 0,
                "cost": 0.0,
                "percentage": 100.0,
                "calls": agent["todayMessages"],
                "costPerMillion": 0.0,
            }
        ],
        "timeSeries": [],
        "sessions": [],
        "timeline": [
            {
                "timestamp": msg["timestamp"],
                "agent": msg["sender"],
                "event": msg.get("content", "")[:120],
                "icon": "message",
            }
            for msg in queue_messages[-20:]
        ],
        "activeTasks": [
            {
                "id": "jarvis-queue",
                "title": "Inter-Agent Queue",
                "status": "inProgress",
                "startedAt": datetime.now().strftime("%Y-%m-%d %H:%M"),
                "progress": min(len(queue_messages), 100),
            },
            {
                "id": "jarvis-friday-bridge",
                "title": "FRIDAY ↔ JARVIS Bridge",
                "status": "inProgress" if queue_messages else "offline",
                "startedAt": datetime.now().strftime("%Y-%m-%d %H:%M"),
                "progress": 100 if queue_messages else 0,
            },
        ],
    }

    with open(OUTPUT_FILE, "w") as f:
        json.dump(output, f, indent=2, default=str)

    print(f"Data written to {OUTPUT_FILE}")
    print(f"  Queue messages: {len(queue_messages)}")
    print(f"  Activity entries: {len(activity_entries)}")
    print(f"  JARVIS status: {agent['status']} — {agent['currentTask']}")
    print(f"  CPU: {system.get('cpuPct', 'N/A')}% | Mem: {system.get('memoryPct', 'N/A')}% | Disk: {system.get('diskPct', 'N/A')}%")


if __name__ == "__main__":
    main()