import argparse
import asyncio
import json
import logging
import os
import re
import shlex

import toml
from ray.job_submission import JobSubmissionClient

from X.run.vars import ENV_VARS, PIP_DEPS, REMOTE_DASHBOARD_URL

# Setup logger
logger = logging.getLogger("ray")


def parse_pip_requirements(line: str) -> list[str]:
    # The pattern means:
    #   (?:\[[^\]]*\]|[^,])+
    #   Match either:
    #       - a bracketed chunk: [anything but ]]
    #       - OR any character that isn't a comma
    #   Repeat 1 or more times.
    # So each match is all characters up to the next top-level comma.
    # For example, the input string "numpy==2.0.0,scipy[extras1,extras2],sympy"
    # is parsed as ["numpy==2.0.0", "scipy[extras1,extras2]", "sympy"]
    pattern = r"(?:\[[^\]]*\]|[^,])+"
    return re.findall(pattern, line)


def generate_pythonpath(base_dir="submodules"):
    # List to hold all the paths
    paths = []

    if not os.path.exists(base_dir):
        logger.warning(f"Base directory {base_dir} does not exist.")
        return ""

    # Iterate through the directories inside submodules
    for submodule in os.listdir(base_dir):
        submodule_path = os.path.join(base_dir, submodule)

        # Check if it's a directory
        if os.path.isdir(submodule_path):
            # Add both submodule and submodule/src paths
            paths.append(submodule_path)
            src_path = os.path.join(submodule_path, "src")
            if os.path.isdir(src_path):
                paths.append(src_path)

    # Add "." for the current directory to make sure the imports are working properly
    # Join the paths with ':'
    pythonpath = ":".join(paths)
    return pythonpath


def get_dependencies_from_toml(toml_file: str) -> list:
    """Extract dependencies from the pyproject.toml file and return them as a list."""
    try:
        parsed_toml = toml.load(toml_file)
        dependencies = parsed_toml.get("project", {}).get("dependencies", [])
        dependencies = _remove_problematic_deps(dependencies)
        logger.info(f"Dependencies extracted: {dependencies}")
        return dependencies
    except FileNotFoundError:
        logger.error(f"File {toml_file} not found.")
        return []
    except toml.TomlDecodeError:
        logger.error(f"Failed to parse {toml_file}.")
        return []


def _remove_problematic_deps(dependencies: list[str]):
    out: list[str] = []
    # remove ray from dependencies. We do this because Ray gets mad if you try to install another version of Ray
    expr = re.compile(r"^\s*ray([^a-zA-Z0-9_]|$)")
    for dep in dependencies:
        if not expr.match(dep):
            out.append(dep)
        else:
            logger.debug(f"Skipping dependency: {dep}")
    return out


def tpus_per_node(tpu_type: str) -> int:
    """Return the number of TPU chips per node for a given TPU type."""
    if tpu_type in {"v4-8", "v5p-8"}:
        return 4
    match = re.search(r"-(\d+)$", tpu_type)
    if not match:
        raise ValueError(f"Cannot parse TPU type: {tpu_type}")
    chips = int(match.group(1))
    if chips > 8:
        raise ValueError("Only single tpu nodes are supported with the CLI")
    return chips


async def submit_and_track_job(
    entrypoint: str,
    dependencies: list,
    env_vars: dict,
    no_wait: bool,
    *,
    entrypoint_num_cpus: float | None = None,
    entrypoint_num_gpus: float | None = None,
    entrypoint_memory: int | None = None,
    entrypoint_resources: dict | None = None,
):
    """Submit a job to Ray and optionally track logs."""

    current_dir = os.getcwd()
    client = JobSubmissionClient(REMOTE_DASHBOARD_URL)
    runtime_dict = {
        "pip": dependencies,
        "working_dir": current_dir,
        "env_vars": env_vars,
        "config": {"setup_timeout_seconds": 1800},
    }

    if len(dependencies) == 0:
        del runtime_dict["pip"]

    logger.info(f"Submitting job with entrypoint: {entrypoint}")
    logger.info(f"Dependencies: {json.dumps(dependencies, indent=4)}")
    logger.info(f"env_vars: {json.dumps(env_vars, indent=4)}")

    logger.info(
        f"Terminal command: \n" f"ray job submit " f"--runtime-env-json '{json.dumps(runtime_dict)}'" f" -- {entrypoint}"
    )
    # Submit the job with runtime environment and entrypoint
    submission_id = client.submit_job(
        entrypoint=entrypoint,
        runtime_env=runtime_dict,
        entrypoint_num_cpus=entrypoint_num_cpus,
        entrypoint_num_gpus=entrypoint_num_gpus,
        entrypoint_memory=entrypoint_memory,
        entrypoint_resources=entrypoint_resources,
    )
    logger.info(f"Job submitted with ID: {submission_id}")
    logger.info(f"Job URL: http://localhost:8265/#/jobs/{submission_id}")

    if no_wait:
        return

    # Stream logs asynchronously
    async for lines in client.tail_job_logs(submission_id):
        print(lines, end="")


def main():
    """Parse command-line arguments and submit the job."""
    parser = argparse.ArgumentParser(description="Submit Ray jobs using the command-line.")
    parser.add_argument("--no_wait", action="store_true", help="Do not wait for the job to finish.")
    parser.add_argument(
        "--env_vars",
        "-e",
        action="append",
        nargs="+",
        metavar=("KEY", "VALUE"),
        help="Set environment variables for the job. If only a KEY is provided, "
        "the VALUE will be set to an empty string.",
    )
    parser.add_argument(
        "--pip_deps", type=parse_pip_requirements, help="List of pip dependencies to " "install before running."
    )
    parser.add_argument(
        "--entrypoint-num-cpus",
        type=float,
        default=None,
        help="Number of CPUs to reserve for the entrypoint command.",
    )
    parser.add_argument(
        "--entrypoint-num-gpus",
        type=float,
        default=None,
        help="Number of GPUs to reserve for the entrypoint command.",
    )
    parser.add_argument(
        "--entrypoint-memory",
        type=int,
        default=None,
        help="Amount of memory to reserve for the entrypoint command.",
    )
    parser.add_argument(
        "--entrypoint-resources",
        type=json.loads,
        default=None,
        help="JSON dictionary describing resources to reserve for the entrypoint command.",
    )
    parser.add_argument(
        "--tpu",
        type=str,
        default=None,
        help="TPU type to reserve for the entrypoint (e.g. v4-8)",
    )
    parser.add_argument("cmd", help="The command to run in the Ray cluster.", nargs=argparse.REMAINDER)

    args = parser.parse_args()

    # Combine the remaining arguments to form the full command
    full_cmd = " ".join(shlex.quote(arg) for arg in args.cmd).strip()
    if not full_cmd.startswith("--"):
        logger.error("Command must start with '--'.")
        exit(1)
    full_cmd = full_cmd[2:]

    # Load and merge environment variables from multiple -e options
    env_vars = {}

    if args.env_vars:
        for item in args.env_vars:
            if len(item) > 2:
                logger.error(
                    f"Too many values provided for environment variable: {' '.join(item)}. "
                    f"Expected at most 2 (KEY VALUE)."
                )
                exit(1)
            elif len(item) == 1:
                # If only the key is provided, set its value to an empty string
                if "=" in item[0]:
                    logger.error(
                        f"Invalid key provided for environment variable: {' '.join(item)}. "
                        f"Key should not contain '='.\n\n"
                        f"You probably meant to do '-e {' '.join(item[0].split('='))}'."
                    )
                    exit(1)
                env_vars[item[0]] = ""
            elif len(item) == 2:
                # If both key and value are provided, store them as a key-value pair
                if "=" in item[0]:
                    logger.error(
                        f"Invalid key provided for environment variable: {' '.join(item)}. "
                        f"Key should not contain '='.\n\n"
                        f"You probably meant to do '-e {' '.join(item[0].split('='))}'."
                    )
                    exit(1)
                env_vars[item[0]] = item[1]

    # Now env_vars is a dictionary with the environment variables set using -e

    env_vars = {**ENV_VARS, **env_vars}

    env_vars["PYTHONPATH"] = generate_pythonpath() + ":" + env_vars.get("PYTHONPATH", "")

    # Convert pyproject.toml to requirements.txt before submission
    pyproject_toml = "pyproject.toml"
    # If we are using the latest docker image then we can skip getting core dependencies from pyproject.toml
    # As they are already installed inside the cluster
    dependencies = []
    dependencies += get_dependencies_from_toml(pyproject_toml)
    dependencies += PIP_DEPS
    dependencies += args.pip_deps if args.pip_deps else []

    entrypoint_resources = args.entrypoint_resources
    if args.tpu:
        try:
            chips = tpus_per_node(args.tpu)
        except ValueError as e:
            logger.error(str(e))
            exit(1)
        tpu_res = {f"TPU-{args.tpu}-head": 1, "TPU": chips}
        entrypoint_resources = (entrypoint_resources or {}) | tpu_res

    # Submit the job and track it asynchronously
    asyncio.run(
        submit_and_track_job(
            full_cmd,
            dependencies,
            env_vars,
            args.no_wait,
            entrypoint_num_cpus=args.entrypoint_num_cpus,
            entrypoint_num_gpus=args.entrypoint_num_gpus,
            entrypoint_memory=args.entrypoint_memory,
            entrypoint_resources=entrypoint_resources,
        )
    )


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
    main()
