"""Python source expertise for coverage.py"""

import os.path
import sys
import tokenize
import zipimport

from coverage import env
from coverage.backward import unicode_class
from coverage.files import FileLocator
from coverage.misc import NoSource, join_regex
from coverage.parser import PythonParser
from coverage.phystokens import source_token_lines, source_encoding
from coverage.plugin import FileReporter


def read_python_source(filename):
    """Read the Python source text from `filename`.

    Returns a str: unicode on Python 3, bytes on Python 2.

    """
    # Python 3.2 provides `tokenize.open`, the best way to open source files.
    if sys.version_info >= (3, 2):
        f = tokenize.open(filename)
    else:
        f = open(filename, "rU")

    with f:
        return f.read()


def get_python_source(filename):
    """Return the source code, as a str."""
    base, ext = os.path.splitext(filename)
    if ext == ".py" and env.WINDOWS:
        exts = [".py", ".pyw"]
    else:
        exts = [ext]

    for ext in exts:
        try_filename = base + ext
        if os.path.exists(try_filename):
            # A regular text file: open it.
            source = read_python_source(try_filename)
            break

        # Maybe it's in a zip file?
        source = get_zip_bytes(try_filename)
        if source is not None:
            if env.PY3:
                source = source.decode(source_encoding(source))
            break
    else:
        # Couldn't find source.
        raise NoSource("No source for code: '%s'." % filename)

    # Python code should always end with a line with a newline.
    if source and source[-1] != '\n':
        source += '\n'

    return source


def get_zip_bytes(filename):
    """Get data from `filename` if it is a zip file path.

    Returns the bytestring data read from the zip file, or None if no zip file
    could be found or `filename` isn't in it.  The data returned will be
    an empty string if the file is empty.

    """
    markers = ['.zip'+os.sep, '.egg'+os.sep]
    for marker in markers:
        if marker in filename:
            parts = filename.split(marker)
            try:
                zi = zipimport.zipimporter(parts[0]+marker[:-1])
            except zipimport.ZipImportError:
                continue
            try:
                data = zi.get_data(parts[1])
            except IOError:
                continue
            assert isinstance(data, bytes)
            return data
    return None


class PythonFileReporter(FileReporter):
    """Report support for a Python file."""

    def __init__(self, morf, coverage=None):
        self.coverage = coverage
        file_locator = coverage.file_locator if coverage else FileLocator()

        if hasattr(morf, '__file__'):
            filename = morf.__file__
        else:
            filename = morf

        # .pyc files should always refer to a .py instead.
        if filename.endswith(('.pyc', '.pyo')):
            filename = filename[:-1]
        elif filename.endswith('$py.class'):   # Jython
            filename = filename[:-9] + ".py"

        super(PythonFileReporter, self).__init__(
            file_locator.canonical_filename(filename)
        )

        if hasattr(morf, '__name__'):
            name = morf.__name__
            name = name.replace(".", os.sep) + ".py"
        else:
            name = file_locator.relative_filename(filename)
        self.name = name

        self._source = None
        self._parser = None
        self._statements = None
        self._excluded = None

    @property
    def parser(self):
        if self._parser is None:
            self._parser = PythonParser(
                filename=self.filename,
                exclude=self.coverage._exclude_regex('exclude'),
            )
        return self._parser

    def statements(self):
        """Return the line numbers of statements in the file."""
        if self._statements is None:
            self._statements, self._excluded = self.parser.parse_source()
        return self._statements

    def excluded_statements(self):
        """Return the line numbers of statements in the file."""
        if self._excluded is None:
            self._statements, self._excluded = self.parser.parse_source()
        return self._excluded

    def translate_lines(self, lines):
        return self.parser.translate_lines(lines)

    def translate_arcs(self, arcs):
        return self.parser.translate_arcs(arcs)

    def no_branch_lines(self):
        no_branch = self.parser.lines_matching(
            join_regex(self.coverage.config.partial_list),
            join_regex(self.coverage.config.partial_always_list)
            )
        return no_branch

    def arcs(self):
        return self.parser.arcs()

    def exit_counts(self):
        return self.parser.exit_counts()

    def source(self):
        if self._source is None:
            self._source = get_python_source(self.filename)
            if env.PY2:
                encoding = source_encoding(self._source)
                self._source = self._source.decode(encoding, "replace")
            assert isinstance(self._source, unicode_class)
        return self._source

    def should_be_python(self):
        """Does it seem like this file should contain Python?

        This is used to decide if a file reported as part of the execution of
        a program was really likely to have contained Python in the first
        place.

        """
        # Get the file extension.
        _, ext = os.path.splitext(self.filename)

        # Anything named *.py* should be Python.
        if ext.startswith('.py'):
            return True
        # A file with no extension should be Python.
        if not ext:
            return True
        # Everything else is probably not Python.
        return False

    def source_token_lines(self):
        return source_token_lines(self.source())