#!/usr/bin/env python
# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""This script fetches and prepares an SDK chroot.
"""

import os
import sys
import urlparse

from chromite.buildbot import constants
from chromite.lib import cgroups
from chromite.lib import commandline
from chromite.lib import cros_build_lib
from chromite.lib import locking
from chromite.lib import osutils
from chromite.lib import toolchain

cros_build_lib.STRICT_SUDO = True


COMPRESSION_PREFERENCE = ('xz', 'bz2')

# TODO(zbehan): Remove the dependency on these, reimplement them in python
MAKE_CHROOT = [os.path.join(constants.SOURCE_ROOT,
                            'src/scripts/sdk_lib/make_chroot.sh')]
ENTER_CHROOT = [os.path.join(constants.SOURCE_ROOT,
                             'src/scripts/sdk_lib/enter_chroot.sh')]

# We need these tools to run. Very common tools (tar,..) are ommited.
NEEDED_TOOLS = ('curl', 'xz', 'unshare')


def GetArchStageTarballs(version):
  """Returns the URL for a given arch/version"""
  suburl = '%s/coreos-sdk-amd64-%s.tar.bz2' % (version, version)
  return [toolchain.GetSdkURL(suburl=suburl)]


def GetStage3Urls(version):
  return GetArchStageTarballs(version)


def FetchRemoteTarballs(storage_dir, urls):
  """Fetches a tarball given by url, and place it in sdk/.

  Args:
    urls: List of URLs to try to download. Download will stop on first success.

  Returns:
    Full path to the downloaded file
  """

  # Note we track content length ourselves since certain versions of curl
  # fail if asked to resume a complete file.
  # pylint: disable=C0301,W0631
  # https://sourceforge.net/tracker/?func=detail&atid=100976&aid=3482927&group_id=976
  for url in urls:
    # http://www.logilab.org/ticket/8766
    # pylint: disable=E1101
    parsed = urlparse.urlparse(url)
    tarball_name = os.path.basename(parsed.path)
    if parsed.scheme in ('', 'file'):
      if os.path.exists(parsed.path):
        return parsed.path
      continue
    content_length = 0
    print 'Attempting download: %s' % url
    result = cros_build_lib.RunCurl(
          ['-I', url], redirect_stdout=True, redirect_stderr=True,
          print_cmd=False)
    successful = False
    for header in result.output.splitlines():
      # We must walk the output to find the 200 code for use cases where
      # a proxy is involved and may have pushed down the actual header.
      if (header.startswith("HTTP/1.0 200") or
          header.startswith("HTTP/1.1 200") or
          header.startswith("HTTP/2.0 200") or
          header.startswith("HTTP/2 200")):
        successful = True
      elif header.lower().startswith("content-length:"):
        content_length = int(header.split(":", 1)[-1].strip())
        if successful:
          break
    if successful:
      break
  else:
    raise Exception('No valid URLs found!')

  tarball_dest = os.path.join(storage_dir, tarball_name)
  current_size = 0
  if os.path.exists(tarball_dest):
    current_size = os.path.getsize(tarball_dest)
    if current_size > content_length:
      osutils.SafeUnlink(tarball_dest)
      current_size = 0

  if current_size < content_length:
    cros_build_lib.RunCurl(
        ['-f', '-L', '-y', '30', '-C', '-', '--output', tarball_dest, url],
        print_cmd=False)

  # Cleanup old tarballs now since we've successfull fetched; only cleanup
  # the tarballs for our prefix, or unknown ones.
  ignored_prefix = ('stage3-' if tarball_name.startswith('cros-sdk-')
                    else 'cros-sdk-')
  for filename in os.listdir(storage_dir):
    if filename == tarball_name or filename.startswith(ignored_prefix):
      continue

    print 'Cleaning up old tarball: %s' % (filename,)
    osutils.SafeUnlink(os.path.join(storage_dir, filename))

  return tarball_dest


def CreateChroot(chroot_path, sdk_tarball, cache_dir,
                 nousepkg=False, nogetbinpkg=False):
  """Creates a new chroot from a given SDK"""

  cmd = MAKE_CHROOT + ['--stage3_path', sdk_tarball,
                       '--chroot', chroot_path,
                       '--cache_dir', cache_dir]
  if nousepkg:
    cmd.append('--nousepkg')
  elif nogetbinpkg:
    cmd.append('--nogetbinpkg')

  try:
    cros_build_lib.RunCommand(cmd, print_cmd=False)
  except cros_build_lib.RunCommandError:
    raise SystemExit('Running %r failed!' % cmd)


def DeleteChroot(chroot_path):
  """Deletes an existing chroot"""
  cmd = MAKE_CHROOT + ['--chroot', chroot_path,
                       '--delete']
  try:
    cros_build_lib.RunCommand(cmd, print_cmd=False)
  except cros_build_lib.RunCommandError:
    raise SystemExit('Running %r failed!' % cmd)


def EnterChroot(chroot_path, cache_dir, chrome_root, chrome_root_mount,
                additional_args):
  """Enters an existing SDK chroot"""
  cmd = ENTER_CHROOT + ['--chroot', chroot_path, '--cache_dir', cache_dir]
  if chrome_root:
    cmd.extend(['--chrome_root', chrome_root])
  if chrome_root_mount:
    cmd.extend(['--chrome_root_mount', chrome_root_mount])
  if len(additional_args) > 0:
    cmd.append('--')
    cmd.extend(additional_args)