#!/usr/bin/env python3

import argparse
import io
import json
import os
import subprocess
import sys
import textwrap
from collections import defaultdict
from dataclasses import asdict, dataclass, field, is_dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple


@dataclass(frozen=True)
class SourceLocation:
    file: Path
    line: int
    column: Optional[int] = field(default=None)


@dataclass
class Report:
    name: str
    msg: str
    locations: List[SourceLocation]
    link: bool = field(default=True)
    severity: str = field(default="warning")


@dataclass(frozen=True)
class Attr:
    # name of attr, e.g. python38Packages.numpy
    name: str
    # location in which the attr is defined, if exist
    location: Optional[SourceLocation]
    # path to the .drv file, if exists
    drv: Optional[Path]
    # path the the output of the drv in the nix store, if exists
    output: Optional[Path]


class JSONEncoder(json.JSONEncoder):
    def default(self, o: Any) -> Any:
        if is_dataclass(o):
            return asdict(o)

        if isinstance(o, Path):
            return str(o)

        # fallback to superclass
        return super().default(o)


def get_check_programs() -> Set[str]:
    # Each rule that is implemented as an external check program rather
    # than an overlay is installed onto our PATH, and the names of the
    # rules put into this environment variable.
    ast_checks = os.environ.get("AST_CHECK_NAMES")

    if ast_checks:
        return set(ast_checks.split(":"))

    return set()


def run_external_checks(
    attrs: List[Attr], excluded_rules: List[str]
) -> Dict[str, List[Report]]:
    rules = get_check_programs() - set(excluded_rules)
    if len(rules) == 0:
        return {}

    encoded_attrs = json.dumps(attrs, cls=JSONEncoder)
    result: Dict[Attr, List[Report]] = defaultdict(list)

    if "no-build-output" not in excluded_rules:
        names_without_outputs = [
            a.name for a in attrs if a.output is None or not a.output.exists()
        ]
        for name in names_without_outputs:
            result[name].append(
                Report(
                    name="no-build-output",
                    msg=f"‘{name}’ has not yet been built. Checks that rely on the presence of build logs are skipped.",
                    link=False,
                    locations=[],
                    severity="notice",
                )
            )

    for rule in rules:
        try:
            json_text = subprocess.check_output([rule], text=True, input=encoded_attrs)
        except subprocess.CalledProcessError:
            print(
                red(
                    textwrap.dedent(
                        f"""
                        Rule ‘{rule}’ failed with input ‘{encoded_attrs}’.
                        This is a bug. Please file a ticket on GitHub.
                        """
                    )
                ),
                file=sys.stderr,
            )
            raise
        if len(json_text) > 0:
            for attr_name, reports in json.loads(json_text).items():
                result[attr_name].extend(load_reports(reports))

    return dict(result)


def concatenate_messages(
    *report_sources: Dict[str, List[Report]]
) -> Dict[Attr, List[Report]]:
    result = defaultdict(list)
    for m in report_sources:
        for attr_name, report in m.items():
            result[attr_name].extend(report)
    return dict(result)


def escape_attr(attr: str) -> str:
    return ".".join(f'"{section}"' for section in attr.split("."))


def nix_eval_json(expr: str, show_trace: bool = False) -> Dict[Attr, List[Report]]:
    args = ["nix-instantiate", "--strict", "--json", "--eval", "-"]

    if show_trace:
        args.append("--show-trace")

    json_text = subprocess.check_output(args, text=True, input=expr)

    def maybe_path(s: Optional[str]) -> Optional[Path]:
        if s:
            return Path(s)
        return None

    def maybe_location(s: Optional[Dict]) -> Optional[SourceLocation]:
        if s and isinstance(s.get("line"), str) and s["line"]:
            s["line"] = int(s["line"])
            return SourceLocation(**s)
        return None

    result: Dict[Attr, List[Report]] = {}
    for name, data in json.loads(json_text).items():
        attr = Attr(
            name=name,
            output=maybe_path(data.get("outputPath")),
            drv=maybe_path(data.get("drvPath")),
            location=maybe_location(data["location"]),
        )
        result[attr] = load_reports(data["report"])

    return result


def load_reports(report_datas: List[Dict[str, Any]]) -> List[Report]:
    reports = []
    for report_data in report_datas:
        locations = []
        for location_data in report_data.pop("locations", []):
            loc = SourceLocation(**location_data)
            locations.append(loc)
        report = Report(
            locations=locations,
            **report_data,
        )
        reports.append(report)
    return reports


def bold(msg: str) -> str:
    return f"[1m{msg}[0m"


def yellow(msg: str) -> str:
    return f"[1;33m{msg}[0m"


def red(msg: str) -> str:
    return f"[1;31m{msg}[0m"


def green(msg: str) -> str:
    return f"[1;32m{msg}[0m"


def indent(text: str, steps=1) -> str:
    return textwrap.indent(text, " " * 4 * steps)


def render_code_location_ansi(loc: SourceLocation) -> str:
    with open(loc.file, "r") as opened_file:
        all_lines = opened_file.read().splitlines()
        line_contents = all_lines[loc.line - 1]
        line_spaces = " " * len(str(loc.line))
        column = loc.column if loc.column is not None else 1
        pointer = " " * (column - 1) + "^"
        rendered_loc_parts = [loc.file, str(loc.line)]
        if loc.column is not None:
            rendered_loc_parts.append(str(loc.column))

        location_lines = [
            f"Near {':'.join(rendered_loc_parts)}:",
            line_spaces + " |",
            str(loc.line) + " | " + line_contents,
            line_spaces + " | " + pointer,
        ]

    return "\n".join(location_lines)


def render_report_ansi(report: Report) -> str:
    color = red if report.severity == "error" else yellow

    message_lines = [
        color(f"{report.severity}: {report.name}"),
        report.msg,
    ] + [render_code_location_ansi(loc) for loc in report.locations]

    if report.link:
        message_lines.append(
            f"See: https://github.com/jtojnar/nixpkgs-hammering/blob/master/explanations/{report.name}.md",
        )

    return "\n".join(message_lines)


def main(args):
    script_dir = Path(__file__).parent
    overlay_generators_path = (script_dir.parent / "overlays").absolute()
    lib_dir = (script_dir.parent / "lib").absolute()

    try_eval = "(value: { inherit value; success = true; })" if args.do_not_catch_errors else "builtins.tryEval"

    attr_messages = []

    for attr in args.attr_paths:
        attr_messages.append(
            textwrap.dedent(
                f"""
                "{attr}" =
                  let
                    result = {try_eval} pkgs.{escape_attr(attr)} or null;
                    maybeReport = {try_eval} (result.value.__nixpkgs-hammering-state.reports or []);
                  in
                    if !(result.success && maybeReport.success) then
                      {{
                        report = [ {{
                          name = "EvalError";
                          msg = "Cannot evaluate attribute ‘{attr}’ in ‘${{packageSet}}’.";
                          severity = "warning";
                          link = false;
                        }} ];
                        location = null;
                        outputPath = null;
                        drvPath = null;
                      }}
                    else if result.value == null then
                      {{
                        report = [ {{
                          name = "AttrPathNotFound";
                          msg = "Packages in ‘${{packageSet}}’ do not contain ‘{attr}’ attribute.";
                          severity = "error";
                          link = false;
                        }} ];
                        location = null;
                        outputPath = null;
                        drvPath = null;
                      }}
                    else
                      {{
                        report = if maybeReport.success then maybeReport.value else [];
                        location =
                          let
                            position = result.value.meta.position or null;
                            posSplit = builtins.split ":" result.value.meta.position;
                          in
                            if position == null then
                              null
                            else {{
                              file = builtins.elemAt posSplit 0;
                              line = builtins.elemAt posSplit 2;
                            }};
                        outputPath = result.value.outPath;
                        drvPath = result.value.drvPath;
                      }};
                """
            )
        )

    # Our overlays need to know the built attributes so that they can check only them.
    # We do it by using functions that return overlays so we need to instantiate them.
    overlay_expressions = []
    for overlay_generator in overlay_generators_path.glob("*"):
        if overlay_generator.stem in args.excluded_rules:
            continue

        overlay = overlay_generator.name
        overlay_expressions.append(f"(import {overlay_generators_path}/{overlay})")

    attr_messages_nix = (
        "{\n" + indent("".join(attr_messages)) + "\n}" if attr_messages else "{ }"
    )
    overlays_nix = (
        "[\n" + indent("\n".join(overlay_expressions)) + "\n]"
        if overlay_expressions
        else "[ ]"
    )
    all_messages_nix = textwrap.dedent(
        f"""
        let
            packageSet = {args.nix_file};
            cleanPkgs = import packageSet {{ }};

            pkgs = import packageSet {{
                overlays = {indent(overlays_nix, 2 + 2).strip()};
            }};
        in {indent(attr_messages_nix, 2 + 0).strip()}
        """
    )

    if args.show_trace:
        print("Nix expression:", all_messages_nix, file=sys.stderr)

    overlay_data = nix_eval_json(all_messages_nix, args.show_trace)
    overlay_reports = {attr.name: reports for attr, reports in overlay_data.items()}
    external_reports = run_external_checks(
        list(overlay_data.keys()), args.excluded_rules
    )
    all_messages = concatenate_messages(overlay_reports, external_reports)

    if args.json:
        print(json.dumps(all_messages, cls=JSONEncoder))
    else:
        for attr_name, reports in all_messages.items():
            print(bold(f"When evaluating attribute ‘{attr_name}’:"), file=sys.stderr)
            if len(reports) > 0:
                print(
                    "\n".join([render_report_ansi(report) for report in reports]),
                    file=sys.stderr,
                )
            else:
                print(green("No issues found."), file=sys.stderr)
            print(file=sys.stderr)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="nixpkgs-hammer",
        description="check package expressions for common mistakes",
    )
    parser.add_argument(
        "-f",
        "--file",
        dest="nix_file",
        metavar="FILE",
        # Absolutize so we can refer to it from Nix.
        type=lambda p: Path(p).resolve(strict=True),
        # Nix defaults to current directory when file not specified.
        default=Path.cwd(),
        help=(
            "evaluate attributes in given path rather than the default "
            "(current working directory). The path needs to be importable "
            "by Nix and the imported value has to accept attribute set "
            "with overlays attribute as an argument."
        ),
    )
    parser.add_argument(
        "--show-trace",
        dest="show_trace",
        action="store_true",
        help="show trace when error occurs",
    )
    parser.add_argument(
        "--do-not-catch-errors",
        dest="do_not_catch_errors",
        action="store_true",
        help="avoid catching evaluation errors (for debugging)",
    )
    parser.add_argument(
        "--json",
        dest="json",
        action="store_true",
        help="Output results as JSON",
    )
    parser.add_argument(
        "-e",
        "--exclude",
        metavar="rule",
        dest="excluded_rules",
        action="append",
        default=[],
        help="rule to exclude (can be passed repeatedly)",
    )
    parser.add_argument(
        "attr_paths",
        metavar="attr-path",
        nargs="+",
        help="Attribute path of package to update",
    )

    args = parser.parse_args()

    main(args)
