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

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 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)

    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):
        self.index[image_id] = { 'filename' : filename,
                                 'checksum' : checksum,
                                 'hashalgo' : hashalgo }
        if url:
            self.index[image_id]['url'] = url
        self.index_dump()
        self.checksum_dump(hashalgo)

    def find_filename(self, image_id):
        entry = self.index.get(image_id)
        if entry is None:
            return None
        filename = entry.get('filename')
        if filename is None:
            return None
        return f'{self.imgdir}/{filename}'


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):
        dl = DistroList(imgdir)
        dl.update(image_id, filename,
                  checksum, hashalgo,
                  url = url)

    @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

    @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
        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 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 """

        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 """

        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 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)
