import random
import shutil
import time
import typing as T
from collections import defaultdict
from pathlib import Path

import cumulusci.core.exceptions as exc
from cumulusci.core.config import OrgConfig
from cumulusci.tasks.bulkdata.generate_from_yaml import GenerateDataFromYaml
from cumulusci.tasks.bulkdata.load import LoadData
from cumulusci.utils.parallel.task_worker_queues.parallel_worker_queue import (
    WorkerQueue,
    WorkerQueueConfig,
)

from .snowfakery_run_until import PortionGenerator
from .snowfakery_working_directory import SnowfakeryWorkingDirectory
from .subtask_configurator import SubtaskConfigurator

# number of portions we will allow to be on-disk waiting to be loaded
# higher numbers use more disk space.
LOAD_QUEUE_SIZE = 15
# more loader workers than generators because they spend so much time
# waiting for responses. 4:1 is an experimentally derived ratio
WORKER_TO_LOADER_RATIO = 4


class SnowfakeryChannelManager:
    """The channels describe the 'shape' of the communication

    The normal case is a single, parallelized, bulk channel,
    multi-threaded on client and server, using a single user
    account.

    Using .load.yml you can add more channels, utilizing
    more user accounts which can speed up throughput in
    a few cases."""

    def __init__(
        self,
        subtask_configurator,
        *,
        project_config,
        logger,
    ):
        self.results_reporter = WorkerQueue.context.Queue()
        self.channels = []
        self.project_config = project_config
        self.logger = logger
        self.subtask_configurator = subtask_configurator
        self.start_time = time.time()

    def add_channel(
        self,
        org_config: OrgConfig,
        num_generator_workers: int,
        num_loader_workers: T.Optional[int],
        working_directory: Path,
        recipe_options: dict,
    ):
        if not num_loader_workers:
            num_loader_workers = num_generator_workers * WORKER_TO_LOADER_RATIO

        self.channels.append(
            Channel(
                project_config=self.project_config,
                org_config=org_config,
                num_generator_workers=num_generator_workers,
                num_loader_workers=num_loader_workers,
                working_directory=working_directory,
                subtask_configurator=self.subtask_configurator,
                logger=self.logger,
                results_reporter=self.results_reporter,
                recipe_options=recipe_options,
            )
        )

    def tick(
        self,
        upload_status,
        template_path: Path,
        tempdir: Path,
        portions: PortionGenerator,
        get_upload_status: T.Callable,
    ):
        """Called every few seconds, to make new data generators if needed."""
        # advance each channel's queues
        for channel in self.channels:
            channel.tick()

        # populate channels that have space
        self.assign_work_to_channels(
            upload_status,
            template_path,
            tempdir,
            portions,
            get_upload_status,
        )

    def assign_work_to_channels(
        self,
        upload_status,
        template_path: Path,
        tempdir: Path,
        portions: PortionGenerator,
        get_upload_status: T.Callable,
    ):
        def channel_with_free_space():
            return [channel for channel in self.channels if not channel.full]

        new_workers = [True]  # initial value to get into the while loop
        channels = channel_with_free_space()
        while (
            channels
            and any(new_workers)
            and not portions.done(upload_status.total_sets_working_on_or_uploaded)
        ):
            # randomize channels so the first one doesn't always grab all slots
            random.shuffle(channels)
            new_workers = [
                channel.make_new_worker(
                    template_path, tempdir, portions, get_upload_status
                )
                for channel in channels
            ]
            channels = channel_with_free_space()

    def get_upload_status(self, batch_size, sets_finished_while_generating_template):
        """Combine information from the different data sources into a single "report".

        Useful for debugging but also for making decisions about what to do next."""
        summed_statuses = defaultdict(int)
        for channel in self.channels:
            for key, value in channel.get_upload_status_for_channel().items():
                summed_statuses[key] += value

        summed_statuses["sets_finished"] += sets_finished_while_generating_template
        return UploadStatus(
            target_count=self.subtask_configurator.run_until.gap,
            elapsed_seconds=int(self.elapsed_seconds()),
            batch_size=batch_size,
            channels=len(self.channels),
            **summed_statuses,
        )

    def elapsed_seconds(self):
        return time.time() - self.start_time

    def failure_descriptions(self) -> T.List[str]:
        """Log failures from sub-processes to main process"""
        ret = []
        for channel in self.channels:
            ret.extend(channel.failure_descriptions())
        return ret

    def check_finished(self) -> bool:
        for channel in self.channels:
            channel.tick()
        return all([channel.check_finished() for channel in self.channels])

    def get_results_report(self, block=False):
        return self.results_reporter.get(block=block)


class Channel:
    def __init__(
        self,
        *,
        project_config,
        org_config,
        num_generator_workers: int,
        num_loader_workers: int,
        working_directory: Path,
        subtask_configurator: SubtaskConfigurator,
        logger,
        recipe_options=None,
        results_reporter=None,
    ):
        self.project_config = project_config
        self.org_config = org_config
        self.num_generator_workers = num_generator_workers
        self.num_loader_workers = num_loader_workers
        self.working_directory = working_directory
        self.subtask_configurator = subtask_configurator
        self.run_until = subtask_configurator.run_until
        self.logger = logger
        self.results_reporter = results_reporter
        self.job_counter = 0
        recipe_options = recipe_options or {}
        self._configure_queues(recipe_options)

    def _configure_queues(self, recipe_options):
        """Configure two ParallelWorkerQueues for datagen and dataload"""
        try:
            # in the future, the connected_app should come from the org
            connected_app = self.project_config.keychain.get_service("connected_app")
        except exc.ServiceNotConfigured:  # pragma: no cover
            # to discuss...when can this happen? What are the consequences?
            connected_app = None

        def data_generator_opts_callback(*args, **kwargs):
            return self.subtask_configurator.data_generator_opts(
                *args, recipe_options=recipe_options, **kwargs
            )

        data_gen_q_config = WorkerQueueConfig(
            project_config=self.project_config,
            org_config=self.org_config,
            connected_app=connected_app,
            redirect_logging=True,
            # processes are better for compute-heavy tasks (in Python)
            spawn_class=WorkerQueue.Process,
            parent_dir=self.working_directory,
            name="data_gen",
            task_class=GenerateDataFromYaml,
            make_task_options=data_generator_opts_callback,
            queue_size=0,
            num_workers=self.num_generator_workers,
        )
        self.data_gen_q = WorkerQueue(data_gen_q_config)

        load_data_q_config = WorkerQueueConfig(
            project_config=self.project_config,
            org_config=self.org_config,
            connected_app=connected_app,
            redirect_logging=True,
            spawn_class=WorkerQueue.Thread,
            parent_dir=self.working_directory,
            name="data_load",
            task_class=LoadData,
            make_task_options=self.subtask_configurator.data_loader_opts,
            queue_size=LOAD_QUEUE_SIZE,
            num_workers=self.num_loader_workers,
            rename_directory=self.data_loader_new_directory_name,
        )
        self.load_data_q = WorkerQueue(load_data_q_config, self.results_reporter)

        self.data_gen_q.feeds_data_to(self.load_data_q)
        return self.data_gen_q, self.load_data_q

    def data_loader_new_directory_name(self, working_directory):
        return data_loader_new_directory_name(working_directory, self.run_until)

    def tick(
        self,
    ):
        """Called every few seconds, to make new data generators if needed."""
        self.data_gen_q.tick()

    @property
    def full(self):
        return self.data_gen_q.full

    def make_new_worker(
        self,
        template_path: Path,
        tempdir: Path,
        portions: PortionGenerator,
        get_upload_status: T.Callable,
    ):
        if (
            self.data_gen_q.num_free_workers and self.data_gen_q.full
        ):  # pragma: no cover
            # TODO: investigate the consequences of taking this branch out
            self.logger.info("Waiting before datagen (load queue is full)")
        else:
            upload_status = get_upload_status(
                portions.next_batch_size,
            )
            self.job_counter += 1
            batch_size = portions.next_batch(
                upload_status.total_sets_working_on_or_uploaded
            )
            if not batch_size:
                self.logger.info("All scheduled portions generated and being uploaded")
                return

            job_dir = self.generator_data_dir(
                self.job_counter, template_path, batch_size, tempdir
            )
            self.data_gen_q.push(job_dir)
            return job_dir

    def generator_data_dir(self, idx, template_path, batch_size, parent_dir):
        """Create a new generator directory with a name based on index and batch_size"""
        assert batch_size > 0
        data_dir = parent_dir / (str(idx) + "_" + str(batch_size))
        shutil.copytree(template_path, data_dir)
        return data_dir

    def get_upload_status_for_channel(self):
        def set_count_from_names(names):
            return sum(int(name.split("_")[1]) for name in names)

        return {
            "sets_queued_to_be_generated": set_count_from_names(
                self.data_gen_q.queued_jobs
            ),
            "sets_being_generated": set_count_from_names(
                self.data_gen_q.inprogress_jobs
            ),
            "sets_queued_for_loading": set_count_from_names(
                self.load_data_q.queued_jobs
            ),
            # note that these may count as already imported in the org
            "sets_being_loaded": set_count_from_names(self.load_data_q.inprogress_jobs),
            "max_num_loader_workers": self.num_loader_workers,
            "max_num_generator_workers": self.num_generator_workers,
            # todo: use row-level result from org load for better accuracy
            "sets_finished": set_count_from_names(self.load_data_q.outbox_jobs),
            "sets_failed": len(self.load_data_q.failed_jobs),
            # TODO: are these two redundant?
            "inprogress_generator_jobs": len(self.data_gen_q.inprogress_jobs),
            "inprogress_loader_jobs": len(self.load_data_q.inprogress_jobs),
            "data_gen_free_workers": self.data_gen_q.num_free_workers,
        }

    def failure_descriptions(self) -> T.List[str]:
        """Log failures from sub-processes to main process"""
        failure_dirs = set(
            self.load_data_q.failed_job_dirs + self.data_gen_q.failed_job_dirs
        )

        def error_from_dir(failure_dir: Path) -> T.Optional[str]:
            f = Path(failure_dir) / "exception.txt"
            if not f.exists():  # pragma: no cover
                return
            return f.read_text(encoding="utf-8").strip().split("\n")[-1]

        errors = [error_from_dir(failure_dir) for failure_dir in failure_dirs]
        return [error for error in errors if error is not None]

    def check_finished(self) -> bool:
        self.data_gen_q.tick()
        still_running = len(self.data_gen_q.workers + self.load_data_q.workers) > 0
        return not still_running


# TODO: This function is actually based on the number generated,
#       because it is called before the load.
#       If there are row errors, it will drift out of correctness
#       Code needs to be updated to rename again after load.
#       Or move away from using directory names for math altogether.
def data_loader_new_directory_name(working_dir: Path, run_until):
    """Change the directory name to reflect the true number of sets created."""

    if not run_until.sobject_name:
        return working_dir

    wd = SnowfakeryWorkingDirectory(working_dir)
    count = wd.get_record_counts()[run_until.sobject_name]

    path, _ = str(working_dir).rsplit("_", 1)
    new_working_dir = Path(path + "_" + str(count))
    return new_working_dir


class UploadStatus(T.NamedTuple):
    """Single "report" of the current status of our processes."""

    batch_size: int
    sets_being_generated: int
    sets_queued_to_be_generated: int
    sets_being_loaded: int
    sets_queued_for_loading: int
    sets_finished: int
    target_count: int
    max_num_loader_workers: int
    max_num_generator_workers: int
    elapsed_seconds: int
    sets_failed: int
    inprogress_generator_jobs: int
    inprogress_loader_jobs: int
    data_gen_free_workers: int
    channels: int

    @property
    def total_in_flight(self):
        return (
            self.sets_queued_to_be_generated
            + self.sets_being_generated
            + self.sets_queued_for_loading
            + self.sets_being_loaded
        )

    @property
    def total_sets_working_on_or_uploaded(
        self,
    ):
        return self.total_in_flight + self.sets_finished

    def _display(self, detailed=False):
        most_important_stats = [
            "target_count",
            "total_sets_working_on_or_uploaded",
            "sets_finished",
            "sets_failed",
        ]
        if self.channels > 1:
            most_important_stats.append("channels")

        queue_stats = [
            "inprogress_generator_jobs",
            "inprogress_loader_jobs",
        ]

        def display_stats(keys):
            def format(a):
                return f"{a.replace('_', ' ').title()}: {getattr(self, a):,}"

            return (
                "\n"
                + "\n".join(
                    format(a)
                    for a in keys
                    if not a[0] == "_" and not callable(getattr(self, a))
                )
                + "\n"
            )

        rc = display_stats(most_important_stats)
        rc += "\n   ** Workers **\n"
        rc += display_stats(queue_stats)
        if detailed:
            rc += "\n   ** Internals **\n"
            rc += display_stats(
                set(dir(self)) - (set(most_important_stats) & set(queue_stats))
            )
        return rc
