# ==================================================================================================
# Copyright 2012 Twitter, Inc.
# --------------------------------------------------------------------------------------------------
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this work except in compliance with the License.
# You may obtain a copy of the License in the LICENSE file, or at:
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==================================================================================================

from __future__ import print_function

from collections import defaultdict
import os
import pprint
import subprocess
import sys
import time

from . import Command

from twitter.common.collections import OrderedSet
from twitter.common.contextutil import pushd
from twitter.common.dirutil import safe_delete, safe_rmtree
from twitter.common.dirutil.chroot import Chroot
from twitter.common.python.pex import PEX

from twitter.pants.base import Address, Config, Target

from twitter.pants.targets import (
  PythonBinary,
  PythonLibrary,
  PythonRequirement,
  PythonTarget,
  PythonThriftLibrary)


SETUP_BOILERPLATE = """
# DO NOT EDIT THIS FILE -- AUTOGENERATED BY PANTS
# Target: %(setup_target)s

from setuptools import setup

setup(**
%(setup_dict)s
)
"""


class SetupPy(Command):
  """Generate setup.py-based Python projects from python_library targets."""

  __command__ = 'setup_py'

  def setup_parser(self, parser, args):
    parser.set_usage("\n"
                     "  %prog setup_py (options) [spec]\n")
    parser.add_option("--run", dest="run", default=None,
                      help="The command to run against setup.py.  Don't forget to quote "
                           "any additional parameters.  If no run command is specified, "
                           "pants will by default generate and dump the source distribution.")

  def __init__(self, root_dir, parser, argv):
    Command.__init__(self, root_dir, parser, argv)

    if not self.args:
      self.error("A spec argument is required")

    address = Address.parse(root_dir, self.args[0])
    self.target = Target.get(address)

    if self.target is None:
      self.error('%s is not a valid target!' % self.args[0])

    if not self.target.provides:
      self.error('Target must provide an artifact.')

    self.dependencies = self.minified_dependencies(self.target)

  def minified_dependencies(self, root):
    depmap = defaultdict(OrderedSet)
    providers = []

    def resolve(trg):
      if isinstance(trg, PythonTarget) and trg.provides:
        providers.append(trg.provides.key)
      for dependency in getattr(trg, 'dependencies', []):
        for prv in providers:
          for dep in dependency.resolve():
            depmap[prv].add(dep)
            resolve(dep)
      if isinstance(trg, PythonTarget) and trg.provides:
        assert providers[-1] == trg.provides.key
        providers.pop()

    resolve(root)

    root_deps = depmap.pop(root.provides.key, {})

    def elide(trg):
      if isinstance(trg, PythonTarget) and trg.provides:
        return
      if any(trg in depset for depset in depmap.values()):
        root_deps.discard(trg)

    root.walk(elide)
    return root_deps

  def write_sources(self, chroot):
    resources = defaultdict(set)
    def write(trg, src):
      chroot.link(os.path.join(trg.target_base, src), os.path.join('src', src))
    def write_resource(trg, resource):
      base, fn = os.path.split(resource)
      resources[base.replace(os.sep, '.')].add(fn)
      chroot.link(os.path.join(trg.target_base, resource), os.path.join('src', resource))
    for source in list(self.target.sources):
      write(self.target, source)
    for resource in list(self.target.resources):
      write_resource(self.target, resource)
    for target in self.dependencies:
      if isinstance(target, PythonLibrary) and not target.provides:
        for source in list(target.sources):
          write(target, source)
        for resource in list(target.resources):
          write_resource(target, resource)
      elif isinstance(target, PythonThriftLibrary):
        print('WARNING: Thrift libraries not yet supported: %s' % target, file=sys.stderr)
    return dict((pkg_name, list(r)) for (pkg_name, r) in resources.items())

  def establish_namespaces(self, chroot, namespace_packages):
    base = os.path.join(chroot.path(), 'src')

    packages = set()

    for root, dirs, files in os.walk(base):
      rel_root = os.path.relpath(root, base)
      if rel_root == '.':
        continue
      package_name = rel_root.replace(os.path.sep, '.')
      packages.add(package_name)
      if '__init__.py' not in files:
        with open(os.path.join(root, '__init__.py'), 'w') as fp:
          fp.write("__import__('pkg_resources').declare_namespace(__name__)"
                   if package_name in namespace_packages else "")

    return packages, namespace_packages

  def write_setup(self, chroot):
    resources = self.write_sources(chroot)
    setup_keywords = self.target.provides._kw
    install_requires = set()

    package_dir = {'': 'src'}
    packages, _ = self.establish_namespaces(chroot, setup_keywords.get('namespace_packages', {}))

    for dep in self.dependencies:
      if isinstance(dep, PythonRequirement):
        install_requires.add(str(dep._requirement))
      elif isinstance(dep, PythonTarget) and dep.provides:
        install_requires.add(dep.provides.key)

    setup_keywords.update(install_requires=list(install_requires))

    if packages:
      setup_keywords.update(
          package_dir=package_dir,
          packages=list(sorted(packages)),
          package_data=resources)

    chroot.write(SETUP_BOILERPLATE % {
      'setup_dict': pprint.pformat(setup_keywords, indent=4),
      'setup_target': repr(self.target)
    }, 'setup.py')

  def execute(self):
    config = Config.load()
    distdir = config.getdefault('pants_distdir')
    setup_dir = os.path.join(distdir, '%s-%s' % (
        self.target.provides._name, self.target.provides._version))
    chroot = Chroot(distdir, name=self.target.provides._name)

    self.write_setup(chroot)
    if os.path.exists(setup_dir):
      import shutil
      shutil.rmtree(setup_dir)
    os.rename(chroot.path(), setup_dir)

    stderr, stdout = [], []
    def read_output(po):
      stdout.append(po.stdout.read())
      stderr.append(po.stderr.read())

    with pushd(setup_dir):
      cmd = '%s setup.py %s' % (sys.executable, self.options.run or 'sdist')
      print('Running "%s" in %s' % (cmd, setup_dir))
      extra_args = {} if self.options.run else dict(stdout=subprocess.PIPE, stderr=subprocess.PIPE)
      po = subprocess.Popen(cmd, shell=True, **extra_args)
      while po.poll() is None:
        read_output(po)
        time.sleep(0.1)
      read_output(po)

    if self.options.run:
      print('Ran %s' % cmd)
      print('Output in %s' % setup_dir)
      return po.returncode
    elif po.returncode != 0:
      print('Failed to run %s!' % cmd)
      for line in ''.join(stdout).splitlines():
        print('stdout: %s' % line)
      for line in ''.join(stderr).splitlines():
        print('stderr: %s' % line)
      return po.returncode

    expected_tgz = '%s-%s.tar.gz' % (self.target.provides._name, self.target.provides._version)
    expected_target = os.path.join(setup_dir, 'dist', expected_tgz)
    dist_tgz = os.path.join(distdir, expected_tgz)
    if not os.path.exists(expected_target):
      print('Could not find expected target %s!' % expected_target)
      sys.exit(1)
    safe_delete(dist_tgz)
    os.rename(expected_target, dist_tgz)
    print('Wrote %s' % dist_tgz)
    safe_rmtree(setup_dir)
