"""
Writes Python egg files.
Supports what's needed for saving and loading components/simulations.
"""

import copy
import os.path
import re
import sys
import zipfile

import pkg_resources
from openmdao.util import eggobserver

__all__ = ('egg_filename', 'write')

# Legal egg strings.
_EGG_NAME_RE = re.compile('[a-zA-Z][_a-zA-Z0-9]*')
_EGG_VERSION_RE = \
    re.compile('([a-zA-Z0-9][_a-zA-Z0-9]*)+(\.[_a-zA-Z0-9][_a-zA-Z0-9]*)*')


def egg_filename(name, version):
    """
    Returns name for egg file as generated by :mod:`setuptools`.

    name: string
        Must be alphanumeric.

    version: string
        Must be alphanumeric.
    """
    assert name and isinstance(name, basestring)
    match = _EGG_NAME_RE.search(name)
    if match is None or match.group() != name:
        raise ValueError('Egg name must be alphanumeric')

    assert version and isinstance(version, basestring)
    match = _EGG_VERSION_RE.search(version)
    if match is None or match.group() != version:
        raise ValueError('Egg version must be alphanumeric')

    name = pkg_resources.to_filename(pkg_resources.safe_name(name))
    version = pkg_resources.to_filename(pkg_resources.safe_version(version))
    return '%s-%s-py%s.egg' % (name, version, sys.version[:3])


def write(name, version, doc, entry_map, src_files, distributions, modules,
          dst_dir, logger, observer=None, compress=True):
    """
    Write egg in the manner of :mod:`setuptools`, with some differences:

    - Writes directly to the zip file, avoiding some intermediate copies.
    - Doesn't compile any Python modules.

    name: string
        Must be an alphanumeric string.

    version: string
        Must be an alphanumeric string.

    doc: string
        Used for the `Summary` and `Description` entries in the egg's metadata.

    entry_map: dict
        A :mod:`pkg_resources` :class:`EntryPoint` map: a dictionary mapping
        group names to dictionaries mapping entry point names to
        :class:`EntryPoint` objects.

    src_files: list
        List of non-Python files to include.

    distributions: list
        List of Distributions this egg depends on. It is used for the `Requires`
        entry in the egg's metadata.

    modules: list
        List of module names not found in a distribution that this egg depends
        on. It is used for the `Requires` entry in the egg's metadata and is
        also recorded in the 'openmdao_orphans.txt' resource.

    dst_dir: string
        The directory to write the egg to.

    logger: Logger
        Used for recording progress, etc.

    observer: callable
        Will be called via an :class:`EggObserver` intermediary.

    Returns the egg's filename.
    """
    observer = eggobserver.EggObserver(observer, logger)

    egg_name = egg_filename(name, version)
    egg_path = os.path.join(dst_dir, egg_name)

    distributions = sorted(distributions, key=lambda dist: dist.project_name)
    modules = sorted(modules)

    sources = []
    files = []
    size = 0  # Approximate (uncompressed) size. Used to set allowZip64 flag.

    # Collect src_files.
    for path in src_files:
        path = os.path.join(name, path)
        files.append(path)
        size += os.path.getsize(path)

    # Collect Python modules.
    for dirpath, dirnames, filenames in os.walk('.', followlinks=True):
        dirs = copy.copy(dirnames)
        for path in dirs:
            if not os.path.exists(os.path.join(dirpath, path, '__init__.py')):
                dirnames.remove(path)
        for path in filenames:
            if path.endswith('.py'):
                path = os.path.join(dirpath[2:], path)  # Skip leading './'
                files.append(path)
                size += os.path.getsize(path)
                sources.append(path)

    # Package info -> EGG-INFO/PKG-INFO
    pkg_info = []
    pkg_info.append('Metadata-Version: 1.1')
    pkg_info.append('Name: %s' % pkg_resources.safe_name(name))
    pkg_info.append('Version: %s' % pkg_resources.safe_version(version))
    pkg_info.append('Summary: %s' % doc.strip().split('\n')[0])
    pkg_info.append('Description: %s' % doc.strip())
    pkg_info.append('Author-email: UNKNOWN')
    pkg_info.append('License: UNKNOWN')
    pkg_info.append('Platform: UNKNOWN')
    for dist in distributions:
        pkg_info.append('Requires: %s (%s)' % (dist.project_name, dist.version))
    for module in modules:
        pkg_info.append('Requires: %s' % module)
    pkg_info = '\n'.join(pkg_info) + '\n'
    sources.append(name + '.egg-info/PKG-INFO')
    size += len(pkg_info)

    # Dependency links -> EGG-INFO/dependency_links.txt
    dependency_links = '\n'
    sources.append(name + '.egg-info/dependency_links.txt')
    size += len(dependency_links)

    # Entry points -> EGG-INFO/entry_points.txt
    entry_points = []
    for entry_group in sorted(entry_map.keys()):
        entry_points.append('[%s]' % entry_group)
        for entry_name in sorted(entry_map[entry_group].keys()):
            entry_points.append('%s' % entry_map[entry_group][entry_name])
        entry_points.append('')
    entry_points = '\n'.join(entry_points) + '\n'
    sources.append(name + '.egg-info/entry_points.txt')
    size += len(entry_points)

    # Unsafe -> EGG-INFO/not-zip-safe
    not_zip_safe = '\n'
    sources.append(name + '.egg-info/not-zip-safe')
    size += len(not_zip_safe)

    # Requirements -> EGG-INFO/requires.txt
    requirements = [str(dist.as_requirement()) for dist in distributions]
    requirements = '\n'.join(requirements) + '\n'
    sources.append(name + '.egg-info/requires.txt')
    size += len(requirements)

    # Modules not part of a distribution -> EGG-INFO/openmdao_orphans.txt
    orphans = '\n'.join(modules) + '\n'
    sources.append(name + '.egg-info/openmdao_orphans.txt')
    size += len(orphans)

    # Top-level names -> EGG-INFO/top_level.txt
    top_level = '%s\n' % name
    sources.append(name + '.egg-info/top_level.txt')
    size += len(top_level)

    # Manifest -> EGG-INFO/SOURCES.txt
    sources.append(name + '.egg-info/SOURCES.txt')
    sources = '\n'.join(sorted(sources)) + '\n'
    size += len(sources)

    # Open zipfile.
    logger.debug('Creating %s', egg_path)
    zip64 = size > zipfile.ZIP64_LIMIT
    compression = zipfile.ZIP_DEFLATED if compress else zipfile.ZIP_STORED
    egg = zipfile.ZipFile(egg_path, 'w', compression, zip64)

    stats = {'completed_files': 0., 'total_files': float(8 + len(files)),
             'completed_bytes': 0., 'total_bytes': float(size)}

    # Write egg info.
    _write_info(egg, 'PKG-INFO', pkg_info, observer, stats)
    _write_info(egg, 'dependency_links.txt', dependency_links, observer, stats)
    _write_info(egg, 'entry_points.txt', entry_points, observer, stats)
    _write_info(egg, 'not-zip-safe', not_zip_safe, observer, stats)
    _write_info(egg, 'requires.txt', requirements, observer, stats)
    _write_info(egg, 'openmdao_orphans.txt', orphans, observer, stats)
    _write_info(egg, 'top_level.txt', top_level, observer, stats)
    _write_info(egg, 'SOURCES.txt', sources, observer, stats)

    # Write collected files.
    for path in sorted(files):
        _write_file(egg, path, observer, stats)

    observer.complete(egg_name)

    egg.close()
    if os.path.getsize(egg_path) > zipfile.ZIP64_LIMIT:
        logger.warning('Egg zipfile requires Zip64 support to unzip.')
    return egg_name


def _write_info(egg, name, info, observer, stats):
    """ Write info string to egg. """
    path = os.path.join('EGG-INFO', name)
    observer.add(path, stats['completed_files'] / stats['total_files'],
                 stats['completed_bytes'] / stats['total_bytes'])
    egg.writestr(path, info)
    stats['completed_files'] += 1
    stats['completed_bytes'] += len(info)


def _write_file(egg, path, observer, stats):
    """ Write file to egg. """
    observer.add(path, stats['completed_files'] / stats['total_files'],
                 stats['completed_bytes'] / stats['total_bytes'])
    egg.write(path)
    stats['completed_files'] += 1
    stats['completed_bytes'] += os.path.getsize(path)
