# TODO: Merge this and launch_sagemaker.py to avoid code duplication 

import argparse
import time
import os
import subprocess
from datetime import datetime
from pathlib import Path

import boto3
import yaml
from botocore.config import Config
from sagemaker import Session as sm_Session
from sagemaker.pytorch import PyTorch
from sagemaker.inputs import FileSystemInput
from sagemaker.debugger import ProfilerConfig, FrameworkProfile, ProfilerRule, rule_configs

try:
    from sagemaker.batch_queueing.queue import Queue
    print(f"Loading SageMaker batch queueing.")
    is_sm_queue = True
except Exception as e:
    print(f"Could not load SageMaker batch queueing: {e}.")
    print(f"Trying to install")
    os.system("bash scripts/setup_sm_batch.sh")
    from sagemaker.batch_queueing.queue import Queue
    print(f"Loading SageMaker batch queueing.")
    is_sm_queue = True
except:
    is_sm_queue = False

is_sm_queue = False

NAME = "cosmos-predict2"
INSTANCE_MAPPER = {
    "p4d": "ml.p4d.24xlarge",
    "p4de": "ml.p4de.24xlarge",
    "p5": "ml.p5.48xlarge",
    "p5en": "ml.p5en.48xlarge",
    "g6e": "ml.g6e.48xlarge",
    "g6e-small": "ml.g6e.12xlarge",
    "g5": "ml.g5.48xlarge",
    "g5-small": "ml.g5.24xlarge"
}


def run_command(command):
    print(f"=> {command}")
    subprocess.run(command, shell=True, check=True)


def get_image(user, instance_type, version="251", build_type="full", profile="default", region="us-east-1"):
    os.environ["AWS_PROFILE"] = f"{profile}"
    account = subprocess.getoutput(
        f"aws --region {region} --profile {profile} sts get-caller-identity --query Account --output text"
    )
    docker_dir = Path(__file__).parent
    if instance_type in INSTANCE_MAPPER.keys():
        algorithm_name = f"{user}-{NAME}-{version}"
        dockerfile_base = docker_dir / f"Dockerfile_{version}"
        dockerfile_update = docker_dir / "Dockerfile_update"
    else:
        raise ValueError(f"Unknown instance_type: {instance_type}")
    fullname = f"{account}.dkr.ecr.{region}.amazonaws.com/{algorithm_name}:latest"
    if build_type == "none":
        return fullname

    login_cmd = f"aws ecr get-login-password --region {region} --profile {profile} | docker login --username AWS --password-stdin"

    if build_type == "full":
        print("Building container")
        commands = [
            # Log in to Sagemaker account to get image.
            f"{login_cmd} 763104351884.dkr.ecr.{region}.amazonaws.com",
            f"docker build -f {dockerfile_base} --build-arg AWS_REGION={region} -t {algorithm_name} .",
            f"docker tag {algorithm_name} {fullname}",
            f"{login_cmd} {fullname}",
            (
                f"aws --region {region} --profile {profile} ecr describe-repositories --repository-names {algorithm_name} || "
                f"aws --region {region} --profile {profile} ecr create-repository --repository-name {algorithm_name}"
            ),
        ]
    elif build_type == "update":
        print("Updating container")
        commands = [
            f"docker build -f {dockerfile_update} --build-arg BASE_DOCKER={algorithm_name} -t {algorithm_name} .",
            f"docker tag {algorithm_name} {fullname}",
            f"{login_cmd} {fullname}",
        ]
    else:
        raise ValueError(f"Unknown build_type: {build_type}")

    # Create command, making sure to exit if any part breaks.
    command = "\n".join([f"{x} || exit 1" for x in commands])
    run_command(command)
    run_command(f"docker push {fullname}")
    print("Sleeping for 5 seconds to ensure push succeeded")
    time.sleep(5)
    return fullname


def main():
    # Use first line of file docstring as description if it exists.
    parser = argparse.ArgumentParser()
    parser.add_argument("--build-type", choices=["full", "update", "none"], help="Build image from scratch")
    parser.add_argument("--local", action="store_true")
    parser.add_argument("--user", required=True, help="User name")
    parser.add_argument("--entry_point", type=str, default="scripts/train_sm.py")
    
    parser.add_argument("--args_file", help="File that stores all args to be passed", type=str, default=None)

    # AWS profile args
    parser.add_argument("--region", default="us-west-2", help="AWS region")
    parser.add_argument("--profile", default="default", help="AWS profile to use")
    parser.add_argument("--arn", default=None, help="If None, reads from SAGEMAKER_ARN env var")
    parser.add_argument(
        "--s3-remote-sync", default=None, help="S3 path to sync to. If none, reads from S3_REMOTE_SYNC env var"
    )

    # Instance args
    parser.add_argument("--instance-count", default=1, type=int, help="Number of instances")
    parser.add_argument("--instance-type", default="p4de", choices=list(INSTANCE_MAPPER.keys()))
    parser.add_argument("--version", default="271", type=str, help="Choose from: (271)")
    parser.add_argument("--spot-instance", action="store_true")

    parser.add_argument('--base-job-name', type=str)
    parser.add_argument('--input-source', choices=['s3', 'lustre', 'local'], default='s3')

    # Jobs Queue
    parser.add_argument("--fss-identifier", default="default", help="Share identifier for FSS queue")
    parser.add_argument("--priority", default=5, type=int, help="Priority of the job")
    parser.add_argument("--queue", type=str, default='ml', help="Job queue")
    parser.add_argument("--name", type=str, default=None)

    args = parser.parse_args()
    main_after_setup_move(args)


def main_after_setup_move(args):
    if args.arn is None:
        assert "SAGEMAKER_ARN" in os.environ, "Please specify --arn or set the SAGEMAKER_ARN environment variable"
        args.arn = os.environ["SAGEMAKER_ARN"]
    
    if args.s3_remote_sync is None:
        assert (
            "S3_REMOTE_SYNC" in os.environ
        ), "Please specify --s3-remote-sync or set the S3_REMOTE_SYNC environment variable"
        args.s3_remote_sync = os.environ["S3_REMOTE_SYNC"]

    image_uri = get_image(
        args.user,
        args.instance_type,
        args.version,
        region=args.region,
        build_type=args.build_type,
        profile=args.profile,
    )

    ##########
    # Create session and make sure of account and region
    ##########
    sagemaker_session = sm_Session(
        boto_session=boto3.session.Session(
            region_name=args.region,
            profile_name=args.profile
        )
    )

    role = args.arn
    # provide a pre-existing role ARN as an alternative to creating a new role
    role_name = role.split(["/"][-1])

    account = '124224456861'

    session = boto3.session.Session(region_name=args.region)
    region = session.region_name

    ##########
    # Configure the training
    ##########
    base_job_name = args.base_job_name # f"{args.user.replace('.', '-')}-{NAME}"

    def get_job_name(base):
        now = datetime.now()
        # Format example: 2023-03-03-10-14-02-324
        now_ms_str = f"{now.microsecond // 1000:03d}"
        date_str = f"{now.strftime('%Y-%m-%d-%H-%M-%S')}-{now_ms_str}"
        job_name = "-".join([base, date_str])
        return job_name

    if args.name is None:
        job_name = get_job_name(base_job_name)
    else:
        job_name = f"{base_job_name}--{args.name}--".replace('_', '-')

    output_root = f"{args.s3_remote_sync}/sagemaker/{args.user}/{NAME}/"
    output_s3 = os.path.join(output_root, job_name)
    
    tags = [
        {
            "Key": "tri.project", 
            "Value": "MM:PJ-0077",
        },
        {
            "Key": "tri.owner.email",
            "Value": "dian.chen@tri.global",
        },
    ]

    max_run = 5 * 24 * 60 * 60
    max_wait = 5 * 24 * 60 * 60 if args.spot_instance else None
    keep_alive_period_in_seconds = 60 * 60 #if not args.spot_instance else None  

    entry_point = args.entry_point
    instance_type = INSTANCE_MAPPER[args.instance_type]
    instance_count = args.instance_count
    train_use_spot_instances = args.spot_instance
    
    checkpoint_s3_uri = os.path.join(
        f's3://tri-ml-sandbox-16011-us-west-2-datasets/sagemaker/{args.user}/{NAME}', job_name)

    checkpoint_local_path = "/opt/ml/checkpoints"
    
    if args.args_file is not None:
        with open(args.args_file, 'r') as fd:
            hyperparameters = yaml.safe_load(fd)
        hyperparameters['s3_path'] = os.path.join(hyperparameters['s3_path'], args.name)
    else:
        hyperparameters = {}
        
    distribution={
        "torch_distributed": {
            "enabled": True,
        }
    }
    code_root = "/opt/ml/code"
    environment = {
        "PYTHONPATH": ":".join([
            f"{code_root}/externals/vidar/",
            f"{code_root}/externals/vidar/externals/webdataset/",
            f"{code_root}/externals/vidar/externals/cv_datasets/",
            "$PYTHONPATH"
        ]),
        "AWS_ACCESS_KEY_ID": os.environ.get('AWS_ACCESS_KEY_ID', ''),
        "AWS_SECRET_ACCESS_KEY": os.environ.get('AWS_SECRET_ACCESS_KEY', ''),
        "WANDB_API_KEY": os.environ.get('WANDB_API_KEY', ''),
        "WANDB_ENTITY": os.environ.get('WANDB_ENTITY', ''),
        "WANDB__SERVICE_WAIT": "300",
        "INSTANCE_COUNT": str(args.instance_count),
        "AWS_S3_USE_CRT": "1",
        "NCCL_DEBUG": "INFO",
        "TOKENIZERS_PARALLELISM": "false",
        "SAGEMAKER_PROGRAM": "externals/vidar/scripts/datasets/webdataset/create_multinode_balanced.py"
    }

    if args.name is not None:
        environment['JOB_NAME'] = args.name

    security_group_ids = {
        'us-east-1': [
            'sg-0afb9fb0e79a54061'
        ],
        'us-west-2': [
            'sg-029d1d476bc087e31',
        ],
    }
    subnets = {
        'us-east-1': [
            'subnet-07bf42d7c9cb929e4',
            'subnet-0e260ba29726b9fbb',
        ],
        'us-west-2': [
            'subnet-0610f766a4cd5cdae', 
            'subnet-029adfb9e225d68f8',
            'subnet-01cc1bfeaf20155b5',
        ]
    }

    print()
    print()
    print('#############################################################')
    print(f'SageMaker Execution Role:       {role}')
    print(f'The name of the Execution role: {role_name[-1]}')
    print(f'SM Queue:                       {is_sm_queue}-{args.priority}-{args.fss_identifier}')
    print(f'AWS region:                     {region}')
    print(f'AWS profile:                    {args.profile}')
    print(f'AWS account:                    {account}')
    print(f'Entry point:                    {entry_point}')
    print(f'Image uri:                      {image_uri}')
    print(f'Job name:                       {job_name}')
    print(f'Configuration file:             {hyperparameters}')
    print(f'Instance count:                 {instance_count}')
    print(f'Instance type:                  {instance_type}')
    print(f'Queue:                          {args.queue}')
    print('#############################################################')
    print()
    print()
 
    estimator = PyTorch(
        entry_point=entry_point,
        sagemaker_session=sagemaker_session,
        base_job_name=base_job_name,
        hyperparameters=hyperparameters,
        role=role,
        image_uri=image_uri,
        instance_count=instance_count,
        instance_type=instance_type,
        train_use_spot_instances=train_use_spot_instances,
        output_path=output_s3,
        job_name=job_name,
        checkpoint_s3_uri=checkpoint_s3_uri,
        checkpoint_local_path=checkpoint_local_path,
        code_location=output_s3,
        distribution=distribution,
        max_run=max_run,
        max_wait=max_wait,
        environment=environment,
        keep_alive_period_in_seconds=keep_alive_period_in_seconds,
        tags=tags,
        subnets=subnets[region],
        security_group_ids=security_group_ids[region],
        train_volume_size=1000,
        enable_sagemaker_metrics=True, 
        logs=True, 
    )

    estimator.fit()


if __name__ == "__main__":
    main()