import argparse
import contextlib
import csv
import logging
import os
import random
import subprocess
import tempfile
from typing import Callable, Dict, Iterable, List

import numpy as np

import ray
from ray.experimental.raysort import constants, logging_utils, sortlib, tracing_utils
from ray.experimental.raysort.types import (
    BlockInfo,
    ByteCount,
    PartId,
    PartInfo,
    Path,
    RecordCount,
)
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

Args = argparse.Namespace

# ------------------------------------------------------------
#     Parse Arguments
# ------------------------------------------------------------


def get_args(*args, **kwargs):
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--ray_address",
        default="auto",
        type=str,
        help="if set to None, will launch a local Ray cluster",
    )
    parser.add_argument(
        "--total_data_size",
        default=1 * 1000 * 1024 * 1024 * 1024,
        type=ByteCount,
        help="total data size in bytes",
    )
    parser.add_argument(
        "--num_mappers",
        default=256,
        type=int,
        help="number of map tasks",
    )
    parser.add_argument(
        "--num_mappers_per_round",
        default=16,
        type=int,
        help="number of map tasks per first-stage merge tasks",
    )
    parser.add_argument(
        "--num_reducers",
        default=16,
        type=int,
        help="number of second-stage reduce tasks",
    )
    parser.add_argument(
        "--num_concurrent_rounds",
        default=4,
        type=int,
        help="max number of rounds of map/merge tasks in flight",
    )
    parser.add_argument(
        "--reducer_input_chunk",
        default=100 * 1024 * 1024,
        type=ByteCount,
        help="bytes to read from each file in reduce tasks",
    )
    parser.add_argument(
        "--skip_sorting",
        default=False,
        action="store_true",
        help="if set, no sorting is actually performed",
    )
    parser.add_argument(
        "--skip_input",
        default=False,
        action="store_true",
        help="if set, mappers will not read data from disk",
    )
    parser.add_argument(
        "--skip_output",
        default=False,
        action="store_true",
        help="if set, reducers will not write out results to disk",
    )
    # Which tasks to run?
    tasks_group = parser.add_argument_group(
        "tasks to run", "if no task is specified, will run all tasks"
    )
    tasks = ["generate_input", "sort", "validate_output"]
    for task in tasks:
        tasks_group.add_argument(f"--{task}", action="store_true")

    args = parser.parse_args(*args, **kwargs)
    # Derive additional arguments.
    args.input_part_size = ByteCount(args.total_data_size / args.num_mappers)
    assert args.num_mappers % args.num_mappers_per_round == 0
    args.num_rounds = int(args.num_mappers / args.num_mappers_per_round)
    args.mount_points = _get_mount_points()
    # If no tasks are specified, run all tasks.
    args_dict = vars(args)
    if not any(args_dict[task] for task in tasks):
        for task in tasks:
            args_dict[task] = True
    return args


def _get_mount_points():
    default_ret = [tempfile.gettempdir()]
    mnt = "/mnt"
    if os.path.exists(mnt):
        ret = [os.path.join(mnt, d) for d in os.listdir(mnt)]
        if len(ret) > 0:
            return ret
    return default_ret


# ------------------------------------------------------------
#     Generate Input
# ------------------------------------------------------------


def _part_info(args: Args, part_id: PartId, kind="input") -> PartInfo:
    node = ray._private.worker.global_worker.node_ip_address
    mnt = random.choice(args.mount_points)
    filepath = _get_part_path(mnt, part_id, kind)
    return PartInfo(part_id, node, filepath)


def _get_part_path(mnt: Path, part_id: PartId, kind="input") -> Path:
    assert kind in {"input", "output", "temp"}
    dir_fmt = constants.DATA_DIR_FMT[kind]
    dirpath = dir_fmt.format(mnt=mnt)
    os.makedirs(dirpath, exist_ok=True)
    filename_fmt = constants.FILENAME_FMT[kind]
    filename = filename_fmt.format(part_id=part_id)
    filepath = os.path.join(dirpath, filename)
    return filepath


@ray.remote
def generate_part(
    args: Args, part_id: PartId, size: RecordCount, offset: RecordCount
) -> PartInfo:
    logging_utils.init()
    pinfo = _part_info(args, part_id)
    subprocess.run(
        [constants.GENSORT_PATH, f"-b{offset}", f"{size}", pinfo.path], check=True
    )
    logging.info(f"Generated input {pinfo}")
    return pinfo


def generate_input(args: Args):
    if args.skip_input:
        return
    size = constants.bytes_to_records(args.input_part_size)
    offset = 0
    tasks = []
    for part_id in range(args.num_mappers):
        tasks.append(generate_part.remote(args, part_id, size, offset))
        offset += size
    assert offset == constants.bytes_to_records(args.total_data_size), args
    logging.info(f"Generating {len(tasks)} partitions")
    parts = ray.get(tasks)
    with open(constants.INPUT_MANIFEST_FILE, "w") as fout:
        writer = csv.writer(fout)
        writer.writerows(parts)


# ------------------------------------------------------------
#     Sort
# ------------------------------------------------------------


def _load_manifest(args: Args, path: Path) -> List[PartInfo]:
    if args.skip_input:
        return [PartInfo(i, None, None) for i in range(args.num_mappers)]
    with open(path) as fin:
        reader = csv.reader(fin)
        return [PartInfo(int(part_id), node, path) for part_id, node, path in reader]


def _load_partition(args: Args, path: Path) -> np.ndarray:
    if args.skip_input:
        return np.frombuffer(
            np.random.bytes(args.input_part_size), dtype=np.uint8
        ).copy()
    return np.fromfile(path, dtype=np.uint8)


def _dummy_sort_and_partition(
    part: np.ndarray, boundaries: List[int]
) -> List[BlockInfo]:
    N = len(boundaries)
    offset = 0
    size = int(np.ceil(part.size / N))
    blocks = []
    for _ in range(N):
        blocks.append((offset, size))
        offset += size
    return blocks


@ray.remote
@tracing_utils.timeit("map")
def mapper(
    args: Args, mapper_id: PartId, boundaries: List[int], path: Path
) -> List[np.ndarray]:
    logging_utils.init()
    part = _load_partition(args, path)
    sort_fn = (
        _dummy_sort_and_partition if args.skip_sorting else sortlib.sort_and_partition
    )
    blocks = sort_fn(part, boundaries)
    return [part[offset : offset + size] for offset, size in blocks]


def _dummy_merge(
    num_blocks: int, _n: int, get_block: Callable[[int, int], np.ndarray]
) -> Iterable[np.ndarray]:
    blocks = [((i, 0), get_block(i, 0)) for i in range(num_blocks)]
    while len(blocks) > 0:
        (m, d), block = blocks.pop(random.randrange(len(blocks)))
        yield block
        d_ = d + 1
        block = get_block(m, d_)
        if block is None:
            continue
        blocks.append(((m, d_), block))


def _merge_impl(
    args: Args,
    M: int,
    pinfo: PartInfo,
    get_block: Callable[[int, int], np.ndarray],
    skip_output=False,
):
    merge_fn = _dummy_merge if args.skip_sorting else sortlib.merge_partitions
    merger = merge_fn(M, get_block)

    if skip_output:
        for datachunk in merger:
            del datachunk
    else:
        with open(pinfo.path, "wb") as fout:
            for datachunk in merger:
                fout.write(datachunk)
    return pinfo


# See worker_placement_groups() for why `num_cpus=0`.
@ray.remote(num_cpus=0, resources={"worker": 1})
@tracing_utils.timeit("merge")
def merge_mapper_blocks(
    args: Args, reducer_id: PartId, mapper_id: PartId, *blocks: List[np.ndarray]
) -> PartInfo:
    part_id = constants.merge_part_ids(reducer_id, mapper_id)
    pinfo = _part_info(args, part_id, kind="temp")
    M = len(blocks)

    def get_block(i, d):
        if i >= M or d > 0:
            return None
        return blocks[i]

    return _merge_impl(args, M, pinfo, get_block)


# See worker_placement_groups() for why `num_cpus=0`.
@ray.remote(num_cpus=0, resources={"worker": 1})
@tracing_utils.timeit("reduce")
def final_merge(
    args: Args, reducer_id: PartId, *merged_parts: List[PartInfo]
) -> PartInfo:
    M = len(merged_parts)

    def _load_block_chunk(pinfo: PartInfo, d: int) -> np.ndarray:
        return np.fromfile(
            pinfo.path,
            dtype=np.uint8,
            count=args.reducer_input_chunk,
            offset=d * args.reducer_input_chunk,
        )

    def get_block(i, d):
        ret = _load_block_chunk(merged_parts[i], d)
        if ret.size == 0:
            return None
        return ret

    pinfo = _part_info(args, reducer_id, "output")
    return _merge_impl(args, M, pinfo, get_block, args.skip_output)


def _node_res(node: str) -> Dict[str, float]:
    return {"resources": {f"node:{node}": 1e-3}}


@contextlib.contextmanager
def worker_placement_groups(args: Args) -> List[ray.PlacementGroupID]:
    """
    Returns one placement group per node with a `worker` resource. To run
    tasks in the placement group, use
    `@ray.remote(num_cpus=0, resources={"worker": 1})`. Ray does not
    automatically reserve CPU resources, so tasks must specify `num_cpus=0`
    in order to run in a placement group.
    """
    pgs = [ray.util.placement_group([{"worker": 1}]) for _ in range(args.num_reducers)]
    ray.get([pg.ready() for pg in pgs])
    try:
        yield pgs
    finally:
        for pg in pgs:
            ray.util.remove_placement_group(pg)


@tracing_utils.timeit("sort", report_time=True)
def sort_main(args: Args):
    parts = _load_manifest(args, constants.INPUT_MANIFEST_FILE)
    assert len(parts) == args.num_mappers
    boundaries = sortlib.get_boundaries(args.num_reducers)

    # The exception of 'ValueError("Resource quantities >1 must be whole numbers.")'
    # will be raised if the `num_cpus` > 1 and not an integer.
    num_cpus = os.cpu_count() / args.num_concurrent_rounds
    if num_cpus > 1.0:
        num_cpus = int(num_cpus)
    mapper_opt = {
        "num_returns": args.num_reducers,
        "num_cpus": num_cpus,
    }  # Load balance across worker nodes by setting `num_cpus`.
    merge_results = np.empty((args.num_rounds, args.num_reducers), dtype=object)

    part_id = 0
    with worker_placement_groups(args) as pgs:
        for round in range(args.num_rounds):
            # Limit the number of in-flight rounds.
            num_extra_rounds = round - args.num_concurrent_rounds + 1
            if num_extra_rounds > 0:
                ray.wait(
                    [f for f in merge_results.flatten() if f is not None],
                    num_returns=num_extra_rounds * args.num_reducers,
                )

            # Submit map tasks.
            mapper_results = np.empty(
                (args.num_mappers_per_round, args.num_reducers), dtype=object
            )
            for _ in range(args.num_mappers_per_round):
                _, node, path = parts[part_id]
                m = part_id % args.num_mappers_per_round
                mapper_results[m, :] = mapper.options(**mapper_opt).remote(
                    args, part_id, boundaries, path
                )
                part_id += 1

            # Submit merge tasks.
            merge_results[round, :] = [
                merge_mapper_blocks.options(
                    scheduling_strategy=PlacementGroupSchedulingStrategy(
                        placement_group=pgs[r]
                    )
                ).remote(args, r, round, *mapper_results[:, r].tolist())
                for r in range(args.num_reducers)
            ]

            # Delete local references to mapper results.
            mapper_results = None

        # Submit second-stage reduce tasks.
        reducer_results = [
            final_merge.options(
                scheduling_strategy=PlacementGroupSchedulingStrategy(
                    placement_group=pgs[r]
                )
            ).remote(args, r, *merge_results[:, r].tolist())
            for r in range(args.num_reducers)
        ]
        reducer_results = ray.get(reducer_results)

    if not args.skip_output:
        with open(constants.OUTPUT_MANIFEST_FILE, "w") as fout:
            writer = csv.writer(fout)
            writer.writerows(reducer_results)

    logging.info(ray._private.internal_api.memory_summary(stats_only=True))


# ------------------------------------------------------------
#     Validate Output
# ------------------------------------------------------------


def _run_valsort(args: List[str]):
    proc = subprocess.run([constants.VALSORT_PATH] + args, capture_output=True)
    if proc.returncode != 0:
        logging.critical("\n" + proc.stderr.decode("ascii"))
        raise RuntimeError(f"Validation failed: {args}")


@ray.remote
def validate_part(path: Path):
    logging_utils.init()
    sum_path = path + ".sum"
    _run_valsort(["-o", sum_path, path])
    logging.info(f"Validated output {path}")
    with open(sum_path, "rb") as fin:
        return os.path.getsize(path), fin.read()


def validate_output(args: Args):
    if args.skip_sorting or args.skip_output:
        return
    partitions = _load_manifest(args, constants.OUTPUT_MANIFEST_FILE)
    results = []
    for _, node, path in partitions:
        results.append(validate_part.options(**_node_res(node)).remote(path))
    logging.info(f"Validating {len(results)} partitions")
    results = ray.get(results)
    total = sum(s for s, _ in results)
    assert total == args.total_data_size, total - args.total_data_size
    all_checksum = b"".join(c for _, c in results)
    with tempfile.NamedTemporaryFile() as fout:
        fout.write(all_checksum)
        fout.flush()
        _run_valsort(["-s", fout.name])
    logging.info("All OK!")


# ------------------------------------------------------------
#     Main
# ------------------------------------------------------------


def init(args: Args):
    if not args.ray_address:
        ray.init(resources={"worker": os.cpu_count()})
    else:
        ray.init(address=args.ray_address)
    logging_utils.init()
    logging.info(args)
    os.makedirs(constants.WORK_DIR, exist_ok=True)
    resources = ray.cluster_resources()
    logging.info(resources)
    args.num_workers = resources["worker"]
    progress_tracker = tracing_utils.create_progress_tracker(args)
    return progress_tracker


def main(args: Args):
    # Keep the actor handle in scope for the duration of the program.
    _progress_tracker = init(args)  # noqa F841

    if args.generate_input:
        generate_input(args)

    if args.sort:
        sort_main(args)

    if args.validate_output:
        validate_output(args)


if __name__ == "__main__":
    main(get_args())
