import os
import json
import subprocess
import sys


def load_env_vars(json_file):
    """
    Load environment variables from a JSON file and set them in the current process.
    """
    try:
        with open(json_file, 'r') as f:
            env_vars = json.load(f)

        if not isinstance(env_vars, dict):
            raise ValueError("JSON file must contain key-value pairs.")

        # Set each key-value pair as an environment variable
        for key, value in env_vars.items():
            os.environ[key] = str(value)
            print(f"Set environment variable: {key}={value}")

    except FileNotFoundError:
        print(f"Error: File not found - {json_file}")
        sys.exit(1)
    except json.JSONDecodeError:
        print(f"Error: Failed to parse JSON file - {json_file}")
        sys.exit(1)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)


def run_docker(image, env_vars, docker_args, jam_track_id):
    """
    Run a Docker container with the specified image, environment variables, and additional arguments.
    """
    try:
        # Build the Docker run command with environment variables
        docker_cmd = ["docker", "run", "--rm"]
        for key, value in env_vars.items():
            docker_cmd.extend(["-e", f"{key}={value}"])

        docker_cmd.extend(["-e", f"JAM_TRACK_ID={jam_track_id}"])
        # Add the Docker image and additional arguments
        docker_cmd.append(image)
        docker_cmd.extend(docker_args)

        # Execute the Docker run command
        print(f"Running Docker command: {' '.join(docker_cmd)}")
        process = subprocess.Popen(docker_cmd, stdout = sys.stdout, stderr = sys.stderr, text=True)
        #result = subprocess.run(docker_cmd, capture_output=False, text=True)
        process.wait()

        # Print the output or handle errors
        if process.returncode == 0:
            print("Docker runner succeeded")
        else:
            print(f"Docker runner failed #{process.returncode}:")

    except Exception as e:
        print(f"Error running Docker: {e}")
        sys.exit(1)


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("Usage: python run_batch_job.py <docker_image> <json_file> [additional_args...]")
        sys.exit(1)

    docker_image = sys.argv[1]
    json_file = sys.argv[2]
    jam_track_id = sys.argv[3]
    docker_args = sys.argv[4:]  # All remaining arguments are passed to the Docker command

    # Load environment variables from JSON
    load_env_vars(json_file)

    # Extract current environment variables (after setting from JSON)
    current_env = {key: value for key, value in os.environ.items() if key in json.load(open(json_file))}
    current_env["AWS_ACCESS_KEY_ID"] = os.environ["AWS_KEY"]
    current_env["AWS_SECRET_ACCESS_KEY"] = os.environ["AWS_SECRET"]

    # Run the Docker container
    run_docker(docker_image, current_env, docker_args, jam_track_id)