"""Counts the lines of source in the tree, reported as a breakdown by project folder."""
import argparse
import os
import re
import sys
import unittest


REPO_ROOT = os.path.abspath(os.path.join(os.path.abspath(__file__), "..", ".."))

# Games\ is missing entirely from a checkout without the private overlay, so it's skipped rather than required.
SOURCE_ROOTS = ["Modules", "Games"]

LANGUAGE_BY_EXTENSION = {
    ".cs": "C#",
    ".fx": "HLSL",
    ".fxh": "HLSL",
    ".fxt": "HLSL",
    ".hlsl": "HLSL",
    ".hlsli": "HLSL",
    ".js": "JavaScript",
    ".lua": "Lua",
}

SOURCE_EXTENSIONS = frozenset(LANGUAGE_BY_EXTENSION)

# Build output and IDE state that lands inside the source tree.
EXCLUDED_DIRECTORY_NAMES = frozenset([".svn", ".vs", "bin", "obj", "packages", "node_modules"])

# Lowercased directory name endings that mark test code, for --no-tests.  Suffixes rather than substrings so that
# UnitTestSharp, a shipped module, isn't mistaken for the tests of one.
TEST_DIRECTORY_SUFFIXES = ("tests", "testing", "testbed", "testcontainer")

# (line comment, block open, block close), keyed by extension.
DEFAULT_COMMENT_SYNTAX = ("//", "/*", "*/")
COMMENT_SYNTAX_BY_EXTENSION = {".lua": ("--", "--[[", "]]")}

# Where the numbers start, leaving room for the deepest names the tree reaches.
NAME_COLUMN_WIDTH = 56


class Counts(object):
    """A tally of lines, split by what each one holds."""
    def __init__(self):
        self.total = 0
        self.code = 0
        self.comment = 0
        self.blank = 0

    def add_line(self, has_code, has_comment):
        """Files a single line under exactly one category, code winning over comment."""
        self.total += 1
        if has_code:
            self.code += 1
        elif has_comment:
            self.comment += 1
        else:
            self.blank += 1

    def add(self, other):
        """Folds another tally into this one."""
        self.total += other.total
        self.code += other.code
        self.comment += other.comment
        self.blank += other.blank


class Breakdown(object):
    """A flat set of tallies, one per bucket, for the summaries that follow the tree."""
    def __init__(self):
        self.buckets = {}

    def add(self, bucket, counts):
        """Folds a tally into its bucket, opening the bucket the first time it's needed."""
        if bucket not in self.buckets:
            self.buckets[bucket] = Counts()
        self.buckets[bucket].add(counts)


def scan_verbatim_string(line, index):
    """Scans a C# @"..." string, in which the only escape is a doubled quote.

    Args:
        line: The line being scanned.
        index: Where in the line to pick the string up, already past any opening quote.

    Returns:
        An (index, still_open) tuple, where still_open says the string runs onto the next line.
    """
    while index < len(line):
        if line[index] != "\"":
            index += 1
        elif index + 1 < len(line) and line[index + 1] == "\"":
            index += 2
        else:
            return index + 1, False
    return len(line), True


def skip_literal(line, index):
    """Steps over a string or character literal so its contents can't be mistaken for a comment.

    Args:
        line: The line being scanned.
        index: The position to examine.  Anything that isn't the start of a literal advances one character.

    Returns:
        An (index, in_verbatim_string) tuple, where in_verbatim_string says a @"..." string runs onto the next line.
    """
    # A verbatim string opens with any mix of @ and $, so long as an @ is in there somewhere.
    scan = index
    is_verbatim = False
    while scan < len(line) and line[scan] in "@$":
        is_verbatim = is_verbatim or line[scan] == "@"
        scan += 1

    if scan == len(line) or line[scan] not in "\"'":
        return index + 1, False

    quote = line[scan]
    scan += 1
    if is_verbatim and quote == "\"":
        return scan_verbatim_string(line, scan)

    while scan < len(line):
        if line[scan] == "\\":
            scan += 2
        elif line[scan] == quote:
            return scan + 1, False
        else:
            scan += 1
    return len(line), False


_TRIGGERS_BY_SYNTAX = {}


def triggers_for(syntax):
    """Returns the characters that could open a comment or a literal, and a pattern that finds the next one.

    Most lines hold none of them, so being able to skip the character-at-a-time scan is worth about 5x overall.
    """
    if syntax not in _TRIGGERS_BY_SYNTAX:
        line_comment, block_open, _ = syntax
        characters = frozenset("\"'@$" + line_comment[0] + block_open[0])
        pattern = re.compile("[{0}]".format(re.escape("".join(sorted(characters)))))
        _TRIGGERS_BY_SYNTAX[syntax] = (characters, pattern)
    return _TRIGGERS_BY_SYNTAX[syntax]


def count_lines(text, syntax):
    """Tallies the lines of one file's text.

    Args:
        text: The whole file, already decoded.
        syntax: A (line comment, block open, block close) tuple.

    Returns:
        A Counts.
    """
    line_comment, block_open, block_close = syntax
    trigger_characters, trigger_pattern = triggers_for(syntax)
    counts = Counts()
    in_block_comment = False
    in_verbatim_string = False

    for line in text.splitlines():
        if not in_block_comment and not in_verbatim_string and not trigger_pattern.search(line):
            counts.add_line(bool(line.strip()), False)
            continue

        has_code = False
        has_comment = False
        index = 0
        while index < len(line):
            if in_block_comment:
                has_comment = True
                end = line.find(block_close, index)
                if end < 0:
                    index = len(line)
                else:
                    index = end + len(block_close)
                    in_block_comment = False
            elif in_verbatim_string:
                has_code = True
                index, in_verbatim_string = scan_verbatim_string(line, index)
            elif line[index] not in trigger_characters:
                # Nothing between here and the next trigger can be anything but whitespace or code.
                match = trigger_pattern.search(line, index)
                end = match.start() if match else len(line)
                has_code = has_code or bool(line[index:end].strip())
                index = end
            # Checked before the line comment because Lua's "--[[" opens with its line comment marker.
            elif line.startswith(block_open, index):
                has_comment = True
                in_block_comment = True
                index += len(block_open)
            elif line.startswith(line_comment, index):
                has_comment = True
                index = len(line)
            else:
                has_code = True
                index, in_verbatim_string = skip_literal(line, index)
        counts.add_line(has_code, has_comment)

    return counts


def read_text(path):
    """Reads a source file, tolerating whichever encoding it happens to have been saved in."""
    with open(path, "rb") as source_file:
        data = source_file.read()
    for bom, encoding in ((b"\xff\xfe", "utf-16-le"), (b"\xfe\xff", "utf-16-be")):
        if data.startswith(bom):
            return data[len(bom):].decode(encoding, errors="replace")
    return data.decode("utf-8-sig", errors="replace")


def count_file(path):
    """Tallies a single source file."""
    extension = os.path.splitext(path)[1].lower()
    return count_lines(read_text(path), COMMENT_SYNTAX_BY_EXTENSION.get(extension, DEFAULT_COMMENT_SYNTAX))


def is_test_directory(name):
    """Whether a directory holds tests rather than the code under test."""
    return name.lower().endswith(TEST_DIRECTORY_SUFFIXES)


def walk_source_tree(root, include_tests):
    """Yields directory, filenames for every directory worth looking at under `root`."""
    for directory, subdirectories, filenames in os.walk(root):
        subdirectories[:] = [name for name in subdirectories
                             if name.lower() not in EXCLUDED_DIRECTORY_NAMES
                             and (include_tests or not is_test_directory(name))]
        yield directory, filenames


def find_node_directories(roots, include_tests):
    """Finds the directories the report has a row for.

    Args:
        roots: Absolute paths of the source roots to scan.
        include_tests: False to leave test projects out entirely.

    Returns:
        A (project_labels, node_directories) tuple.  project_labels maps a project's directory to the name to show it
        under; node_directories additionally holds every folder along the way, which is what groups a module's
        projects together under one row.
    """
    project_labels = {}
    for root in roots:
        for directory, filenames in walk_source_tree(root, include_tests):
            projects = [name for name in filenames if name.lower().endswith(".csproj")]
            if projects:
                project_labels[directory] = os.path.splitext(projects[0])[0]

    node_directories = set(project_labels)
    for project_directory in project_labels:
        directory = os.path.dirname(project_directory)
        while directory.startswith(REPO_ROOT) and directory != REPO_ROOT:
            node_directories.add(directory)
            directory = os.path.dirname(directory)

    return project_labels, node_directories


def find_owning_directory(directory, node_directories):
    """Returns the nearest ancestor of `directory` (itself included) that the report has a row for, or None."""
    while directory.startswith(REPO_ROOT) and directory != REPO_ROOT:
        if directory in node_directories:
            return directory
        directory = os.path.dirname(directory)
    return None


class Node(object):
    """One row of the report: a project, or a folder grouping several of them."""
    def __init__(self, name):
        self.name = name
        self.children = {}
        self.own = Counts()
        self.own_file_count = 0

    def child(self, name):
        """Returns the named child, creating it the first time it's asked for."""
        if name not in self.children:
            self.children[name] = Node(name)
        return self.children[name]

    def totals(self):
        """Returns this row's Counts, including everything nested under it."""
        counts = Counts()
        counts.add(self.own)
        for child in self.children.values():
            counts.add(child.totals())
        return counts

    def file_count(self):
        """Returns how many files this row covers, including everything nested under it."""
        return self.own_file_count + sum(child.file_count() for child in self.children.values())


class Survey(object):
    """Everything one pass over the source tree produces: the folder tree, plus the summaries cut across it."""
    def __init__(self):
        self.tree = Node("Total")
        self.languages = Breakdown()
        self.kinds = Breakdown()


def build_survey(roots, include_tests):
    """Reads every source file under `roots` and files its lines into a tree of project folders."""
    project_labels, node_directories = find_node_directories(roots, include_tests)
    survey = Survey()

    for root in roots:
        for directory, filenames in walk_source_tree(root, include_tests):
            sources = [name for name in filenames if os.path.splitext(name)[1].lower() in SOURCE_EXTENSIONS]
            if not sources:
                continue

            owner = find_owning_directory(directory, node_directories)
            if owner is None:
                continue

            relative = os.path.relpath(directory, REPO_ROOT).split(os.sep)
            kind = "Tests" if any(is_test_directory(component) for component in relative) else "Code"

            node = survey.tree
            for component in os.path.relpath(owner, REPO_ROOT).split(os.sep):
                node = node.child(component)
            node.name = project_labels.get(owner, node.name)

            for filename in sources:
                counts = count_file(os.path.join(directory, filename))
                node.own.add(counts)
                node.own_file_count += 1
                survey.languages.add(LANGUAGE_BY_EXTENSION[os.path.splitext(filename)[1].lower()], counts)
                survey.kinds.add(kind, counts)

    return survey


def measure(counts, use_code):
    """Returns whichever of a tally's numbers the report is showing."""
    return counts.code if use_code else counts.total


def format_row(depth, name, value, share=None):
    """Renders one row, dot-leadered out to the number column, optionally followed by its share of the total."""
    label = "{0}{1} ".format("  " * depth, name)
    row = "{0} {1:>9,}".format(label.ljust(NAME_COLUMN_WIDTH, "."), value)
    return row if share is None else "{0}   {1:>5.1f}%".format(row, share)


def render(node, use_code, max_depth, depth=0):
    """Yields the tree's rows, largest subtree first."""
    yield format_row(depth, node.name, measure(node.totals(), use_code))

    if max_depth is not None and depth >= max_depth:
        return

    for child in sorted(node.children.values(), key=lambda child: measure(child.totals(), use_code), reverse=True):
        for row in render(child, use_code, max_depth, depth + 1):
            yield row

    # Files that sit in a grouping folder rather than in any one of its projects, shared between them.
    if node.children and node.own.total:
        yield format_row(depth + 1, "(shared between these)", measure(node.own, use_code))


def render_breakdown(breakdown, bucket_names, use_code, total):
    """Yields a flat summary section, each row carrying its share of the whole."""
    for name in bucket_names:
        value = measure(breakdown.buckets[name], use_code)
        yield format_row(1, name, value, 100.0 * value / total if total else 0.0)


def largest_first(breakdown, use_code):
    """Returns a breakdown's bucket names, biggest first."""
    return sorted(breakdown.buckets,
                  key=lambda name: measure(breakdown.buckets[name], use_code), reverse=True)


def report(arguments):
    """Counts the tree and prints the breakdown.  Returns a process exit code."""
    roots = [os.path.join(REPO_ROOT, name) for name in SOURCE_ROOTS]
    roots = [root for root in roots if os.path.isdir(root)]
    if not roots:
        print("No source roots found under {0}.".format(REPO_ROOT))
        return 1

    survey = build_survey(roots, include_tests=not arguments.no_tests)
    counts = survey.tree.totals()
    total = measure(counts, arguments.code)

    print("{0}: {1:,} files, {2:,} lines ({3:,} code, {4:,} comment, {5:,} blank)".format(
        REPO_ROOT, survey.tree.file_count(), counts.total, counts.code, counts.comment, counts.blank))
    print("Counting {0}.{1}".format("code lines only" if arguments.code else "every line",
                                    "  Tests excluded." if arguments.no_tests else ""))
    print()
    for row in render(survey.tree, use_code=arguments.code, max_depth=arguments.depth):
        print(row)

    print()
    print("By language")
    for row in render_breakdown(survey.languages, largest_first(survey.languages, arguments.code),
                                arguments.code, total):
        print(row)

    # Nothing to compare once --no-tests has already thrown one side of it away.
    if not arguments.no_tests:
        print()
        print("Code vs tests (unit tests, testbeds and performance tests)")
        kinds = [name for name in ("Code", "Tests") if name in survey.kinds.buckets]
        for row in render_breakdown(survey.kinds, kinds, arguments.code, total):
            print(row)

    return 0


class CountLinesTests(unittest.TestCase):
    """Covers the line scanner, the one piece here with anywhere to hide a bug."""
    def count(self, text, extension=".cs"):
        counts = count_lines(text, COMMENT_SYNTAX_BY_EXTENSION.get(extension, DEFAULT_COMMENT_SYNTAX))
        return counts.total, counts.code, counts.comment, counts.blank

    def test_categorises_each_line_once(self):
        self.assertEqual((4, 1, 2, 1), self.count("// a\nint x = 0;\n\n/* b */"))

    def test_code_wins_over_a_trailing_comment(self):
        self.assertEqual((1, 1, 0, 0), self.count("int x = 0; // why"))

    def test_block_comment_spans_lines(self):
        self.assertEqual((3, 0, 3, 0), self.count("/* one\n   two\n   three */"))

    def test_code_after_a_block_comment_closes_counts_as_code(self):
        self.assertEqual((2, 1, 1, 0), self.count("/* one\n   two */ int x = 0;"))

    def test_comment_markers_inside_a_string_are_not_comments(self):
        self.assertEqual((2, 2, 0, 0), self.count("var path = \"//server/share\";\nint x = 0;"))

    def test_block_opener_inside_a_string_does_not_swallow_the_file(self):
        self.assertEqual((2, 2, 0, 0), self.count("var glob = \"/*.cs\";\nint x = 0;"))

    def test_escaped_quote_does_not_end_a_string(self):
        self.assertEqual((2, 2, 0, 0), self.count("var quoted = \"say \\\" /*\";\nint x = 0;"))

    def test_verbatim_string_spans_lines(self):
        self.assertEqual((3, 3, 0, 0), self.count("var svg = @\"<path\n   d = /* not a comment\n\">;"))

    def test_interpolated_verbatim_string_spans_lines(self):
        self.assertEqual((2, 2, 0, 0), self.count("var svg = $@\"<path\n   {value} \";"))

    def test_doubled_quote_does_not_end_a_verbatim_string(self):
        self.assertEqual((2, 2, 0, 0), self.count("var quoted = @\"a \"\" /* b\";\nint x = 0;"))

    def test_blank_line_inside_a_verbatim_string_counts_as_blank(self):
        self.assertEqual((3, 2, 0, 1), self.count("var text = @\"one\n\ntwo\";"))

    def test_char_literal_holding_a_quote(self):
        self.assertEqual((2, 2, 0, 0), self.count("var quote = '\"';\nint x = 0;"))

    def test_lua_line_comment(self):
        self.assertEqual((2, 1, 1, 0), self.count("-- a comment\nlocal x = 0", extension=".lua"))

    def test_lua_block_comment(self):
        self.assertEqual((2, 0, 2, 0), self.count("--[[ one\n     two ]]", extension=".lua"))

    def test_whitespace_only_line_is_blank(self):
        self.assertEqual((1, 0, 0, 1), self.count("   \t  "))


def main():
    """Parses the command line and runs the requested job."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--code", action="store_true",
                        help="Count only lines holding code, skipping blank and comment-only ones.")
    parser.add_argument("--no-tests", action="store_true",
                        help="Leave out test projects (unit tests, testbeds and performance tests).")
    parser.add_argument("--depth", type=int, default=None,
                        help="Stop the tree at this many levels.  1 shows just the source roots, 2 the modules.")
    parser.add_argument("--self-test", action="store_true", help="Run this script's own tests and exit.")
    arguments = parser.parse_args()

    if arguments.self_test:
        suite = unittest.TestLoader().loadTestsFromTestCase(CountLinesTests)
        return 0 if unittest.TextTestRunner(verbosity=2).run(suite).wasSuccessful() else 1

    return report(arguments)


if __name__ == "__main__":
    sys.exit(main())