# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
from azure.ai.ml.entities import Environment, CommandJob
import os
from omegaconf import OmegaConf
import hydra
from hydra.core.hydra_config import HydraConfig
from shrike.pipeline.aml_connect import azureml_connect
from shrike._core import (
    is_eyesoff_helper,
    b64_decode,
    b64_encode,
    experimental,
    O365_FEED,
)
from shrike.pipeline import AMLPipelineHelper


def merge_serialize_config_with_default(serialized_config):
    default_config = AMLPipelineHelper._default_config()
    cfg = OmegaConf.merge(
        default_config, OmegaConf.structured(b64_decode(serialized_config))
    )
    return cfg


def get_or_create_env(ml_client, env_name, is_eyesoff):
    try:
        env = ml_client.environemnts.get(name=env_name)
    except Exception:
        pip_packages = [
            "ray>=1.7.0",
            "shrike[pipeline]>=1.28.0",
        ]
        if is_eyesoff:
            pip_packages.append(f"--index-url {O365_FEED}")
        conda_file_dict = {
            "name": env_name,
            "dependencies": [
                "python=3.8",
                "pip",
                {
                    "pip": pip_packages,
                },
            ],
        }
        env = Environment(
            name=env_name,
            conda_file=conda_file_dict,
        )

    return env


@experimental()
@hydra.main(config_name="default")
def ray_actor_on_shrike(cfg):
    default_config = AMLPipelineHelper._default_config()
    cfg = OmegaConf.merge(default_config, cfg)
    env_name = "ray_env"

    aml = cfg.aml
    ml_client = azureml_connect(
        aml_subscription_id=aml.subscription_id,
        aml_resource_group=aml.resource_group,
        aml_workspace_name=aml.workspace_name,
        aml_auth=aml.get("auth", "interactive"),
        aml_tenant=aml.get("tenant", None),
        aml_force=aml.get("force", False),
    )
    env = get_or_create_env(
        ml_client,
        env_name,
        is_eyesoff_helper(aml.get("tenant", None), aml.subscription_id),
    )
    compute_target = cfg.compute.default_compute_target

    script = os.path.join(HydraConfig.get().runtime.cwd, cfg.run.script)
    with open(script) as file:
        contents = file.read()

    command_job = CommandJob(
        code=".",
        command=f"echo '{contents}' > script.py && python script.py --serialized_config {b64_encode(OmegaConf.to_yaml(cfg))}",
        compute=compute_target,
        environment=env,
    )
    job = ml_client.jobs.create_or_update(
        command_job, experiment_name=cfg.run.experiment_name
    )
    print(job.studio_url)
    ml_client.jobs.stream(job.name)
