#!/usr/bin/env python
try:
    import requests
except ImportError:
    print(" >> Please install python-requests")
    raise SystemExit(1)
import argparse
import json
import os
import re
import sys
import urlparse
from HTMLParser import HTMLParser

KVM_REGEX = '.*KVM.*x86_64-.*\.qcow2$'
KVM_KUBIC_REGEX = '.*x86_64-.*Stack-hardware-x86_64.*\.qcow2$'
DOCKER_REGEX = '%(docker_image)s.*x86_64-.*\.tar\.xz$'
OPENSTACK_REGEX = '.*OpenStack-Cloud.*x86_64-.*\.qcow2$'
ISO_REGEX = '.*-DVD-x86_64-.*-Media1\.iso$'

# Link to the package version tracker to allow exploring image contents
PACKAGE_VERSION_TRACKER_IMAGE_URL_TPL = 'http://nodes.provo-prod.caasp.suse.net:30526/image/{}'


class ImageFinder(HTMLParser):

    def __init__(self, args):
        HTMLParser.__init__(self)
        self.images = set()

        if args.type == "docker":
            self.regexp = DOCKER_REGEX % {"docker_image": args.image_name}
        elif args.type == "kvm":
            if args.url == "channel://kubic":
                self.regexp = KVM_KUBIC_REGEX
            else:
                self.regexp = KVM_REGEX
        elif args.type == "openstack":
            self.regexp = OPENSTACK_REGEX
        elif args.type == "iso":
            self.regexp = ISO_REGEX
        else:
            print(" >> Unknown Image Type: " + args.type)
            raise SystemExit(1)

    def get_image(self):
        if len(self.images) > 0:
            return self.images.pop().strip(" ")
        else:
            return None

    def handle_data(self, data):
        m = re.search(self.regexp, data)
        if m:
            self.images.add(m.group(0))


def use_local_file(args):
    expected_path = get_expected_path(args)
    actual_path = urlparse.urlparse(args.url).path

    if os.path.isfile(actual_path):
        link_file(actual_path, expected_path)
    else:
        print(" >> File not found: " + args.url)
        raise SystemExit(2)

def use_remote_file(args):
    download_file(args)
    link_file(get_actual_path(args), get_expected_path(args))

def print_package_version_tracker_link(url):
    """Link to the package version tracker to allow exploring image contents"""
    image_fn = url.rsplit('/', 1)[-1]
    url = PACKAGE_VERSION_TRACKER_IMAGE_URL_TPL.format(image_fn)
    print(" >> To explore the image contents follow:")
    print(" >>   {}".format(url))
    print(" >>")
    pass

def use_channel_file(args):
    remote_url = get_channel_url(args)
    remote_canonical_url = get_canonical_url(args, remote_url)
    expected_path = urlparse.urlparse(args.url).netloc
    download_file(args, remote_canonical_url)
    link_file(get_actual_path(args, remote_canonical_url), get_expected_path(args, urlparse.urlparse(args.url).netloc))
    print_package_version_tracker_link(remote_canonical_url)

## Util functions

# Local File Handling functions

def get_filename(url):
    return urlparse.urlparse(url).path.split('/')[-1]

def get_expected_path(args, filename=None):
    filename = filename or get_filename(args.url)

    if args.type == "docker":
        return os.path.abspath(
            "%(path)s/%(type)s-%(docker_image)s-%(filename)s" %
            {'path': args.path, 'filename': filename, 'type': args.type,
             'docker_image': args.image_name}
        )
    else:
        return os.path.abspath(
            "%(path)s/%(type)s-%(filename)s" %
            {'path': args.path, 'filename': filename, 'type': args.type}
        )

def get_actual_path(args, url=None):
    url = url or args.url
    return  os.path.abspath(
        "%(path)s/%(filename)s" %
        {'path': args.path, 'filename': get_filename(url)}
    )

def link_file(actual_path, expected_path):
    print(" >> File on Disk  : " + actual_path)
    if not expected_path == actual_path:
        print(" >> Static Path : " + expected_path)
        if os.path.islink(expected_path):
            os.unlink(expected_path)
        os.symlink(actual_path, expected_path)

# Remote Downloading functions

def get_canonical_url(args, url):
    proxies = {
      'http': args.proxy,
      'https': args.proxy,
    }

    image = requests.head(url, proxies=proxies)

    if image.status_code == 302 or image.status_code == 301:
        return get_canonical_url(args, image.headers.get('Location'))
    elif image.status_code == 200:
        return url
    else:
        raise Exception("Cannot find image's real location: " + url)

def find_sha256(text):
    # Find the regex by looking for a length-64 hex word
    sha_regex = re.compile(r'\b([0-9A-Fa-f]{64})\b')
    sha_match = sha_regex.findall(text)
    if len(sha_match) == 1:
        return sha_match[0]
    return None  # None will indicate that a sha256 match was not found or that too many were found

def get_sha256(url, proxies):
    print(" >> Downloading Sha256 File")
    partition = url.rpartition('/')
    sha256file = requests.head(url + '.sha256', proxies=proxies)

    if sha256file.status_code == 404:
        # try to find this file in /SHA256SUMS within the same directory
        for line in requests.get(partition[0] + "/SHA256SUMS", proxies=proxies).text.split('\n'):
            # find sha256 for this file
            columns = line.split(" ")
            if partition[2] in columns:
                sha256 = get_sha256(line)
                if sha256 is not None:
                    return sha256
        raise Exception("Cannot find remote image SHA256 in %s" % partition[0] + "/SHA256SUMS")
    elif sha256file.status_code == 302 or sha256file.status_code == 301:
        location = sha256file.headers.get('Location')
        return get_sha256(location[:location.rfind('.sha256')], proxies=proxies)
    elif sha256file.status_code == 200:
        # this will work if the SHA256 is stored as a PGP SIGNED MESSAGE in the .sha256 file
        sha256 = find_sha256(requests.get(url + '.sha256', proxies=proxies).text)
        if sha256 is not None:
            return sha256
        raise Exception("Cannot find remote image SHA256 in %s" % url + '.sha256')
    else:
        raise Exception("Cannot find remote image SHA256")

def download_file(args, canonical_url=None):
    url = canonical_url or args.url
    expected_path = get_expected_path(args)
    actual_path = get_actual_path(args, url)

    if not os.path.isfile(actual_path):
        print(" >> Downloading File")
        print(" >> Remote File         : " + url)
        print(" >> Local File (On Disk): " + actual_path)

        proxies = {
          'http': args.proxy,
          'https': args.proxy,
        }

        try:
            remote_sha_pre = get_sha256(url, proxies=proxies)

            proxy_flag =  '--no-proxy' if args.proxy == '' else '-e use_proxy=yes -e http_proxy=' + args.proxy

            os.system(
                "wget %(proxy_flag)s %(url)s -O %(file)s --progress=dot:giga" %
                { "url": url, "file": actual_path, "proxy_flag": proxy_flag}
            )

            local_sha = os.popen('sha256sum %s' % actual_path).read().split(' ')[0]
            remote_sha_post = get_sha256(url, proxies=proxies)

            print(" >> Local SHA:                  %s" % local_sha)
            print(" >> Remote SHA (Pre Download):  %s" % remote_sha_pre)
            print(" >> Remote SHA (Post Download): %s" % remote_sha_post)

            if local_sha not in [remote_sha_post, remote_sha_pre]:
                print(" >> Download corrupted - please retry.")
                raise SystemExit(3)

        except:
            print(" >> Deleting failed download")
            os.remove(actual_path)
            raise

# Remote Parsing functions

def get_channel_url(args):

    channel = urlparse.urlparse(args.url).netloc

    parser = ImageFinder(args)

    if args.download_urls_fname is None:
        # generate location relative to this script (not `cwd`)
        mypath = os.path.realpath(__file__)
        relpath = '../../misc-files/download-urls.json'
        download_urls_fname = os.path.join(mypath, relpath)
    else:
        download_urls_fname = args.download_urls_fname

    download_urls_fname = os.path.abspath(download_urls_fname)

    try:
        with open(download_urls_fname) as f:
            j = json.load(f)
        base = j["baseurl"]
    except Exception as e:
        raise Exception("Unable to extract baseurl: {}".format(e))

    channel = base.get(channel, False)
    if not channel:
        raise Exception("Unknown channel: %s" % channel)

    base_url = channel.get(args.location, channel.get("default", False))
    if not base_url:
        raise Exception("Malformed channel entry")

    if args.type == "docker":
        base_url += 'images_container_base'
    elif args.type in ["kvm", "openstack"]:
        base_url += 'images'
    elif args.type == "iso":
        base_url += 'images/iso'
    else:
        print(" >> Unknown Image Type: " + args.type)
        raise SystemExit(1)

    proxies = {
      'http': args.proxy,
      'https': args.proxy,
    }

    r = requests.get(base_url, proxies=proxies)
    parser.feed(r.text)

    image = parser.get_image()

    if image is None:
        print(" >> No matching images found")

    return "%(base)s/%(image)s" % {
        "base": base_url,
        "image": image
    }



if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='Download CaaSP Image')
    parser.add_argument('--type', choices=["docker", "kvm", "openstack", "iso"], help="Type of image to download", required=True)
    parser.add_argument('--path', help="Where should the image be downloaded and linked (default: '../downloads')", default='../downloads')
    parser.add_argument('--image-name', help='Name of the Docker derived image to download (eg: "sles12-velum-development").')
    parser.add_argument('--proxy', help="Proxy server to use", default='')
    parser.add_argument('--location', help="Download Location")
    parser.add_argument('--download-urls-fname',
            help="Path of the download URLs JSON file (default: <script_location>/../misc-files/download-urls.json)",
            default=None)
    parser.add_argument('url', metavar='url', help='URL of image to download')
    parser.add_argument('--expand-channel', action='store_true', help='Show URL of channel repository')
    args = parser.parse_args()

    if urlparse.urlparse(args.url).scheme in ['http', 'https']:
        use_remote_file(args)

    elif urlparse.urlparse(args.url).scheme == "file":
        use_local_file(args)

    elif urlparse.urlparse(args.url).scheme == "channel":
        if args.type in ["docker", "kvm", "openstack", "iso"]:
            if args.expand_channel != False:
                sys.stdout.write("%(repo)s%(ext)s" % {'repo': get_channel_url(args).split("images")[0], 'ext': 'standard'})
                sys.exit(0)
            else:
                use_channel_file(args)
        else:
            print(" >> Unknown Image Type: " + args.type)
            raise SystemExit(1)
    else:
        print(" >> Unknown URL Type: " + args.url)
        raise SystemExit(2)
