import os
import time
import json
import subprocess

BASE_DIR = "/usr/local/google/home/jerxie/Projects/PersonalProject/cortex-hub/tmp/swarm"
REQUESTS_DIR = os.path.join(BASE_DIR, "requests")
RESPONSES_DIR = os.path.join(BASE_DIR, "responses")
SESSION_NAME = "jetski_swarm"

def process_request(file_path):
    try:
        with open(file_path, 'r') as f:
            req = json.load(f)
        
        req_id = req.get("request_id")
        target = req.get("target_agent")
        prompt = req.get("prompt")
        resp_file = req.get("response_file")
        
        # Mapping for Grid layout
        if target.startswith("agent_"):
            try:
                agent_num = int(target.split("_")[1])
                target = f"grid.{agent_num - 1}"
                print(f"Mapped {req.get('target_agent')} to {target}")
            except ValueError:
                pass
        
        print(f"Processing {req_id} for {target}...")
        
        # Send keys
        cmd = f"tmux send-keys -t {SESSION_NAME}:{target} \"{prompt}\" C-m"
        subprocess.run(cmd, shell=True, check=True)
        
        # Wait for agent to process (simple wait for now)
        time.sleep(5) 
        
        # Capture pane
        cap_cmd = f"tmux capture-pane -p -t {SESSION_NAME}:{target}"
        result = subprocess.run(cap_cmd, shell=True, capture_output=True, text=True, check=True)
        
        # Write response
        resp_data = {
            "request_id": req_id,
            "status": "success",
            "output": result.stdout
        }
        
        # Ensure target file directory exists
        os.makedirs(os.path.dirname(resp_file), exist_ok=True)
        
        with open(resp_file, 'w') as f:
            json.dump(resp_data, f, indent=2)
            
        # Delete request
        os.remove(file_path)
        print(f"Completed {req_id}")
        
    except Exception as e:
        print(f"Error processing {file_path}: {e}")

def main():
    print("Swarm Orchestrator started polling...")
    while True:
        try:
            files = os.listdir(REQUESTS_DIR)
            for file in files:
                if file.endswith(".json"):
                    file_path = os.path.join(REQUESTS_DIR, file)
                    process_request(file_path)
        except Exception as e:
            print(f"Polling error: {e}")
        time.sleep(2)

if __name__ == "__main__":
    main()