"""
run-clang-tidy is a tool to automate clang-tidy invocations over a project.
It was inspired by https://raw.githubusercontent.com/llvm/llvm-project/main/clang-tools-extra/clang-tidy/tool/run-clang-tidy.py

The key difference between this script and the llvm project, is the ability
to run checks on a per-file basis. The run-clang-tidy.yaml file is a new
configuration which enables specifying clang-tidy checks based on regex file names.

The structure of the configuration is as follows:
   
  # A list of path regex to limit where files are searched.
  # A file must match at least one regex to have clang-tidy run.
  # Avoids running clang-tidy against generated code in the build dir. 
  search_paths:
    - "^src/.*"

  # A list of path regex to exclude from the search path.
  # Any match here excludes clang-tidy from being run on the file.
  excludes:
    - "src/some_problematic_file\\.cpp"

  # A list of objects containing a check field and an optional set of search filters.
  # Defines which checks are enabled for clang-tidy invocations.
  # If a search filter is not supplied, the check is run against all files.
  checks:
    - value: "cppcoreguidelines-*"

    # In this example, we can enable cppcoreguidelines but disable magic number checks 
    # only on test files.
    - value: "-cppcoreguidelines-avoid-magic-numbers"
      search_filters:
        - ".*\\.t\\.cpp"

  # A list of objects containing a check field and an optional set search filters.
  # Defines which check failures are reported as errors.
  # Typically we set this to include all files by setting a single '*' value
  # Follows the same rules as the `checks` field.
  warnings_as_errors:
    - value: "*"

When running the script, it will by default use the search_paths field for determining
files to invoke clang-tidy on. If posargs are supplied, this replaces the search_paths
values for the script run.
"""

import argparse
import asyncio
import collections
import dataclasses
import json
import logging
import multiprocessing
import os
import re
import sys
import time
from typing import Any, Awaitable, Callable, TypeVar

import yaml

logger = logging.getLogger(__name__)


T = TypeVar("T")


async def run_with_semaphore(
    semaphore: asyncio.Semaphore,
    f: Callable[..., Awaitable[T]],
    *args: Any,
) -> T:
    async with semaphore:
        return await f(*args)


@dataclasses.dataclass
class ClangTidyRequest:
    filename: str
    build_path: str
    reports_path: str
    checks: list[str]
    warnings_as_errors: list[str]

    @property
    def command(self) -> list[str]:
        cmd = [
            "clang-tidy",
            self.filename,
            f"-p={self.build_path}",
            f"-checks={','.join(self.checks)}",
            f"-export-fixes={self.reports_path}",
        ]
        if self.warnings_as_errors:
            cmd.append(f"-warnings-as-errors={','.join(self.warnings_as_errors)}")
        return cmd


@dataclasses.dataclass
class ClangTidyResult:
    filename: str
    command: list[str]
    duration: float
    exitcode: int
    raw_output: str
    diagnostics: list[dict[str, Any]]


async def run_clang_tidy(
    request: ClangTidyRequest,
    on_start: Callable[[ClangTidyRequest], None],
    on_end: Callable[[ClangTidyResult], None],
) -> ClangTidyResult:
    on_start(request)

    start_time = time.time()
    process = await asyncio.create_subprocess_exec(
        *request.command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
    )
    stdout, _ = await process.communicate()
    end_time = time.time()
    exitcode = process.returncode if process.returncode is not None else -1

    diagnostics = []
    if os.path.exists(request.reports_path):
        with open(request.reports_path, "r") as f:
            report = yaml.safe_load(f)
            diagnostics = report.get("Diagnostics", [])
        if os.path.getsize(request.reports_path) == 0:
            os.remove(request.reports_path)

    result = ClangTidyResult(
        filename=request.filename,
        command=request.command,
        duration=end_time - start_time,
        exitcode=exitcode,
        raw_output=stdout.decode(),
        diagnostics=diagnostics,
    )
    on_end(result)
    return result


async def main() -> int:
    parser = argparse.ArgumentParser(description="run-clang-tidy")
    parser.add_argument(
        "-c", "--config-path", type=str, help="Path to the run-clang-tidy config file"
    )
    parser.add_argument(
        "-p", "--build-path", type=str, help="Path to the build directory"
    )
    parser.add_argument(
        "-j",
        "--jobs",
        type=int,
        default=multiprocessing.cpu_count(),
        help="Number of parallel jobs",
    )
    parser.add_argument("-r", "--reports-path", help="Output directory for clang-tidy")
    parser.add_argument(
        "search_paths", nargs="*", help="Search paths for files to analyze"
    )
    args = parser.parse_args()

    if not os.path.isfile(args.config_path):
        logger.error(f"Error: Config file {args.config_path} does not exist.")
        return 1

    if not args.build_path or not os.path.isdir(args.build_path):
        logger.error(f"Error: Build directory {args.build_path} does not exist.")
        return 1

    with open(args.config_path, "r") as file:
        config = yaml.safe_load(file)

    reports_path = os.path.join(args.reports_path, time.strftime("%Y%m%dT%H%M%S"))
    os.makedirs(reports_path, exist_ok=True)

    logger.info(f"Using config: {args.config_path}")
    logger.info(f"Using build directory: {args.build_path}")
    logger.info(f"Using reports directory: {reports_path}")

    # Load the database and extract all files.
    with open(os.path.join(args.build_path, "compile_commands.json")) as f:
        database = json.load(f)

    cwd = os.getcwd()
    files = {os.path.abspath(os.path.join(e["directory"], e["file"])) for e in database}
    files = {os.path.relpath(f, cwd) for f in files}

    search_paths = args.search_paths or config.get("search_paths", [])
    files = {f for f in files if any(re.match(p, f) for p in search_paths)}

    excludes = config.get("excludes", [])
    exclusions = {f for f in files if any(re.match(p, f) for p in excludes)}
    if exclusions:
        logger.warning(f"Excluding files: {' '.join(exclusions)}")
    files -= exclusions

    checks = {
        filename: [
            item["value"]
            for item in config.get("checks", [])
            if not item.get("search_filters")
            or any(re.match(p, filename) for p in item["search_filters"])
        ]
        for filename in files
    }

    warnings_as_errors = {
        filename: [
            item["value"]
            for item in config.get("warnings_as_errors", [])
            if not item.get("search_filters")
            or any(re.match(p, filename) for p in item["search_filters"])
        ]
        for filename in files
    }

    started = 0
    total = len(files)

    def on_start(request: ClangTidyRequest) -> None:
        nonlocal started
        started += 1
        logger.info(f"[{started}/{total}] {' '.join(request.command)}")

    def on_end(result: ClangTidyResult) -> None:
        if result.exitcode == 0:
            logger.info(
                f"clang-tidy succeeded for {result.filename}"
                f" in {result.duration:.2f} seconds"
            )
        else:
            logger.error(
                f"clang-tidy failed for {result.filename}"
                f" in {result.duration:.2f} seconds"
                f"\n{' '.join(result.command)}"
                f"\n{result.raw_output}"
            )

    requests = [
        ClangTidyRequest(
            filename=filename,
            checks=checks[filename],
            warnings_as_errors=warnings_as_errors[filename],
            build_path=args.build_path,
            reports_path=os.path.join(
                reports_path, f"{os.path.basename(filename)}.yaml"
            ),
        )
        for filename in files
    ]
    semaphore = asyncio.Semaphore(args.jobs)
    tasks = [
        asyncio.create_task(
            run_with_semaphore(semaphore, run_clang_tidy, r, on_start, on_end)
        )
        for r in requests
    ]
    results = await asyncio.gather(*tasks)
    failed_results = [r for r in results if r.exitcode != 0]

    if not failed_results:
        logger.info("No warnings reported. Success! :)")
        return 0

    logger.error("The following files failed:")
    for result in sorted(failed_results, key=lambda r: r.filename):
        diagnostic_counts: dict[str, int] = collections.defaultdict(int)
        for diagnostic in result.diagnostics:
            diagnostic_counts[diagnostic["DiagnosticName"]] += 1

        logger.error(f"{result.filename} ({len(result.diagnostics)} warnings)")
        for name, count in sorted(diagnostic_counts.items()):
            logger.error(f"  {name} ({count} warnings)")
    return 1


class ColorFormatter(logging.Formatter):
    def format_time(self, record: Any, datefmt: Any | None = None) -> str:
        ct = self.converter(record.created)
        t = time.strftime("%Y-%m-%d %H:%M:%S", ct)
        s = f"{t}.{int(record.msecs):03d}{time.strftime('%z', ct)}"
        return s

    def format(self, record: Any) -> str:
        log_msg = super().format(record)
        timestamp = self.format_time(record)
        level = record.levelname
        if record.levelno == logging.ERROR:
            prefix = f"\033[91m{timestamp} [{level}]\033[0m"  # Red
        elif record.levelno == logging.WARNING:
            prefix = f"\033[93m{timestamp} [{level}]\033[0m"  # Yellow
        elif record.levelno == logging.INFO:
            prefix = f"\033[92m{timestamp} [{level}]\033[0m"  # Green
        else:
            prefix = f"{timestamp} [{level}]"
        return f"{prefix}: {log_msg}"


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(message)s")

    handler = logging.StreamHandler()
    handler.setFormatter(ColorFormatter())
    logger.addHandler(handler)
    logger.propagate = False

    try:
        sys.exit(asyncio.run(main()))
    except KeyboardInterrupt:
        logger.info("Process interrupted by user")
        sys.exit(1)
