""" base class for image download """
import os
import re
import sys
import glob
import json
import time
import hashlib
import datetime
import configparser

from collections import UserDict

import requests

def default_imgdir():
    """ return default distro image directory """
    imgdir = os.environ.get('VIRT_DISTRO_IMGDIR')
    if imgdir is not None:
        return imgdir
    home = os.environ.get('HOME')
    if home is not None:
        return f'{home}/.cache/virt-distro'
    return 'distro'

class DistroEntry(UserDict):
    """ one downloaded distro image """

    def __init__(self, data, image_id, imgdir):
        data['id']      = image_id
        data['imgdir']  = imgdir
        (name, version, arch, kind) = image_id.rsplit('-', 3)
        data['name']    = name
        data['version'] = version
        data['arch']    = arch
        data['kind']    = kind
        super().__init__(data)

    @property
    def name(self):
        return self.get('name')

    @property
    def version(self):
        return self.get('version')

    @property
    def arch(self):
        return self.get('arch')

    @property
    def kind(self):
        return self.get('kind')

    @property
    def imgdir(self):
        return self.get('imgdir')

    @property
    def filename(self):
        return self.get('filename')

    @property
    def filepath(self):
        return f'{self.imgdir}/{self.filename}'

    def repoconfig(self):
        repos = self.get('rpmrepos')
        if not repos:
            return None
        config = ''
        for (name, url) in repos.items():
            reponame = f'{self.name}-{self.version}-{name}'
            config += f'[{reponame}]\n'
            config += f'name={reponame}\n'
            config += f'baseurl={url}\n'
            config += 'gpgcheck=false\n'
            config += 'skip_if_unavailable=true\n'
            config += '\n'
        return config

    def kickstart(self):
        repos = self.get('rpmrepos')
        if not repos:
            return None
        config = ''
        first = True
        for (name, url) in repos.items():
            if first:
                config += f'url --url={url}\n'
                first = False
            else:
                config += f'repo --name={name} --baseurl={url}\n'
        return config

class DistroList:
    """ list of downloaded distro images """

    def __init__(self, imgdir):
        self.imgdir = imgdir
        self.filename = f'{self.imgdir}/index.json'
        self.index = {}
        self.index_load()

    def index_load(self):
        if not os.path.exists(self.filename):
            return
        with open(self.filename, 'r', encoding = 'utf-8') as f:
            self.index = json.load(f)

    def index_dump(self):
        with open(self.filename, 'w', encoding = 'utf-8') as f:
            json.dump(self.index, f, indent = 4)

    def checksum_dump(self, hashalgo = 'sha256'):
        outfile = f'{self.imgdir}/checksums.{hashalgo}'
        with open(outfile, 'w', encoding = 'utf-8') as f:
            for name in self.index.keys():
                if self.index.get(name).get('hashalgo') != hashalgo:
                    continue
                filename = self.index.get(name).get('filename')
                checksum = self.index.get(name).get('checksum')
                print(f'{checksum} {filename}', file = f)

    @staticmethod
    def mkid(name_ver, kind, arch = None):
        if arch is None:
            arch = os.uname().machine
        return f'{name_ver}-{arch}-{kind}'

    # pylint: disable=too-many-arguments,too-many-positional-arguments
    def update(self, image_id, filename,
               checksum, hashalgo = 'sha256',
               url = None, rpmrepos = None):
        self.index[image_id] = { 'filename' : filename,
                                 'checksum' : checksum,
                                 'hashalgo' : hashalgo }
        if url:
            self.index[image_id]['url'] = url
        if rpmrepos:
            self.index[image_id]['rpmrepos'] = rpmrepos
        self.index_dump()
        self.checksum_dump(hashalgo)

    def remove_stale(self):
        print('# --- check for stale images ---')
        files = list(map(lambda e: e.filename, self.list_entries()))
        for f in glob.glob(f'{self.imgdir}/*'):
            name = f.removeprefix(f'{self.imgdir}/')
            if not os.path.isfile(f):
                continue
            if name == 'index.json':
                continue
            if name.startswith('checksum'):
                continue
            if name in files:
                print(f'# latest: {name}')
            else:
                print(f'# stale : {name}')
                os.remove(f)

    def list_entries(self):
        for image_id, entry in self.index.items():
            yield DistroEntry(entry, image_id, self.imgdir)

    def find_entry(self, image_id):
        entry = self.index.get(image_id)
        if entry is None:
            return None
        return DistroEntry(entry, image_id, self.imgdir)

    def find_filename(self, image_id):
        entry = self.find_entry(image_id)
        if entry is None:
            return None
        return entry.filepath


class FetchDistroCore:
    """ base class for image download """

    arch = os.uname().machine
    retries = 5
    timeout = 20

    def __init__(self, arch = None):
        if arch is not None:
            self.arch = arch

    @staticmethod
    def index_load(imgdir):
        dl = DistroList(imgdir)
        return dl.index

    # pylint: disable=too-many-arguments,too-many-positional-arguments
    @staticmethod
    def index_update(imgdir, image_id, filename,
                     checksum, hashalgo = 'sha256',
                     url = None, rpmrepos = None):
        dl = DistroList(imgdir)
        dl.update(image_id, filename,
                  checksum, hashalgo,
                  url = url,
                  rpmrepos = rpmrepos)

    @staticmethod
    def index_find(imgdir, image_id):
        dl = DistroList(imgdir)
        return dl.find_filename(image_id)

    @staticmethod
    def request(url, stream, start = None, end = None):
        hdr = {}
        if start and end:
            hdr['range'] = f'bytes={start}-{end-1}'
        r = FetchDistroCore.retries
        for i in range(r):
            try:
                req = requests.get(url, headers = hdr,
                                   stream = stream,
                                   timeout = FetchDistroCore.timeout)
            except requests.exceptions.RequestException as e:
                print(f'# request connect error [{i}/{r}]: {e}')
                print(f'# retry in {FetchDistroCore.timeout} seconds')
                time.sleep(FetchDistroCore.timeout)
                continue
            req.raise_for_status()
            if i > 0:
                print(f'# ... connected [{i}/{r}]')
            return req
        return None

    # pylint: disable=too-many-statements
    @staticmethod
    def fetch_image_checksum(url, out, checksum):
        """ download one file, verify checksum """

        print(f'# latest: {out}')
        print(f'# sha256: {checksum}')
        if os.path.exists(out):
            print('# skip download, file exists')
            return
        print('# downloading file ...')

        req = FetchDistroCore.request(url, True)
        if req is None:
            raise RuntimeError('connect failed')
        count = 0
        countnext = 0
        length = int(req.headers.get("content-length", 0))
        lengthm = int(length / (1024*1024))
        hasher = hashlib.sha256()

        with open(out, 'wb') as fh:
            i = req.iter_content(chunk_size = 1024 * 64)
            while True:
                try:
                    chunk = next(i)
                except StopIteration:
                    chunk = None
                except requests.exceptions.RequestException as e:
                    print('')
                    print(f'# request read error: {e}')
                    print('# try reconnect ...')
                    req = FetchDistroCore.request(url, True, count, length)
                    if req is None:
                        os.remove(out)
                        raise RuntimeError('re-connect failed') from e
                    i = req.iter_content(chunk_size = 1024 * 64)
                    continue

                if chunk is None:
                    break
                fh.write(chunk)
                hasher.update(chunk)
                count += len(chunk)
                countm = int(count / (1024*1024))
                if not sys.stdout.isatty():
                    if length == 0 or count > countnext:
                        print(f'# {countm:4}M / {lengthm:4}M', flush = True)
                        countnext = count + 1024 * 1024 * 10
                elif length == 0:
                    print(f'# {countm:4}M', end='\r')
                else:
                    percent = int(count * 100 / length)
                    print(f'# {countm:4}M / {lengthm:4}M - {percent:2}%', end='\r')

        print()
        if hasher.hexdigest() != checksum:
            print(f'# sha256: {hasher.hexdigest()} (mismatch)')
            os.remove(out)
            raise RuntimeError('sha256 mismatch')

        print('# sha256: check ok')

    @staticmethod
    def checksum_find_latest(url, pattern, checksum = 'SHA256'):
        """ find file (and sha256) by downloading CHECKSUM file """

        print(f'# fetch {url} ...')
        req = FetchDistroCore.request(url, False)
        if req is None:
            raise RuntimeError('request failed')
        matches = []
        for line in req.text.splitlines():
            if not line.startswith(checksum):
                continue
            if not pattern in line:
                continue
            if 'latest' in line:
                continue
            matches.append(line)
        if len(matches) == 0:
            return (None, None)
        latest = sorted(matches, reverse = True)[0]
        regex = checksum + r' \((.+)\) = ([a-fA-F0-9]{64})'
        m = re.match(regex, latest)
        filename = m.group(1)
        checksum = m.group(2)
        return (filename, checksum)

    @staticmethod
    def indexhtml_find_file(url, pattern):
        """ find file by searching index.html """

        print(f'# fetch {url} ...')
        req = FetchDistroCore.request(url, False)
        if req is None:
            raise RuntimeError('request failed')
        regex = rf'href="([^"]*{re.escape(pattern)}[^"]*)"'
        found = re.search(regex, req.text)
        if not found:
            return None
        # some mirrors return absolute not relative links
        elems = found.group(1).split('/')
        return elems[-1]

    @staticmethod
    def find_mirror(url):
        """ find location by following mirror network redirects """
        retries = 5
        while retries > 0:
            print(f'# {url}')
            r = requests.head(url, allow_redirects = False,
                              timeout = FetchDistroCore.timeout)
            if r.status_code == 200:
                print('# -> OK')
                return url
            if r.status_code in (301, 302):
                print(f'# -> Redirect [{r.status_code}]')
                url = r.headers['location']
            else:
                print(f'# -> Oops [{r.status_code}]')
                return None
            retries -= 1
        print('# too many redirects')
        return None

class TreeInfo:
    """ handle .treeinfo files """

    def __init__(self, url):
        self.info = configparser.ConfigParser()
        req = FetchDistroCore.request(url, False)
        if req is None:
            raise RuntimeError('request failed')
        self.info.read_string(req.text)

    def sha256(self, image):
        line = self.info.get('checksums', image)
        if not line:
            return None
        (algo, checksum) = line.split(':')
        if algo == 'sha256':
            return checksum
        return None

    def release_date(self):
        ts = self.info.get('general', 'timestamp')
        date = datetime.date.fromtimestamp(int(ts))
        return date.isoformat()

    def distro_name(self):
        name = self.info.get('release', 'short').lower()
        version = self.info.get('release', 'version').lower()
        return f'{name}-{version}'

class FetchNetinst(FetchDistroCore):
    """ class for fedora / centos / rhel netinst iso download """

    def __init__(self, mirror, arch = None):
        super().__init__(arch)
        self.mirror = mirror
        # filled by treeinfo()
        self.info = None
        self.netinst = None
        self.outfile = None
        self.image_id = None
        self.checksum = None

    def baseurl(self, _release):
        return None

    def rpmrepos(self, _booturl):
        return None

    def treeinfo(self, release):
        # pylint: disable=assignment-from-none
        base = self.baseurl(release)
        boot = 'images/boot.iso'
        self.info = TreeInfo(f'{base}.treeinfo')
        self.netinst = f'{base}{boot}'
        self.checksum = self.info.sha256(boot)
        reldate = self.info.release_date()
        distname = self.info.distro_name()
        self.outfile = f'{distname}-{self.arch}-netinst-{reldate}.iso'
        self.image_id = f'{distname}-{self.arch}-netinst'

    def fetch_image(self, release, imgdir = 'distro'):
        self.treeinfo(release)
        print(f'# ----- {self.outfile} -----')
        url = self.find_mirror(self.netinst)
        out = f'{imgdir}/{self.outfile}'
        os.makedirs(imgdir, exist_ok = True)
        self.fetch_image_checksum(url, out, self.checksum)
        self.index_update(imgdir, self.image_id, self.outfile, self.checksum,
                          url = url, rpmrepos = self.rpmrepos(url))
