# ... (imports)
import sys
import argparse
import os

# Import AegisSandwich
try:
    from .sandwich import AegisSandwich
except ImportError:
    AegisSandwich = None

# ... (other functions: scan_file, protect_shell, report_status, check_updates)

def main():
    # Handle 'run' command manually first to consume remainder args correctly
    if len(sys.argv) > 1 and sys.argv[1] == "run":
        if not AegisSandwich:
            print("[Error] AegisSandwich dependencies (psutil) not found. Install 'aegisflow[run]' or ensure psutil is installed.")
            sys.exit(1)
            
        cmd_args = sys.argv[2:]
        if not cmd_args:
            print("Usage: aegis run <command> [args...]")
            sys.exit(1)
            
        sandwich = AegisSandwich(cmd_args)
        sys.exit(sandwich.run())

    parser = argparse.ArgumentParser(description="AegisFlow Governance Layer")
    subparsers = parser.add_subparsers(dest="command")
    
    # scan command
    scan_parser = subparsers.add_parser("scan", help="Scan a file for threats")
    scan_parser.add_argument("path", help="Path to file or directory")
    
    # protect command
    protect_parser = subparsers.add_parser("protect", help="Start protected shell")
    
    # report command
    report_parser = subparsers.add_parser("report", help="Show security audit report")
    
    # update command
    update_parser = subparsers.add_parser("update", help="Check for threat feed updates")
    
    # run command (in help only, handled above)
    run_parser = subparsers.add_parser("run", help="Wrap and monitor an agent process")

    args = parser.parse_args()
    
    if args.command == "scan":
        if os.path.isfile(args.path):
            scan_file(args.path)
        else:
            print(f"Directory scanning not yet implemented. Please point to a file.")
    elif args.command == "protect":
        protect_shell()
    elif args.command == "report":
        report_status()
    elif args.command == "update":
        check_updates()
    else:
        parser.print_help()

if __name__ == "__main__":
    main()
