import argparse
import datetime
import importlib
import importlib.resources
import logging
import logging.handlers
import os
import glob
import random
from shutil import copyfile
import sys
import time
import typing

import pendulum

# Plotman libraries
from plotman import (
    analyzer,
    archive,
    configuration,
    interactive,
    manager,
    plot_util,
    reporting,
    csv_exporter,
)
from plotman import resources as plotman_resources
from plotman.job import Job


class PlotmanArgParser:
    def add_idprefix_arg(self, subparser: argparse.ArgumentParser) -> None:
        subparser.add_argument(
            "idprefix", type=str, nargs="+", help="disambiguating prefix of plot ID"
        )

    def parse_args(self) -> typing.Any:
        parser = argparse.ArgumentParser(description="Chia plotting manager.")
        sp = parser.add_subparsers(dest="cmd")

        sp.add_parser("version", help="print the version")

        p_status = sp.add_parser("status", help="show current plotting status")
        p_status.add_argument(
            "--json", action="store_true", help="export status report in json format"
        )

        sp.add_parser(
            "prometheus",
            help="show current plotting status in prometheus readable format",
        )

        sp.add_parser("dirs", help="show directories info")

        p_interactive = sp.add_parser(
            "interactive", help="run interactive control/monitoring mode"
        )
        p_interactive.add_argument(
            "--autostart-plotting",
            action="store_true",
            default=None,
            dest="autostart_plotting",
        )
        p_interactive.add_argument(
            "--no-autostart-plotting",
            action="store_false",
            default=None,
            dest="autostart_plotting",
        )
        p_interactive.add_argument(
            "--autostart-archiving",
            action="store_true",
            default=None,
            dest="autostart_archiving",
        )
        p_interactive.add_argument(
            "--no-autostart-archiving",
            action="store_false",
            default=None,
            dest="autostart_archiving",
        )

        sp.add_parser("dsched", help="print destination dir schedule")

        sp.add_parser("plot", help="run plotting loop")

        sp.add_parser("archive", help="move completed plots to farming location")

        p_export = sp.add_parser(
            "export", help="exports metadata from the plot logs as CSV"
        )
        p_export.add_argument(
            "-o",
            dest="save_to",
            default=None,
            type=str,
            help="save to file. Optional, prints to stdout by default",
        )

        p_config = sp.add_parser(
            "config", help="display or generate plotman.yaml configuration"
        )
        sp_config = p_config.add_subparsers(dest="config_subcommand")
        sp_config.add_parser(
            "generate", help="generate a default plotman.yaml file and print path"
        )
        sp_config.add_parser("path", help="show path to current plotman.yaml file")

        p_details = sp.add_parser("details", help="show details for job")
        self.add_idprefix_arg(p_details)

        p_logs = sp.add_parser("logs", help="fetch the logs for job")

        p_logs.add_argument(
            "-f", "--follow", action="store_true", help="Follow log output"
        )
        self.add_idprefix_arg(p_logs)

        p_files = sp.add_parser("files", help="show temp files associated with job")
        self.add_idprefix_arg(p_files)

        p_kill = sp.add_parser("kill", help="kill job (and cleanup temp files)")
        p_kill.add_argument(
            "-f",
            "--force",
            action="store_true",
            default=False,
            help="Don't ask for confirmation before killing the plot job",
        )
        self.add_idprefix_arg(p_kill)

        p_suspend = sp.add_parser("suspend", help="suspend job")
        self.add_idprefix_arg(p_suspend)

        p_resume = sp.add_parser("resume", help="resume suspended job")
        self.add_idprefix_arg(p_resume)

        p_analyze = sp.add_parser(
            "analyze", help="analyze timing stats of completed jobs"
        )

        p_analyze.add_argument(
            "--clipterminals",
            action="store_true",
            help="Ignore first and last plot in a logfile, useful for "
            "focusing on the steady-state in a staggered parallel "
            "plotting test (requires plotting  with -n>2)",
        )
        p_analyze.add_argument("--bytmp", action="store_true", help="slice by tmp dirs")
        p_analyze.add_argument(
            "--bybitfield",
            action="store_true",
            help="slice by bitfield/non-bitfield sorting",
        )
        p_analyze.add_argument(
            "logfile", type=str, nargs="+", help="logfile(s) to analyze"
        )

        args = parser.parse_args()
        return args


def get_term_width(cfg: configuration.PlotmanConfig) -> int:
    default_columns = 120  # 80 is typically too narrow.
    if cfg.user_interface.use_stty_size:
        try:
            (rows_string, columns_string) = os.popen("stty size", "r").read().split()
            columns = int(columns_string)
        except:
            columns = default_columns
    else:
        columns = default_columns
    return columns


class Iso8601Formatter(logging.Formatter):
    def formatTime(
        self, record: logging.LogRecord, datefmt: typing.Optional[str] = None
    ) -> str:
        time = pendulum.from_timestamp(timestamp=record.created, tz="local")
        return time.isoformat(timespec="microseconds")


def main() -> None:
    random.seed()

    pm_parser = PlotmanArgParser()
    args = pm_parser.parse_args()

    if args.cmd == "version":
        import pkg_resources

        print(pkg_resources.get_distribution("plotman"))
        return

    elif args.cmd == "config":
        config_file_path = configuration.get_path()
        if args.config_subcommand == "path":
            if os.path.isfile(config_file_path):
                print(config_file_path)
                return
            print(
                f"No 'plotman.yaml' file exists at expected location: '{config_file_path}'"
            )
            print(f"To generate a default config file, run: 'plotman config generate'")
            return
        if args.config_subcommand == "generate":
            if os.path.isfile(config_file_path):
                overwrite = None
                while overwrite not in {"y", "n"}:
                    overwrite = input(
                        f"A 'plotman.yaml' file already exists at the default location: '{config_file_path}' \n\n"
                        "\tInput 'y' to overwrite existing file, or 'n' to exit without overwrite."
                    ).lower()
                    if overwrite == "n":
                        print("\nExited without overrwriting file")
                        return

            # Copy the default plotman.yaml (packaged in plotman/resources/) to the user's config file path,
            # creating the parent plotman file/directory if it does not yet exist
            with importlib.resources.path(
                plotman_resources, "plotman.yaml"
            ) as default_config:
                config_dir = os.path.dirname(config_file_path)

                os.makedirs(config_dir, exist_ok=True)
                copyfile(default_config, config_file_path)
                print(f"\nWrote default plotman.yaml to: {config_file_path}")
                return

        if not args.config_subcommand:
            print("No action requested, add 'generate' or 'path'.")
            return

    config_path = configuration.get_path()
    config_text = configuration.read_configuration_text(config_path)
    preset_target_definitions_text = importlib.resources.read_text(
        plotman_resources,
        "target_definitions.yaml",
    )

    cfg = configuration.get_validated_configs(
        config_text, config_path, preset_target_definitions_text
    )

    with cfg.setup():
        root_logger = logging.getLogger()
        root_handler = logging.handlers.RotatingFileHandler(
            backupCount=10,
            encoding="utf-8",
            filename=cfg.logging.application,
            maxBytes=10_000_000,
        )
        root_formatter = Iso8601Formatter(fmt="%(asctime)s: %(message)s")
        root_handler.setFormatter(root_formatter)
        root_logger.addHandler(root_handler)
        root_logger.setLevel(logging.INFO)
        root_logger.info("Start root logger")

        disk_space_logger = logging.getLogger("disk_space")
        disk_space_logger.propagate = False
        disk_space_handler = logging.handlers.RotatingFileHandler(
            backupCount=10,
            encoding="utf-8",
            filename=cfg.logging.disk_spaces,
            maxBytes=10_000_000,
        )
        disk_space_formatter = Iso8601Formatter(fmt="%(asctime)s: %(message)s")
        disk_space_handler.setFormatter(disk_space_formatter)
        disk_space_logger.addHandler(disk_space_handler)
        disk_space_logger.setLevel(logging.INFO)
        disk_space_logger.info("Start disk space logger")

        #
        # Stay alive, spawning plot jobs
        #
        if args.cmd == "plot":
            print("...starting plot loop")
            while True:
                (started, msg) = manager.maybe_start_new_plot(
                    cfg.directories, cfg.scheduling, cfg.plotting, cfg.logging
                )

                # TODO: report this via a channel that can be polled on demand, so we don't spam the console
                if started:
                    print("%s" % (msg))
                else:
                    print("...sleeping %d s: %s" % (cfg.scheduling.polling_time_s, msg))
                root_logger.info("[plot] %s", msg)

                time.sleep(cfg.scheduling.polling_time_s)

        #
        # Analysis of completed jobs
        #
        elif args.cmd == "analyze":

            analyzer.analyze(
                args.logfile,
                args.clipterminals,
                args.bytmp,
                args.bybitfield,
                get_term_width(cfg),
            )

        #
        # Exports log metadata to CSV
        #
        elif args.cmd == "export":
            logfilenames = glob.glob(os.path.join(cfg.logging.plots, "*.plot.log"))
            if args.save_to is None:
                csv_exporter.generate(logfilenames=logfilenames, file=sys.stdout)
            else:
                with open(args.save_to, "w", encoding="utf-8") as file:
                    csv_exporter.generate(logfilenames=logfilenames, file=file)

        else:
            jobs = Job.get_running_jobs(cfg.logging.plots)

            # Status report
            if args.cmd == "status":
                if args.json:
                    # convert jobs list into json
                    result = reporting.json_report(jobs)
                else:
                    result = "{0}\n\n{1}\n\nUpdated at: {2}".format(
                        reporting.status_report(jobs, get_term_width(cfg)),
                        reporting.summary(jobs),
                        datetime.datetime.today().strftime("%c"),
                    )
                print(result)

            # Prometheus report
            if args.cmd == "prometheus":
                print(reporting.prometheus_report(jobs))

            # Directories report
            elif args.cmd == "dirs":
                print(
                    reporting.dirs_report(
                        jobs,
                        cfg.directories,
                        cfg.archiving,
                        cfg.scheduling,
                        get_term_width(cfg),
                    )
                )

            elif args.cmd == "interactive":
                interactive.run_interactive(
                    cfg=cfg,
                    autostart_plotting=args.autostart_plotting,
                    autostart_archiving=args.autostart_archiving,
                )

            # Start running archival
            elif args.cmd == "archive":
                if cfg.archiving is None:
                    start_msg = (
                        "archiving not configured but is required for this command"
                    )
                    print(start_msg)
                    root_logger.info("[archive] %s", start_msg)
                else:
                    start_msg = "...starting archive loop"
                    print(start_msg)
                    root_logger.info("[archive] %s", start_msg)
                    firstit = True
                    while True:
                        if not firstit:
                            print(
                                "Sleeping %d s until next iteration..."
                                % (cfg.scheduling.polling_time_s)
                            )
                            time.sleep(cfg.scheduling.polling_time_s)
                            jobs = Job.get_running_jobs(cfg.logging.plots)
                        firstit = False

                        archiving_status, log_messages = archive.spawn_archive_process(
                            cfg.directories, cfg.archiving, cfg.logging, jobs
                        )
                        if log_messages:
                            for log_message in log_messages:
                                print(log_message)
                                root_logger.info("[archive] %s", log_message)
                        else:
                            root_logger.info("[archive] %s", archiving_status)

            # Debugging: show the destination drive usage schedule
            elif args.cmd == "dsched":
                for (d, ph) in manager.dstdirs_to_furthest_phase(jobs).items():
                    print("  %s : %s" % (d, str(ph)))

            #
            # Job control commands
            #
            elif args.cmd in ["details", "logs", "files", "kill", "suspend", "resume"]:
                print(args)

                selected = []

                # TODO: clean up treatment of wildcard
                if args.idprefix[0] == "all":
                    selected = jobs
                else:
                    # TODO: allow multiple idprefixes, not just take the first
                    selected = manager.select_jobs_by_partial_id(jobs, args.idprefix[0])
                    if len(selected) == 0:
                        print("Error: %s matched no jobs." % args.idprefix[0])
                    elif len(selected) > 1:
                        print('Error: "%s" matched multiple jobs:' % args.idprefix[0])
                        for j in selected:
                            print("  %s" % j.plotter.common_info().plot_id)
                        selected = []

                for job in selected:
                    if args.cmd == "details":
                        print(job.status_str_long())

                    elif args.cmd == "logs":
                        job.print_logs(args.follow)

                    elif args.cmd == "files":
                        temp_files = job.get_temp_files()
                        for f in temp_files:
                            print("  %s" % f)

                    elif args.cmd == "kill":
                        info = job.plotter.common_info()
                        # First suspend so job doesn't create new files
                        print(
                            "Pausing PID %d, plot id %s" % (job.proc.pid, info.plot_id)
                        )
                        job.suspend()

                        temp_files = job.get_temp_files()

                        print(
                            "Will kill pid %d, plot id %s"
                            % (job.proc.pid, info.plot_id)
                        )
                        print("Will delete %d temp files" % len(temp_files))

                        if args.force:
                            conf = "y"
                        else:
                            conf = input('Are you sure? ("y" to confirm): ')

                        if conf != "y":
                            print(
                                "Canceled.  If you wish to resume the job, do so manually."
                            )
                        else:
                            print("killing...")

                            job.cancel()

                            print("cleaning up temp files...")

                            for f in temp_files:
                                os.remove(f)

                    elif args.cmd == "suspend":
                        print(f"Suspending {job.plotter.common_info().plot_id}")
                        job.suspend()
                    elif args.cmd == "resume":
                        print(f"Resuming {job.plotter.common_info().plot_id}")
                        job.resume()
