# vim: set ts=4

# Copyright 2017 Rémi Duraffort
# This file is part of lavacli.
#
# lavacli is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# lavacli is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with lavacli.  If not, see <http://www.gnu.org/licenses/>

import argparse
import contextlib
import datetime
import json
import pathlib
import sys
import time
import xmlrpc.client
from urllib.parse import urlparse

from lavacli.utils import exc2str, log_yaml, safe_yaml
from lavacli.utils import fetch_jobs, print_error_summary
from collections import defaultdict


def configure_log_options(parser):
    parser.add_argument(
        "--filters",
        default=None,
        type=str,
        help="comma separated list of levels to show",
    )
    parser.add_argument(
        "--polling",
        default=5,
        type=int,
        help="polling interval in seconds [default=5s]",
    )
    parser.add_argument(
        "--raw", default=False, action="store_true", help="print raw logs"
    )
    parser.add_argument("--start", default=0, type=int, help="start at the given line")
    parser.add_argument("--end", default=None, type=int, help="end at the given line")


def configure_parser(parser, version):
    sub = parser.add_subparsers(dest="sub_sub_command", help="Sub commands")
    sub.required = True

    # "cancel"
    jobs_cancel = sub.add_parser("cancel", help="cancel a job")
    jobs_cancel.add_argument("job_id", nargs="+", help="job id")

    if version >= (2023, 10):
        # "autocancel"
        jobs_autocancel = sub.add_parser(
            "autocancel",
            help="cancel jobs that take longer than job timeout and the specific duration",
        )
        jobs_autocancel.add_argument(
            "--duration", type=int, required=True, help="duration in minutes"
        )
        jobs_autocancel.add_argument(
            "--dry-run",
            action="store_true",
            default=False,
            help="do not cancel the jobs",
        )

    # "config"
    jobs_config = sub.add_parser("config", help="job configuration")
    jobs_config.add_argument("job_id", help="job id")
    jobs_config.add_argument(
        "--dest", default=".", help="save files into this directory"
    )

    # "definition"
    jobs_definition = sub.add_parser("definition", help="job definition")
    jobs_definition.add_argument("job_id", help="job id")

    # "find-errors"
    jobs_find_errors = sub.add_parser("find-errors", help="Find errors in recent jobs")
    jobs_find_errors.add_argument(
        "--limit", type=int, default=100, help="Number of jobs to fetch"
    )
    jobs_find_errors.add_argument(
        "--since", type=int, default=None, help="Jobs from the last N minutes"
    )

    # "list"
    jobs_list = sub.add_parser("list", help="list jobs")
    jobs_list.add_argument(
        "--state",
        type=str,
        default=None,
        choices=[
            "SUBMITTED",
            "SCHEDULING",
            "SCHEDULED",
            "RUNNING",
            "CANCELING",
            "FINISHED",
        ],
        help="filter jobs by state",
    )
    jobs_list.add_argument(
        "--health",
        type=str,
        default=None,
        choices=["UNKNOWN", "COMPLETE", "INCOMPLETE", "CANCELED"],
        help="filter jobs by health",
    )
    jobs_list.add_argument(
        "--since",
        type=int,
        default=0,
        help="Filter by jobs which completed in the last N minutes.",
    )
    jobs_list.add_argument(
        "--verbose",
        dest="verbose",
        action="store_true",
        help="If verbose is True, add extra keys, including error_type.",
    )
    jobs_list.add_argument(
        "--start", type=int, default=0, help="skip the N first jobs [default=0]"
    )
    jobs_list.add_argument(
        "--limit", type=int, default=25, help="limit to N jobs [default=25]"
    )
    out_format = jobs_list.add_mutually_exclusive_group()
    out_format.add_argument(
        "--json",
        dest="output_format",
        default=None,
        action="store_const",
        const="json",
        help="print as json",
    )
    out_format.add_argument(
        "--yaml",
        dest="output_format",
        default=None,
        action="store_const",
        const="yaml",
        help="print as yaml",
    )

    # "logs"
    jobs_logs = sub.add_parser("logs", help="get logs")
    jobs_logs.add_argument("job_id", help="job id")
    jobs_logs.add_argument(
        "--no-follow",
        default=False,
        action="store_true",
        help="do not keep polling until the end of the job",
    )
    configure_log_options(jobs_logs)

    # "queue"
    jobs_queue = sub.add_parser("queue", help="job queue")
    jobs_queue.add_argument("device_types", nargs="*", help="filter by device-types")
    out_format = jobs_queue.add_mutually_exclusive_group()
    out_format.add_argument(
        "--json",
        dest="output_format",
        default=None,
        action="store_const",
        const="json",
        help="print as json",
    )
    out_format.add_argument(
        "--yaml",
        dest="output_format",
        action="store_const",
        const="yaml",
        default=None,
        help="print as yaml",
    )
    jobs_queue.add_argument(
        "--start", type=int, default=0, help="skip the N first jobs [default=0]"
    )
    jobs_queue.add_argument(
        "--limit", type=int, default=25, help="limit to N jobs [default=25]"
    )

    # "resubmit"
    jobs_resubmit = sub.add_parser("resubmit", help="resubmit a job")
    jobs_resubmit.add_argument("job_id", help="job id")
    jobs_resubmit.add_argument(
        "--url",
        action="store_true",
        dest="print_url",
        default=False,
        help="Print the full url",
    )
    jobs_resubmit.add_argument(
        "--follow",
        default=True,
        dest="no_follow",
        action="store_false",
        help="resubmit and poll for the logs",
    )
    configure_log_options(jobs_resubmit)

    # "run"
    jobs_run = sub.add_parser("run", help="run the job")
    jobs_run.add_argument(
        "definition", type=argparse.FileType("r"), help="job definition"
    )
    jobs_run.add_argument(
        "--filters",
        default=None,
        type=str,
        help="comma separated list of levels to show",
    )
    jobs_run.add_argument(
        "--no-follow",
        default=False,
        action="store_true",
        help="do not keep polling until the end of the job",
    )
    jobs_run.add_argument(
        "--polling",
        default=5,
        type=int,
        help="polling interval in seconds [default=5s]",
    )
    jobs_run.add_argument(
        "--raw", default=False, action="store_true", help="print raw logs"
    )
    jobs_run.add_argument(
        "--start", default=0, type=int, help="start at the given line"
    )
    jobs_run.add_argument("--end", default=None, type=int, help="end at the given line")

    # "show"
    jobs_show = sub.add_parser("show", help="job details")
    jobs_show.add_argument("job_id", help="job id")
    out_format = jobs_show.add_mutually_exclusive_group()
    out_format.add_argument(
        "--json",
        dest="output_format",
        default=None,
        action="store_const",
        const="json",
        help="print as json",
    )
    out_format.add_argument(
        "--yaml",
        dest="output_format",
        action="store_const",
        const="yaml",
        default=None,
        help="print as yaml",
    )

    # "submit"
    jobs_submit = sub.add_parser("submit", help="submit a new job")
    jobs_submit.add_argument(
        "definition", nargs="+", type=argparse.FileType("r"), help="job definition"
    )
    jobs_submit.add_argument(
        "--url",
        action="store_true",
        dest="print_url",
        default=False,
        help="Print the full url",
    )
    jobs_submit.add_argument(
        "--follow",
        default=True,
        dest="no_follow",
        action="store_false",
        help="resubmit and poll for the logs",
    )
    configure_log_options(jobs_submit)

    if version >= (2019, 3):
        # "validate"
        jobs_validate = sub.add_parser("validate", help="validate the job definition")
        jobs_validate.add_argument(
            "definition", type=argparse.FileType("r"), help="job definition"
        )
        jobs_validate.add_argument(
            "--strict", action="store_true", default=False, help="check in strict mode"
        )

    # "wait"
    jobs_wait = sub.add_parser("wait", help="wait for the job to finish")
    jobs_wait.add_argument("job_id", help="job id")
    jobs_wait.add_argument(
        "--polling",
        default=5,
        type=int,
        help="polling interval in seconds [default=5s]",
    )
    jobs_wait.add_argument(
        "--timeout",
        default=0,
        type=int,
        help="Maximum time to wait in seconds, 0 to disable [default=0s]",
    )


def help_string():
    return "manage jobs"


def handle_cancel(proxy, options, _):
    ret = 0
    for job_id in options.job_id:
        try:
            proxy.scheduler.jobs.cancel(job_id)
        except xmlrpc.client.Error as exc:
            print(
                f"Unable to cancel job {job_id}: {exc2str(exc, options.uri)}",
                file=sys.stderr,
            )
            ret = 1
    return ret


def handle_autocancel(proxy, options, _):
    # Jobs that still running within its timeout shouldn't be canceled.
    skip_list = []
    # A loop is needed as the list api only returns the last 100 jobs.
    while True:
        jobs = proxy.scheduler.jobs.list(
            "RUNNING",
            None,
            0,
            100,
            None,
            False,
            options.duration,
        )

        if not jobs:
            break

        job_ids = [job["id"] for job in jobs]
        if set(job_ids).issubset(skip_list):
            break
        jobs = [job for job in jobs if job["id"] not in skip_list]

        duration = datetime.timedelta(minutes=options.duration)
        for job in jobs:
            job_id = job["id"]
            definition = proxy.scheduler.jobs.definition(job_id)
            data = safe_yaml.load(definition)["timeouts"]["job"]
            job_timeout = datetime.timedelta(
                days=data.get("days", 0),
                hours=data.get("hours", 0),
                minutes=data.get("minutes", 0),
                seconds=data.get("seconds", 0),
            )
            if duration > job_timeout:
                try:
                    if not options.dry_run:
                        proxy.scheduler.jobs.cancel(job_id)
                    print(
                        "[CANCELED] %s: [%s] (%s) - %s"
                        % (
                            job_id,
                            job["submitter"],
                            job["description"],
                            job["device_type"],
                        )
                    )
                except xmlrpc.client.Error as exc:
                    print(
                        f"Unable to cancel {job_id}: {exc2str(exc, options.uri)}",
                        file=sys.stderr,
                    )
                    return 1
            else:
                print(
                    "[  SKIP  ] %s: [%s] (%s) - %s - timeout not reached yet"
                    % (
                        job_id,
                        job["submitter"],
                        job["description"],
                        job["device_type"],
                    )
                )
                skip_list.append(job_id)

    return 0


def handle_config(proxy, options, _):
    config = proxy.scheduler.jobs.configuration(options.job_id)
    definition, device, dispatcher, env, env_dut = config

    # Create the directory if needed
    dest = pathlib.Path(options.dest)
    with contextlib.suppress(FileExistsError):
        dest.mkdir(mode=0o755)

    # Save the files
    with (dest / "definition.yaml").open("w") as f_out:
        f_out.write(definition)
    if device is not None:
        with (dest / "device.yaml").open("w") as f_out:
            f_out.write(device)
    if dispatcher is not None:
        with (dest / "dispatcher.yaml").open("w") as f_out:
            f_out.write(dispatcher)
    if env is not None:
        with (dest / "env.yaml").open("w") as f_out:
            f_out.write(env)
    if env_dut is not None:
        with (dest / "env.dut.yaml").open("w") as f_out:
            f_out.write(env_dut)
    return 0


def handle_definition(proxy, options, _):
    print(proxy.scheduler.jobs.definition(options.job_id))
    return 0


def handle_list(proxy, options, config):
    jobs = proxy.scheduler.jobs.list(
        options.state,
        options.health,
        options.start,
        options.limit,
        options.since,
        options.verbose,
    )

    if options.output_format == "json":
        print(json.dumps(jobs))
    elif options.output_format == "yaml":
        safe_yaml.dump(jobs, sys.stdout)
    else:
        print(f"Jobs (from {1 + options.start} to {options.start + options.limit}):")
        for job in jobs:
            if options.verbose:
                if job["error_type"]:
                    print(
                        "* %s: %s,%s [%s] (%s) - %s %s <%s> <%s> %s: %s"
                        % (
                            job["id"],
                            job["state"],
                            job["health"],
                            job["submitter"],
                            job["description"],
                            job["device_type"],
                            job["actual_device"],
                            job["start_time"],
                            job["end_time"],
                            job["error_type"],
                            job["error_msg"],
                        )
                    )
                else:
                    print(
                        "* %s: %s,%s [%s] (%s) - %s %s <%s> <%s>"
                        % (
                            job["id"],
                            job["state"],
                            job["health"],
                            job["submitter"],
                            job["description"],
                            job["device_type"],
                            job["actual_device"],
                            job["start_time"],
                            job["end_time"],
                        )
                    )

            else:
                print(
                    "* %s: %s,%s [%s] (%s) - %s"
                    % (
                        job["id"],
                        job["state"],
                        job["health"],
                        job["submitter"],
                        job["description"],
                        job["device_type"],
                    )
                )
    return 0


if sys.stdout.isatty():
    COLORS = {
        "exception": "\033[1;31m",
        "error": "\033[1;31m",
        "warning": "\033[1;33m",
        "info": "\033[1;37m",
        "debug": "\033[0;37m",
        "target": "\033[32m",
        "input": "\033[0;35m",
        "feedback": "\033[0;33m",
        "results": "\033[1;34m",
        "dt": "\033[0;90m",
        "end": "\033[0m",
    }
else:
    COLORS = {
        "exception": "",
        "error": "",
        "warning": "",
        "info": "",
        "debug": "",
        "target": "",
        "input": "",
        "feedback": "",
        "results": "",
        "dt": "",
        "end": "",
    }


def print_logs(logs, raw, filters):
    filters = [] if filters is None else filters.split(",")
    if raw:
        for line in logs:
            if filters and not line["lvl"] in filters:
                continue
            sys.stdout.write("- ")
            log_yaml.dump(
                line,
                sys.stdout,
            )
    else:
        for line in logs:
            timestamp = line["dt"].split(".")[0]
            level = line["lvl"]
            if filters and level not in filters:
                continue
            if isinstance(line["msg"], dict) and "sending" in line["msg"].keys():
                level = "input"
                msg = str(line["msg"]["sending"])
            elif isinstance(line["msg"], bytes):
                msg = line["msg"].decode("utf-8", errors="replace")
            else:
                msg = str(line["msg"])
            msg = msg.rstrip("\n")

            ns = ""
            if level == "feedback" and "ns" in line:
                ns = " <" + COLORS["feedback"] + line["ns"] + COLORS["end"] + ">"
            print(
                COLORS["dt"]
                + timestamp
                + COLORS["end"]
                + ns
                + " "
                + COLORS[level]
                + msg
                + COLORS["end"]
            )


def _download_logs(proxy, version, job_id, start, end):
    (finished, data) = proxy.scheduler.jobs.logs(job_id, start, end)
    logs = safe_yaml.load(str(data))
    return (finished, logs)


def handle_logs(proxy, options, config):
    # Loop
    lines = options.start
    while True:
        try:
            (finished, logs) = _download_logs(
                proxy, config["version"], options.job_id, lines, options.end
            )
            if logs:
                print_logs(logs, options.raw, options.filters)
                lines += len(logs)
            # Loop only if the job is not finished
            if (
                finished
                or options.no_follow
                or (options.end is not None and lines >= options.end)
            ):
                break
        except xmlrpc.client.ProtocolError as exc:
            # Ignore temporary Internal Server Error as we already got some logs
            if lines > options.start and exc.errcode == 500:
                pass

        # Wait some time
        time.sleep(options.polling)

    # Print the failure comment if the job is finished
    if finished:
        details = proxy.scheduler.jobs.show(options.job_id)
        if details.get("failure_comment"):
            print_logs(
                [
                    {
                        "dt": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                        "lvl": "info",
                        "msg": "[lavacli] Failure comment: %s"
                        % details["failure_comment"],
                    }
                ],
                options.raw,
                options.filters,
            )
    return 0


def handle_queue(proxy, options, config):
    dts = options.device_types if options.device_types else None
    data = proxy.scheduler.jobs.queue(dts, options.start, options.limit)
    if options.output_format == "json":
        print(json.dumps(data))
    elif options.output_format == "yaml":
        safe_yaml.dump(data, sys.stdout)
    else:
        print(f"Jobs (from {1+options.start} to {options.start + options.limit}):")
        for job in data:
            print(
                f"* {job['id']}: {job['submitter']} ({job['description'] or ''}) - {job['requested_device_type']}"
            )
    return 0


def handle_resubmit(proxy, options, config):
    job_id = proxy.scheduler.jobs.resubmit(options.job_id)

    if options.no_follow:
        prefix = ""
        if options.print_url:
            parsed = urlparse(options.uri)
            host = parsed.netloc.split("@")[-1]
            prefix = f"{parsed.scheme}://{host}/scheduler/job/"

        if isinstance(job_id, list):
            for job in job_id:
                print(prefix + str(job))
        else:
            print(prefix + str(job_id))

    else:
        print_logs(
            [
                {
                    "dt": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                    "lvl": "info",
                    "msg": "[lavacli] Job %s submitted" % job_id,
                }
            ],
            options.raw,
            options.filters,
        )

        follow_logs(job_id, proxy, options, config)
    return 0


def follow_logs(job_id, proxy, options, config):
    # Add the job_id to options for handle_logs
    # For multinode, print something and loop on all jobs
    if isinstance(job_id, list):
        for job in job_id:
            print_logs(
                [
                    {
                        "dt": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                        "lvl": "info",
                        "msg": "[lavacli] Seeing %s logs" % job,
                    }
                ],
                options.raw,
                options.filters,
            )
            options.job_id = job
            handle_logs(proxy, options, config)
    else:
        options.job_id = str(job_id)
        handle_logs(proxy, options, config)


def handle_run(proxy, options, config):
    job_id = proxy.scheduler.jobs.submit(options.definition.read())
    print_logs(
        [
            {
                "dt": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                "lvl": "info",
                "msg": "[lavacli] Job %s submitted" % job_id,
            }
        ],
        options.raw,
        options.filters,
    )

    # Add the job_id to options for handle_logs
    # For multinode, print something and loop on all jobs
    if isinstance(job_id, list):
        for job in job_id:
            print_logs(
                [
                    {
                        "dt": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                        "lvl": "info",
                        "msg": "[lavacli] Seeing %s logs" % job,
                    }
                ],
                options.raw,
                options.filters,
            )
            options.job_id = job
            handle_logs(proxy, options, config)
    else:
        options.job_id = job_id
        handle_logs(proxy, options, config)
    return 0


def handle_show(proxy, options, config):
    job = proxy.scheduler.jobs.show(options.job_id)

    if options.output_format == "json":
        job["submit_time"] = job["submit_time"].value or None
        job["start_time"] = job["start_time"].value or None
        job["end_time"] = job["end_time"].value or None
        print(json.dumps(job))
    elif options.output_format == "yaml":
        job["submit_time"] = job["submit_time"].value or None
        job["start_time"] = job["start_time"].value or None
        job["end_time"] = job["end_time"].value or None
        safe_yaml.dump(job, sys.stdout)
    else:
        print("id          : %s" % job["id"])
        print("description : %s" % job["description"])
        print("submitter   : %s" % job["submitter"])
        print("device-type : %s" % job["device_type"])
        print("device      : %s" % job["device"])
        print("health-check: %s" % job["health_check"])
        print("state       : %s" % job["state"])
        print("Health      : %s" % job["health"])
        if job.get("failure_comment"):
            print("failure     : %s" % job["failure_comment"])
        if job.get("error_type"):
            print("error       : %s: %s" % (job["error_type"], job["error_msg"]))
        print("pipeline    : %s" % job["pipeline"])
        print("tags        : %s" % str(job["tags"]))
        print("visibility  : %s" % job["visibility"])
        print("submit time : %s" % job["submit_time"])
        print("start time  : %s" % job["start_time"])
        print("end time    : %s" % job["end_time"])
    return 0


def handle_submit(proxy, options, config):
    prefix = ""
    if options.print_url:
        parsed = urlparse(options.uri)
        host = parsed.netloc.split("@")[-1]
        prefix = f"{parsed.scheme}://{host}/scheduler/job/"

    for definition in options.definition:
        try:
            job_id = proxy.scheduler.jobs.submit(definition.read())
            if options.no_follow:
                if isinstance(job_id, list):
                    for job in job_id:
                        print(prefix + str(job))
                else:
                    print(prefix + str(job_id))
        except xmlrpc.client.Error as exc:
            print(
                f"Unable to submit {definition.name}: {exc2str(exc, options.uri)}",
                file=sys.stderr,
            )
            return 1

    if not options.no_follow:
        follow_logs(job_id, proxy, options, config)

    return 0


def handle_validate(proxy, options, _):
    ret = proxy.scheduler.jobs.validate(options.definition.read(), options.strict)
    if not ret:
        return 0
    print("key: %s" % ret["path"])
    print("msg: %s" % ret["msg"])
    return 1


def handle_wait(proxy, options, _):
    job = proxy.scheduler.jobs.show(options.job_id)
    old_state = ""
    time_elapsed = 0
    while job["state"] != "Finished":
        if old_state != job["state"]:
            if old_state:
                sys.stdout.write("\n")
            sys.stdout.write(job["state"])
        else:
            sys.stdout.write(".")
        sys.stdout.flush()
        old_state = job["state"]
        time.sleep(options.polling)
        job = proxy.scheduler.jobs.show(options.job_id)
        if options.timeout > 0:
            time_elapsed = time_elapsed + options.polling
            if time_elapsed > options.timeout:
                print("timeout!")
                return 1
    if old_state != job["state"] and old_state:
        sys.stdout.write("\n")
    return 0


def handle_find_errors(proxy, options, config):
    """Main function to fetch and process LAVA jobs."""

    parsed = urlparse(options.uri)
    host = parsed.netloc.split("@")[-1]
    url = f"{parsed.scheme}://{host}"

    job_numbers = fetch_jobs(proxy, None, "INCOMPLETE", options.limit, options.since)
    if not job_numbers:
        print("No jobs found.")
        return

    print(f"\nTotal jobs found: {len(job_numbers)}")

    error_summary = defaultdict(lambda: defaultdict(list))

    for job_id in job_numbers:
        try:
            job = proxy.scheduler.jobs.show(job_id)
        except xmlrpc.client.Error as exc:
            print(f"Error fetching job {job_id}: {exc}")
            continue

        device_name = job.get("device") or job.get("device_type")
        job_url = f"{url}/scheduler/job/{job_id}"

        error_msg = job.get("error_msg", "").strip()
        if error_msg:
            error_summary[error_msg][device_name].append(job_url)

    print_error_summary(error_summary, summary=True)
    print_error_summary(error_summary, summary=False)


def handle(proxy, options, config):
    handlers = {
        "cancel": handle_cancel,
        "autocancel": handle_autocancel,
        "config": handle_config,
        "definition": handle_definition,
        "list": handle_list,
        "logs": handle_logs,
        "queue": handle_queue,
        "resubmit": handle_resubmit,
        "run": handle_run,
        "show": handle_show,
        "submit": handle_submit,
        "validate": handle_validate,
        "wait": handle_wait,
        "find-errors": handle_find_errors,
    }
    return handlers[options.sub_sub_command](proxy, options, config)
