"""
Genesis V2 Swarm Manager
=========================
Manages agent nodes via tmux persistent sessions.
"""

import json
import subprocess
import logging
from pathlib import Path
from dataclasses import dataclass, field, asdict
from typing import Optional, Dict, List

logger = logging.getLogger("genesis_v2.core.swarm")

STATE_FILE = Path(__file__).parent.parent.parent / ".genesis_swarm_state.json"


@dataclass
class SwarmNode:
    """Represents an agent node in the swarm."""
    name: str
    model: str
    provider: str = "google"
    mission: str = ""
    status: str = "idle"       # idle, running, stopped
    session_type: str = "api"  # api, tmux
    
    def to_dict(self) -> dict:
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: dict) -> "SwarmNode":
        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})


class SwarmManager:
    """Manages a fleet of agent nodes."""
    
    def __init__(self):
        self.nodes: Dict[str, SwarmNode] = {}
        self._load_state()
    
    def _load_state(self):
        """Load persisted swarm state."""
        if STATE_FILE.exists():
            try:
                data = json.loads(STATE_FILE.read_text())
                for name, node_data in data.get("nodes", {}).items():
                    self.nodes[name] = SwarmNode.from_dict(node_data)
            except (json.JSONDecodeError, Exception) as e:
                logger.error(f"Failed to load swarm state: {e}")
    
    def _save_state(self):
        """Persist swarm state to disk."""
        data = {"nodes": {name: node.to_dict() for name, node in self.nodes.items()}}
        STATE_FILE.write_text(json.dumps(data, indent=2))
    
    def get_tmux_sessions(self) -> List[str]:
        """Get list of active tmux session names."""
        try:
            result = subprocess.run(
                ["tmux", "list-sessions", "-F", "#{session_name}"],
                capture_output=True, text=True
            )
            if result.returncode == 0:
                return [s.strip() for s in result.stdout.strip().split("\n") if s.strip()]
        except Exception:
            pass
        return []
    
    def list_nodes(self) -> List[SwarmNode]:
        """List all nodes, updating tmux status."""
        active_sessions = self.get_tmux_sessions()
        for name, node in self.nodes.items():
            if node.session_type == "tmux":
                if name in active_sessions:
                    node.status = "running"
                else:
                    node.status = "stopped"
        self._save_state()
        return list(self.nodes.values())
    
    def spawn_node(self, name: str, model: str, provider: str = "google", mission: str = "", use_tmux: bool = True) -> SwarmNode:
        node = SwarmNode(name, model, provider, mission)
        if use_tmux:
            node.session_type = "tmux"
            # Ensure logs directory exists
            log_dir = "/mnt/e/genesis-system/logs"
            subprocess.run(["mkdir", "-p", log_dir])
            log_file = f"{log_dir}/{name}.log"
            # Build the worker command with proper env and logging
            worker_cmd = (
                f"export PYTHONPATH=/mnt/e/genesis-system/lib:/mnt/e/genesis-system; "
                f"echo '[GENESIS] Mission: {mission}' | tee {log_file}; "
                f"echo '[GENESIS] Started: $(date)' | tee -a {log_file}; "
                f"python3 /mnt/e/genesis-system/run_genesis.py "
                f"--mission '{mission}' 2>&1 | tee -a {log_file}; "
                f"echo '[GENESIS] Exited: $(date)' | tee -a {log_file}; "
                f"echo 'Press Enter to close...'; read"
            )
            # Use remain-on-exit so session stays visible even if worker crashes
            subprocess.run(["tmux", "new-session", "-d", "-s", name, "bash", "-c", worker_cmd])
            subprocess.run(["tmux", "set-option", "-t", name, "remain-on-exit", "on"], capture_output=True)
            node.status = "running"
        else:
            node.session_type = "api"
            node.status = "idle"
            
        self.nodes[name] = node
        self._save_state()
        return node
    
    def kill_node(self, name: str) -> bool:
        """Kill a node and its tmux session."""
        if name in self.nodes:
            node = self.nodes[name]
            if node.session_type == "tmux":
                subprocess.run(["tmux", "kill-session", "-t", name], capture_output=True)
            node.status = "stopped"
            self._save_state()
            return True
        return False
    
    def attach_node(self, name: str) -> bool:
        """Attach to a tmux session (interactive only)."""
        if name in self.nodes and self.nodes[name].session_type == "tmux":
            subprocess.run(["tmux", "attach-session", "-t", name])
            return True
        return False


# Global instance
manager = SwarmManager()