"""Job run command."""

import json
from typing import Dict, List, Mapping, Optional, Sequence, Union

import rich_click as click

from lightning_sdk.cli.resource_completion import complete_studio
from lightning_sdk.cli.utils.json_output import echo_json
from lightning_sdk.cli.utils.logging import LightningCommand
from lightning_sdk.cli.utils.teamspace_option import resolve_teamspace, teamspace_option
from lightning_sdk.job import Job
from lightning_sdk.machine import Machine

_MACHINE_VALUES = tuple(
    [machine.name for machine in Machine.__dict__.values() if isinstance(machine, Machine) and machine._include_in_cli]
)


@click.command("run", cls=LightningCommand)
@click.option("--name", default=None, help="The name of the job. Needs to be unique within the teamspace.")
@click.option(
    "--num-machines",
    "--num_machines",
    default=1,
    show_default=True,
    type=click.IntRange(min=1),
    help="The number of machines to run on.",
)
@click.option(
    "--machine",
    default="CPU",
    show_default=True,
    type=click.Choice(_MACHINE_VALUES),
    help="The machine type to run the job on.",
)
@click.option(
    "--command",
    default=None,
    help=(
        "The command to run inside your job. "
        "Required if using a studio. "
        "Optional if using an image. "
        "If not provided for images, will run the container entrypoint and default command."
    ),
)
@click.option(
    "--studio",
    default=None,
    shell_complete=complete_studio,
    help=(
        "The studio env to run the job with. Mutually exclusive with image. "
        "If both --studio and --image are omitted while running inside a Studio, defaults to that Studio "
        "(if its teamspace matches --teamspace)."
    ),
)
@click.option("--image", default=None, help="The docker image to run the job with. Mutually exclusive with studio.")
@teamspace_option
@click.option(
    "--cloud",
    default=None,
    help=(
        "Cloud provider or cloud account to run the job on. Defaults to the studio cloud account or teamspace default."
    ),
)
@click.option(
    "--env",
    "-e",
    default=[""],
    help=("Environment variable to set inside the job. Should be of format KEY=VALUE"),
    multiple=True,
)
@click.option(
    "--interruptible",
    is_flag=True,
    flag_value=True,
    default=False,
    help="Whether the job should run on interruptible instances. They are cheaper but can be preempted.",
)
@click.option(
    "--image-credentials",
    "--image_credentials",
    default=None,
    help=(
        "The credentials used to pull the image. "
        "Required if the image is private. "
        "This should be the name of the respective credentials secret created on the Lightning AI platform."
    ),
)
@click.option(
    "--cloud-account-auth",
    "--cloud_account_auth",
    is_flag=True,
    default=False,
    help=(
        "Whether to authenticate with the cloud account to pull the image. "
        "Required if the registry is part of a cloud provider (e.g. ECR)."
    ),
)
@click.option(
    "--entrypoint",
    default=None,
    show_default=True,
    help=(
        "The entrypoint of your docker container. "
        "Default runs the provided command in a standard shell. "
        "To use the pre-defined entrypoint of the provided image, set this to an empty string. "
        "Only applicable when submitting docker jobs."
    ),
)
@click.option(
    "--path-mapping",
    "--path_mapping",
    default=[""],
    help=(
        "Maps path inside of containers to paths inside data-connections. "
        "Should be of form <CONTAINER_PATH_1>:<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1> and "
        "omitting the path inside the connection defaults to the connections root. "
        "Can be specified multiple times for multiple mappings"
    ),
    multiple=True,
)
@click.option(
    "--path-mappings",
    "--path_mappings",
    default="",
    help=(
        "Maps path inside of containers to paths inside data-connections. "
        "Should be a comma separated list of form: "
        "<MAPPING_1>,<MAPPING_2>,... "
        "where each mapping is of the form "
        "<CONTAINER_PATH_1>:<CONNECTION_NAME_1>:<PATH_WITHIN_CONNECTION_1> and "
        "omitting the path inside the connection defaults to the connections root. "
        "Instead of a comma-separated list, consider passing --path-mapping multiple times."
    ),
)
@click.option(
    "--tags",
    "--tag",
    "tags",
    default=(),
    multiple=True,
    help=(
        "Tag to apply to the job. "
        "Can be a comma-separated list or passed multiple times. "
        "Tags that don't exist in the teamspace yet are created."
    ),
)
@click.option("--json", "as_json", is_flag=True, default=False, help="Output the created job as JSON.")
@click.option(
    "--keep-minutes",
    "--keep_minutes",
    default=None,
    type=click.IntRange(min=0),
    help=("Minutes to keep the machine after the job stops. Omit or pass 0 to release the machine immediately."),
)
def run_job(
    name: Optional[str] = None,
    num_machines: int = 1,
    machine: str = "CPU",
    command: Optional[str] = None,
    studio: Optional[str] = None,
    image: Optional[str] = None,
    teamspace: Optional[str] = None,
    org: Optional[str] = None,
    user: Optional[str] = None,
    cloud: Optional[str] = None,
    env: Sequence[str] = (),
    interruptible: bool = False,
    image_credentials: Optional[str] = None,
    cloud_account_auth: bool = False,
    entrypoint: str = "sh -c",
    path_mapping: Sequence[str] = (),
    path_mappings: str = "",
    tags: Sequence[str] = (),
    as_json: bool = False,
    keep_minutes: Optional[int] = None,
) -> None:
    """Run async workloads using a docker image or studio.

    Pass --num-machines greater than 1 to run a multi-machine job.
    """
    if not name:
        from datetime import datetime

        timestr = datetime.now().strftime("%b-%d-%H_%M")
        name = f"job-{timestr}"

    machine_enum: Union[str, Machine]
    try:
        machine_enum = getattr(Machine, machine.upper(), Machine(machine, machine))
    except KeyError:
        machine_enum = machine

    resolved_teamspace = resolve_teamspace(teamspace, org, user)

    path_mappings_dict = _resolve_path_mapping(path_mappings=path_mappings)
    for mapping in path_mapping:
        path_mappings_dict.update(_resolve_path_mapping(path_mappings=mapping))

    env_dict = {}
    for value in env:
        env_dict.update(_resolve_envs(value))

    job = Job.run(
        name=name,
        machine=machine_enum,
        command=command,
        studio=studio,
        image=image,
        teamspace=resolved_teamspace,
        org=org,
        user=user,
        cloud=cloud,
        env=env_dict,
        interruptible=interruptible,
        image_credentials=image_credentials,
        cloud_account_auth=cloud_account_auth,
        entrypoint=entrypoint,
        path_mappings=path_mappings_dict,
        num_machines=num_machines,
        tags=_resolve_tags(tags),
        keep_minutes=keep_minutes,
    )

    if as_json:
        echo_json(
            {
                "id": job.id,
                "name": job.name,
                "teamspace": f"{resolved_teamspace.owner.name}/{resolved_teamspace.name}",
                "link": job.link,
            }
        )
        return

    click.echo(f"Submitted job {job.name}.")


def _resolve_path_mapping(path_mappings: str) -> Dict[str, str]:
    path_mappings = path_mappings.strip()

    if not path_mappings:
        return {}

    path_mappings_dict = {}
    for mapping in path_mappings.split(","):
        if not mapping.strip():
            continue

        splits = str(mapping).split(":", 1)
        if len(splits) != 2:
            raise RuntimeError(
                "Mapping needs to be of form <CONTAINER_PATH>:<CONNECTION_NAME>[:<PATH_WITHIN_CONNECTION>], "
                f"but got {mapping}"
            )

        path_mappings_dict[splits[0].strip()] = splits[1].strip()

    return path_mappings_dict


def _resolve_tags(tags: Sequence[str]) -> List[str]:
    """Flatten repeated and comma-separated --tags values into the list the SDK expects."""
    resolved: List[str] = []
    for value in tags:
        resolved.extend(tag.strip() for tag in value.split(",") if tag.strip())

    return resolved


def _resolve_envs(envs: str) -> Dict[str, str]:
    if not envs:
        return {}

    try:
        env_dict = json.loads(envs)
        if isinstance(env_dict, Mapping):
            return dict(env_dict)

        raise ValueError(f"Env {envs} cannot be parsed as environment variable")
    except json.decoder.JSONDecodeError as error:
        env_dict = {}
        splits = envs.split("=", 1)
        if len(splits) == 2:
            key, value = splits
            env_dict.update({key: value})
            return env_dict

        raise ValueError(f"Env {envs} cannot be parsed as environment variable: {error!s}") from error