import os
import subprocess
import multiprocessing
import time
from typing import Dict, Any, List

class SwarmCoordinator:
    def __init__(self, base_repo: str, num_agents: int = 32):
        self.base_repo = os.path.abspath(base_repo)
        self.num_agents = min(num_agents, 32)
        self.worktrees: Dict[int, str] = {}
        self.agent_states: Dict[int, str] = {}
        self.results: Dict[int, Any] = {}
        
        if not os.path.isdir(self.base_repo):
            raise ValueError(f"Base repository path does not exist: {self.base_repo}")
        
        # Verify it's a git repository
        if not os.path.isdir(os.path.join(self.base_repo, ".git")):
            raise ValueError(f"Base repository is not a git repository: {self.base_repo}")

    def create_worktrees(self) -> None:
        worktree_base = os.path.join(self.base_repo, ".worktrees")
        os.makedirs(worktree_base, exist_ok=True)
        
        for i in range(self.num_agents):
            worktree_path = os.path.join(worktree_base, f"agent_{i}")
            if not os.path.exists(worktree_path):
                subprocess.run([
                    "git", "worktree", "add", worktree_path,
                    "--detach"
                ], cwd=self.base_repo, check=True, capture_output=True)
            self.worktrees[i] = worktree_path
            self.agent_states[i] = "pending"

    def _run_agent(self, agent_id: int, worktree_path: str) -> None:
        """Execute agent logic in isolated worktree"""
        try:
            os.chdir(worktree_path)
            
            # Simulate agent processing (replace with actual agent logic)
            time.sleep(0.1)  # Simulate work
            
            self.results[agent_id] = {
                "agent_id": agent_id,
                "worktree": worktree_path,
                "status": "completed",
                "timestamp": time.strftime("%Y-%m-%d %H:%M:%S")
            }
            self.agent_states[agent_id] = "completed"
        except Exception as e:
            self.results[agent_id] = {
                "error": str(e),
                "agent_id": agent_id,
                "worktree": worktree_path
            }
            self.agent_states[agent_id] = "failed"

    def run_agents(self) -> None:
        processes = []
        for agent_id, worktree_path in self.worktrees.items():
            p = multiprocessing.Process(
                target=self._run_agent,
                args=(agent_id, worktree_path),
                name=f"Agent_{agent_id}"
            )
            p.start()
            processes.append(p)
            self.agent_states[agent_id] = "running"

        # Wait for all processes to complete
        for p in processes:
            p.join()

    def merge_results(self) -> Dict[str, Any]:
        """Merge results with intelligent aggregation"""
        completed = sum(1 for state in self.agent_states.values() if state == "completed")
        
        # Merge results into a single structure
        merged = {
            "total_agents": self.num_agents,
            "completed_agents": completed,
            "agent_results": self.results,
            "summary": {
                "latest": self.results.get(max(self.results.keys()), {}).get("timestamp", "N/A"),
                "failed": [k for k, v in self.agent_states.items() if v == "failed"]
            }
        }
        
        return merged

    def run(self) -> Dict[str, Any]:
        """Orchestrate full agent execution cycle"""
        self.create_worktrees()
        self.run_agents()
        return self.merge_results()

if __name__ == "__main__":
    # Example usage
    coordinator = SwarmCoordinator(
        base_repo=os.getcwd(),
        num_agents=4  # For demonstration (max 32)
    )
    results = coordinator.run()
    print("Swarm Coordinator Results:")
    print(results)