import os
import time
import json
import argparse
import uuid

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")

def main():
    parser = argparse.ArgumentParser(description="Swarm CLI")
    parser.add_argument("--target", required=True, help="Target agent")
    parser.add_argument("--prompt", required=True, help="Instruction prompt")
    parser.add_argument("--mode", choices=["normal", "yolo"], default="normal", help="Mode: normal or yolo")
    
    args = parser.parse_args()
    
    req_id = f"req_{uuid.uuid4().hex[:8]}"
    resp_file = os.path.join(RESPONSES_DIR, f"res_{req_id}.json")
    
    req_data = {
        "request_id": req_id,
        "target_agent": args.target,
        "prompt": args.prompt,
        "response_file": resp_file
    }
    
    print(f"Target: {args.target}")
    print(f"Prompt: {args.prompt}")
    print(f"Mode: {args.mode}")
    
    if args.mode == "normal":
        try:
            confirm = input("Are you sure you want to send this command? [y/N]: ")
            if confirm.lower() != 'y':
                print("Aborted.")
                return
        except EOFError:
            print("No interactive terminal found, aborting in normal mode.")
            return
            
    os.makedirs(REQUESTS_DIR, exist_ok=True)
    req_file = os.path.join(REQUESTS_DIR, f"{req_id}.json")
    
    with open(req_file, 'w') as f:
        json.dump(req_data, f, indent=2)
        
    print(f"Request {req_id} sent.")
    print(f"Waiting for response in {resp_file}...")
    
    # Wait for response
    start_time = time.time()
    timeout = 30 # seconds
    while time.time() - start_time < timeout:
        if os.path.exists(resp_file):
            print("Response received:")
            try:
                with open(resp_file, 'r') as f:
                    resp = json.load(f)
                    print(resp.get("output"))
                return
            except json.JSONDecodeError:
                # File might be still being written
                time.sleep(0.5)
                continue
        time.sleep(1)
        
    print("Timed out waiting for response.")

if __name__ == "__main__":
    main()