import os
import sys

from coverage.data import CoverageData

from pex.pex_builder import BOOTSTRAP_DIR


def _iter_filter(root_dir, data_dict):
    root_fragment = os.path.join(root_dir, "pex/")
    pex_fragment = "/%s/pex/" % BOOTSTRAP_DIR

    for filename, records in data_dict.items():
        # already acceptable coverage
        if filename.startswith(root_fragment):
            yield (filename, dict((record, None) for record in records))
            continue

        # possible it's coverage from within a pex environment
        try:
            bi = filename.index(pex_fragment)
        except ValueError:
            continue

        # rewrite to make up for fact that each pex environment is ephemeral.
        yield (
            os.path.join(root_dir, "pex", filename[bi + len(pex_fragment) :]),
            dict((record, None) for record in records),
        )


def combine_pex_coverage(root_dir, coverage_file_iter, unlink=True):
    combined = CoverageData(basename=".coverage")

    for filename in coverage_file_iter:
        cov = CoverageData(basename=filename)
        cov.read()
        combined.add_line_data(dict(_iter_filter(root_dir, cov.line_data())))
        combined.add_arc_data(dict(_iter_filter(root_dir, cov.arc_data())))
        if unlink:
            os.unlink(filename)

    # filter out non-pex files
    prefix = os.path.join(root_dir, "pex/")
    non_pex = [filename for filename in combined.lines if not filename.startswith(prefix)]
    for filename in non_pex:
        combined.lines.pop(filename)

    non_pex = [filename for filename in combined.arcs if not filename.startswith(prefix)]
    for filename in non_pex:
        combined.arcs.pop(filename)

    combined.write()
    return combined.filename


def main(args):
    script = os.path.realpath(args[0])
    root = os.path.realpath(os.path.join(os.path.dirname(script), ".."))

    coverage_iterator = [
        os.path.join(root, filename)
        for filename in os.listdir(root)
        if filename.startswith(".coverage")
    ]

    combine_pex_coverage(root, coverage_iterator, unlink=False)


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