#!/usr/bin/env python3
"""
Genesis Agent Heartbeat System
Each Claude Code instance calls this to report its status.
Other agents and the dashboard read the shared status file.
"""

import json
import os
import time
import socket
import hashlib
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional, List, Dict, Any
import threading

# Paths
GENESIS_ROOT = Path("/mnt/e/genesis-system")
STATUS_FILE = GENESIS_ROOT / "data" / "agent_status.json"
HEARTBEAT_LOG = GENESIS_ROOT / "logs" / "agent_heartbeats.jsonl"
LOCK_FILE = GENESIS_ROOT / "data" / ".agent_status.lock"

# Heartbeat interval
HEARTBEAT_INTERVAL = 5  # seconds


def get_agent_id() -> str:
    """Generate unique agent ID based on PID and working directory."""
    cwd = os.getcwd()
    pid = os.getpid()
    # Create short hash for readability
    unique = f"{cwd}-{pid}-{socket.gethostname()}"
    short_hash = hashlib.md5(unique.encode()).hexdigest()[:6]

    # Determine agent type from path
    if "rwl-1" in cwd:
        return f"claude-rwl-1-{short_hash}"
    elif "rwl-2" in cwd:
        return f"claude-rwl-2-{short_hash}"
    elif "rwl-3" in cwd:
        return f"claude-rwl-3-{short_hash}"
    elif "genesis-system" in cwd:
        return f"claude-main-{short_hash}"
    else:
        return f"claude-unknown-{short_hash}"


def get_git_branch() -> str:
    """Get current git branch."""
    try:
        import subprocess
        result = subprocess.run(
            ["git", "branch", "--show-current"],
            capture_output=True, text=True, cwd=os.getcwd()
        )
        return result.stdout.strip() or "unknown"
    except:
        return "unknown"


def get_modified_files() -> List[str]:
    """Get list of currently modified files."""
    try:
        import subprocess
        result = subprocess.run(
            ["git", "diff", "--name-only"],
            capture_output=True, text=True, cwd=os.getcwd()
        )
        files = result.stdout.strip().split("\n")
        return [f for f in files if f][:10]  # Limit to 10 files
    except:
        return []


def load_status() -> Dict[str, Any]:
    """Load current agent status from shared file."""
    if STATUS_FILE.exists():
        try:
            with open(STATUS_FILE, 'r') as f:
                return json.load(f)
        except:
            pass
    return {"agents": {}, "activity_log": [], "file_locks": {}}


def save_status(status: Dict[str, Any]) -> None:
    """Save agent status to shared file with simple locking."""
    # Simple file-based locking
    lock_acquired = False
    for _ in range(10):  # Try 10 times
        try:
            if not LOCK_FILE.exists():
                LOCK_FILE.touch()
                lock_acquired = True
                break
        except:
            pass
        time.sleep(0.1)

    try:
        STATUS_FILE.parent.mkdir(parents=True, exist_ok=True)
        with open(STATUS_FILE, 'w') as f:
            json.dump(status, f, indent=2, default=str)
    finally:
        if lock_acquired:
            try:
                LOCK_FILE.unlink()
            except:
                pass


def log_activity(agent_id: str, action: str, details: str = "") -> None:
    """Log activity to heartbeat log."""
    entry = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "agent_id": agent_id,
        "action": action,
        "details": details
    }
    HEARTBEAT_LOG.parent.mkdir(parents=True, exist_ok=True)
    with open(HEARTBEAT_LOG, 'a') as f:
        f.write(json.dumps(entry) + "\n")


class AgentHeartbeat:
    """Manages heartbeat for a single Claude Code agent."""

    def __init__(self):
        self.agent_id = get_agent_id()
        self.current_task: Optional[str] = None
        self.current_story: Optional[str] = None
        self.status: str = "idle"
        self._stop_event = threading.Event()
        self._thread: Optional[threading.Thread] = None

    def update_task(self, task: str, story: Optional[str] = None) -> None:
        """Update current task being worked on."""
        self.current_task = task
        self.current_story = story
        self.status = "active"
        self._send_heartbeat()
        log_activity(self.agent_id, "task_started", f"{task} - {story or 'N/A'}")

    def complete_task(self) -> None:
        """Mark current task as complete."""
        if self.current_task:
            log_activity(self.agent_id, "task_completed", self.current_task)
        self.current_task = None
        self.current_story = None
        self.status = "idle"
        self._send_heartbeat()

    def claim_file(self, filepath: str) -> bool:
        """Claim a file for editing (prevents conflicts)."""
        status = load_status()
        file_locks = status.get("file_locks", {})

        # Check if file is already claimed by another agent
        if filepath in file_locks:
            owner = file_locks[filepath].get("agent_id")
            if owner and owner != self.agent_id:
                # Check if lock is stale (>5 minutes)
                lock_time = file_locks[filepath].get("timestamp", "")
                try:
                    lock_dt = datetime.fromisoformat(lock_time.replace("Z", "+00:00"))
                    if (datetime.now(timezone.utc) - lock_dt).seconds < 300:
                        return False  # File is locked by another active agent
                except:
                    pass

        # Claim the file
        file_locks[filepath] = {
            "agent_id": self.agent_id,
            "timestamp": datetime.now(timezone.utc).isoformat()
        }
        status["file_locks"] = file_locks
        save_status(status)
        log_activity(self.agent_id, "file_claimed", filepath)
        return True

    def release_file(self, filepath: str) -> None:
        """Release a file claim."""
        status = load_status()
        file_locks = status.get("file_locks", {})
        if filepath in file_locks and file_locks[filepath].get("agent_id") == self.agent_id:
            del file_locks[filepath]
            status["file_locks"] = file_locks
            save_status(status)
            log_activity(self.agent_id, "file_released", filepath)

    def _send_heartbeat(self) -> None:
        """Send heartbeat update to shared status file."""
        status = load_status()

        # Update this agent's status
        status["agents"][self.agent_id] = {
            "agent_id": self.agent_id,
            "status": self.status,
            "current_task": self.current_task,
            "current_story": self.current_story,
            "branch": get_git_branch(),
            "working_dir": os.getcwd(),
            "modified_files": get_modified_files(),
            "last_heartbeat": datetime.now(timezone.utc).isoformat(),
            "pid": os.getpid()
        }

        # Add to activity log (keep last 50 entries)
        activity_log = status.get("activity_log", [])
        activity_log.insert(0, {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "agent_id": self.agent_id,
            "status": self.status,
            "task": self.current_task
        })
        status["activity_log"] = activity_log[:50]

        # Clean up stale agents (no heartbeat in 60 seconds)
        now = datetime.now(timezone.utc)
        for aid, agent in list(status["agents"].items()):
            try:
                last_hb = datetime.fromisoformat(agent["last_heartbeat"].replace("Z", "+00:00"))
                if (now - last_hb).seconds > 60:
                    status["agents"][aid]["status"] = "stale"
            except:
                pass

        save_status(status)

    def start_background_heartbeat(self) -> None:
        """Start background thread for periodic heartbeats."""
        def heartbeat_loop():
            while not self._stop_event.is_set():
                self._send_heartbeat()
                self._stop_event.wait(HEARTBEAT_INTERVAL)

        self._thread = threading.Thread(target=heartbeat_loop, daemon=True)
        self._thread.start()
        log_activity(self.agent_id, "agent_started", f"Branch: {get_git_branch()}")

    def stop(self) -> None:
        """Stop the heartbeat thread."""
        self._stop_event.set()
        self.status = "offline"
        self._send_heartbeat()
        log_activity(self.agent_id, "agent_stopped", "")


# Global heartbeat instance
_heartbeat: Optional[AgentHeartbeat] = None


def get_heartbeat() -> AgentHeartbeat:
    """Get or create the global heartbeat instance."""
    global _heartbeat
    if _heartbeat is None:
        _heartbeat = AgentHeartbeat()
        _heartbeat.start_background_heartbeat()
    return _heartbeat


def report_task(task: str, story: Optional[str] = None) -> None:
    """Convenience function to report current task."""
    get_heartbeat().update_task(task, story)


def complete_task() -> None:
    """Convenience function to mark task complete."""
    get_heartbeat().complete_task()


def claim_file(filepath: str) -> bool:
    """Convenience function to claim a file."""
    return get_heartbeat().claim_file(filepath)


def release_file(filepath: str) -> None:
    """Convenience function to release a file."""
    get_heartbeat().release_file(filepath)


# CLI interface
if __name__ == "__main__":
    import sys

    if len(sys.argv) < 2:
        print("Usage: python agent_heartbeat.py <command> [args]")
        print("Commands:")
        print("  status     - Show all agent statuses")
        print("  start      - Start heartbeat (runs in foreground)")
        print("  report     - Report task: report <task> [story]")
        print("  complete   - Mark task complete")
        print("  claim      - Claim file: claim <filepath>")
        print("  release    - Release file: release <filepath>")
        sys.exit(1)

    cmd = sys.argv[1]

    if cmd == "status":
        status = load_status()
        print(json.dumps(status, indent=2))

    elif cmd == "start":
        hb = get_heartbeat()
        print(f"Heartbeat started for {hb.agent_id}")
        print("Press Ctrl+C to stop...")
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            hb.stop()
            print("\nStopped.")

    elif cmd == "report":
        task = sys.argv[2] if len(sys.argv) > 2 else "unknown"
        story = sys.argv[3] if len(sys.argv) > 3 else None
        report_task(task, story)
        print(f"Reported: {task}")

    elif cmd == "complete":
        complete_task()
        print("Task marked complete")

    elif cmd == "claim":
        filepath = sys.argv[2] if len(sys.argv) > 2 else ""
        if claim_file(filepath):
            print(f"Claimed: {filepath}")
        else:
            print(f"CONFLICT: {filepath} is locked by another agent!")
            sys.exit(1)

    elif cmd == "release":
        filepath = sys.argv[2] if len(sys.argv) > 2 else ""
        release_file(filepath)
        print(f"Released: {filepath}")