#!/usr/bin/env python3

import sys
import json

def main():
    if len(sys.argv) < 2:
        print("Usage: unified_job.py {config|exec} [args...]", file=sys.stderr)
        sys.exit(1)
    
    command = sys.argv[1]
    
    if command == "config":
        handle_config(sys.argv[2:])
    elif command == "exec":
        handle_exec(sys.argv[2:])
    else:
        print(f"Unknown command: {command}", file=sys.stderr)
        print("Usage: unified_job.py {config|exec} [args...]", file=sys.stderr)
        sys.exit(1)

def handle_config(args):
    if len(args) < 1:
        print("Config mode requires partition ref", file=sys.stderr)
        sys.exit(1)
    
    partition_ref = args[0]
    
    config = {
        "configs": [{
            "outputs": [{"str": partition_ref}],
            "inputs": [],
            "args": ["Hello", "gorgeous", partition_ref],
            "env": {"PARTITION_REF": partition_ref}
        }]
    }
    
    print(json.dumps(config))

def handle_exec(args):
    print("What a time to be alive.")
    print(f"Partition ref: {os.getenv('PARTITION_REF', 'unknown')}")
    print(f"Args: {args}")

if __name__ == "__main__":
    import os
    main()