#
# upip - Package manager for MicroPython
#
# Copyright (c) 2015-2018 Paul Sokolovsky
#
# Licensed under the MIT license.
#
import sys
import os
import errno
import json
import zlib
import time
import shlex
import subprocess
import shutil
# import upip_host_tarfile as tarfile
# gc.collect()


debug = False
index_urls = ["https://micropython.org/pi", "https://pypi.org/pypi"]
install_path = None
cleanup_files = []
gzdict_sz = 16 + 15

file_buf = bytearray(512)


class NotFoundError(Exception):
    pass


def op_split(path):
    if path == "":
        return ("", "")
    r = path.rsplit("/", 1)
    if len(r) == 1:
        return ("", path)
    head = r[0]
    if not head:
        head = "/"
    return (head, r[1])


def op_basename(path):
    return op_split(path)[1]

# Expects *file* name


def _makedirs(name, mode=0o777):
    ret = False
    s = ""
    comps = name.rstrip("/").split("/")[:-1]
    if comps[0] == "":
        s = "/"
    for c in comps:
        if s and s[-1] != "/":
            s += "/"
        s += c
        try:
            os.mkdir(s)
            ret = True
        except OSError as e:
            if e.args[0] != errno.EEXIST and e.args[0] != errno.EISDIR:
                raise e
            ret = False
    return ret


def save_file(fname, subf):
    global file_buf
    with open(fname, "wb") as outf:
        while True:
            sz = subf.readinto(file_buf)
            if not sz:
                break
            outf.write(file_buf, sz)


def install_tar(f, prefix):
    meta = {}
    for info in f:
        # print(info)
        fname = info.name
        try:
            fname = fname[fname.index("/") + 1:]
        except ValueError:
            fname = ""

        save = True
        for p in ("setup.", "PKG-INFO", "README"):
                #print(fname, p)
                if fname.startswith(p) or ".egg-info" in fname:
                    if fname.endswith("/requires.txt"):
                        meta["deps"] = f.extractfile(info).read()
                    save = False
                    if debug:
                        print("Skipping", fname)
                    break

        if save:
            outfname = prefix + fname
            # if info.type != tarfile.DIRTYPE:
            #     if debug:
            #         print("Extracting " + outfname)
            #     _makedirs(outfname)
            #     subf = f.extractfile(info)
            #     save_file(outfname, subf)
    return meta


def expandhome(s):
    if "~/" in s:
        h = os.getenv("HOME")
        s = s.replace("~/", h + "/")
    return s


import ssl
import socket
warn_ssl = True


def url_open(url):
    global warn_ssl

    if debug:
        print(url)

    proto, _, host, urlpath = url.split('/', 3)
    try:
        ai = socket.getaddrinfo(host, 443, 0, socket.SOCK_STREAM)
    except OSError as e:
        fatal("Unable to resolve %s (no Internet?)" % host, e)
    # print("Address infos:", ai)
    ai = ai[0]

    s = socket.socket(ai[0], ai[1], ai[2])
    try:
        # print("Connect address:", addr)
        s.connect(ai[-1])

        if proto == "https:":
            context = ssl.create_default_context()

            s = context.wrap_socket(s, server_hostname=host)
            if warn_ssl:
                print("Warning: %s SSL certificate is not validated" % host)
                warn_ssl = False

        # MicroPython rawsocket module supports file interface directly
        s.write("GET /{} HTTP/1.0\r\nHost: {}\r\n\r\n".format(urlpath, host).encode('utf-8'))
        l = s.recv()
        protover, status, msg = l.split(None, 2)
        if status != b"200":
            if status == b"404" or status == b"301":
                raise NotFoundError("Package not found")
            raise ValueError(status)
        server_metadata = l.split(b'\n')[:-1]
        json_data = l.split(b'\n')[-1]
        while 1:
            l = s.recv()
            if not l:
                # raise ValueError("Unexpected EOF in HTTP headers")
                time.sleep(0.2)
                break
            if l == b'\r\n':
                break
    except Exception as e:
        s.close()
        raise e

    # print(server_metadata)
    return json_data


def get_pkg_metadata(name):
    for url in index_urls:
        try:
            f = url_open("%s/%s/json" % (url, name))
        except NotFoundError:
            continue
        try:
            return json.loads(f)
        finally:
            pass
    raise NotFoundError("Package not found")


def fatal(msg, exc=None):
    print("Error:", msg)
    if exc and debug:
        raise exc
    sys.exit(1)


def install_pkg(pkg_spec, install_path, read_pkg_info=False):
    kw = ["setup.", "PKG-INFO", "README"]
    print('Collecting: {}'.format(pkg_spec))
    data = get_pkg_metadata(pkg_spec)
    # print(data)
    latest_ver = data["info"]["version"]
    packages = data["releases"][latest_ver]
    del data
    # gc.collect()
    assert len(packages) == 1
    package_url = packages[0]["url"]
    if read_pkg_info:
        print("Reading PKG-INFO %s %s from %s" % (pkg_spec, latest_ver, package_url))
    else:
        print("Installing %s %s from %s" % (pkg_spec, latest_ver, package_url))
    curl_cmd_str = "curl -O '{}'".format(package_url)
    curl_cmd = shlex.split(curl_cmd_str)
    try:
        proc = subprocess.call(curl_cmd)
        # print('Done!')
    except KeyboardInterrupt:
        print('Operation cancelled')
    try:
        tarfile = [f for f in os.listdir() if pkg_spec.replace('micropython-', '') in f and f.endswith('.tar.gz')][0]
        shutil.unpack_archive(tarfile)
        os.remove(tarfile)
    except Exception as e:
        print(e)
    try:
        dir_pkg = [f for f in os.listdir() if pkg_spec.replace('micropython-', '') in f and not f.endswith('.py')][0]
        pkg_egg_info_dir = [f for f in os.listdir(dir_pkg) if f.endswith('.egg-info')][0]
        with open(dir_pkg+"/"+pkg_egg_info_dir+'/PKG-INFO', 'r') as pkginfo:
            pkg_info = pkginfo.read()
            # print(pkg_info)
    except Exception as e:
        print(e)
    if read_pkg_info:
        print(pkg_info)
        shutil.rmtree(dir_pkg)
    else:
        pkg_req = None
        try:
            print('Package requirements: ')
            with open(dir_pkg+"/"+pkg_egg_info_dir+'/requires.txt', 'r') as pkgreq:
                pkg_req = pkgreq.read().split('\n')
                for pk in pkg_req:
                    if pk != '':
                        print(pk)
        except Exception as e:
            print('None')

        if pkg_req:
            for pk in pkg_req:
                if pk != '':
                    req_content = install_pkg(pk, ".")

        pkg_content = [f for f in os.listdir(dir_pkg) if not any([f.startswith(p) for p in kw]) and ".egg-info" not in f]
        # print(pkg_content)
        dir_lib = 'lib'
        if dir_lib not in os.listdir():
            os.mkdir(dir_lib)
        for pkg_file in pkg_content:
            if os.stat(dir_pkg+'/'+pkg_file)[0] & 0x4000:
                if pkg_file not in os.listdir(dir_lib):
                    os.mkdir('/'.join([dir_lib, pkg_file]))
                rcopytree(dir_pkg+'/'+pkg_file, '/'.join([dir_lib, pkg_file]))
        for pkg_file in pkg_content:
            if not os.stat(dir_pkg+'/'+pkg_file)[0] & 0x4000:
                shutil.copyfile(dir_pkg+'/'+pkg_file, dir_lib+'/'+pkg_file)

        shutil.rmtree(dir_pkg)
        print('Successfully installed {}'.format(dir_pkg))
        return pkg_content, dir_pkg


def rcopytree(src, dst):
    for pkg_file in os.listdir(src):
        if os.stat(src+'/'+pkg_file)[0] & 0x4000:
            if pkg_file in os.listdir(dst):
                pass
            else:
                # print('Dest dir: ', dst+'/'+pkg_file)
                os.mkdir(dst+'/'+pkg_file)
            rcopytree(src+'/'+pkg_file, '/'.join([dst, pkg_file]))

    for pkg_file in os.listdir(src):
        if not os.stat(src+'/'+pkg_file)[0] & 0x4000:
            # print('Source file: ', src+'/'+pkg_file)
            shutil.copyfile(src+'/'+pkg_file, dst+'/'+pkg_file)


def install(to_install, install_path=None):
    # Calculate gzip dictionary size to use
    global gzdict_sz
    sz = 65538
    if sz <= 65536:
        gzdict_sz = 16 + 12

    if install_path is None:
        install_path = get_install_path()
    if install_path[-1] != "/":
        install_path += "/"
    if not isinstance(to_install, list):
        to_install = [to_install]
    print("Installing to: " + install_path)
    # sets would be perfect here, but don't depend on them
    installed = []
    try:
        while to_install:
            if debug:
                print("Queue:", to_install)
            pkg_spec = to_install.pop(0)
            if pkg_spec in installed:
                continue
            meta = install_pkg(pkg_spec, install_path)
            installed.append(pkg_spec)
            if debug:
                print(meta)
            deps = meta.get("deps", "").rstrip()
            if deps:
                deps = deps.decode("utf-8").split("\n")
                to_install.extend(deps)
    except Exception as e:
        print("Error installing '{}': {}, packages may be partially installed".format(
                pkg_spec, e),
              file=sys.stderr)


def get_install_path():
    global install_path
    if install_path is None:
        # sys.path[0] is current module's path
        install_path = sys.path[1]
    install_path = expandhome(install_path)
    return install_path


def cleanup():
    for fname in cleanup_files:
        try:
            os.unlink(fname)
        except OSError:
            print("Warning: Cannot delete " + fname)


def help():
    print("""\
upip - Simple PyPI package manager for MicroPython
Usage: micropython -m upip install [-p <path>] <package>... | -r <requirements.txt>
import upip; upip.install(package_or_list, [<path>])

If <path> is not given, packages will be installed into sys.path[1]
(can be set from MICROPYPATH environment variable, if current system
supports that).""")
    print("Current value of sys.path[1]:", sys.path[1])
    print("""\

Note: only MicroPython packages (usually, named micropython-*) are supported
for installation, upip does not support arbitrary code in setup.py.
""")


def main():
    global debug
    global index_urls
    global install_path
    install_path = None

    if len(sys.argv) < 2 or sys.argv[1] == "-h" or sys.argv[1] == "--help":
        help()
        return

    if sys.argv[1] != "install":
        fatal("Only 'install' command supported")

    to_install = []

    i = 2
    while i < len(sys.argv) and sys.argv[i][0] == "-":
        opt = sys.argv[i]
        i += 1
        if opt == "-h" or opt == "--help":
            help()
            return
        elif opt == "-p":
            install_path = sys.argv[i]
            i += 1
        elif opt == "-r":
            list_file = sys.argv[i]
            i += 1
            with open(list_file) as f:
                while True:
                    l = f.readline()
                    if not l:
                        break
                    if l[0] == "#":
                        continue
                    to_install.append(l.rstrip())
        elif opt == "-i":
            index_urls = [sys.argv[i]]
            i += 1
        elif opt == "--debug":
            debug = True
        else:
            fatal("Unknown/unsupported option: " + opt)

    to_install.extend(sys.argv[i:])
    if not to_install:
        help()
        return

    install(to_install)

    if not debug:
        cleanup()


if __name__ == "__main__":
    main()
