""" ubuntu """
import os
import json

from virt.qemutest.distro.core import FetchDistroCore

class FetchUbuntuCloud(FetchDistroCore):
    """ class for ubuntu cloud image download """

    releases = "https://cloud-images.ubuntu.com/releases"
    arch_map = {
        'x86_64': 'amd64',
        'aarch64': 'arm64',
    }

    def __init__(self, arch = None):
        super().__init__(arch)
        self.ubuntu_arch = self.arch_map.get(self.arch, self.arch)

    def image_id(self, release):
        return f'ubuntu-{release}-{self.arch}-cloud'

    def fetch_json(self, path):
        url = f'{self.releases}/{path}'
        print(f'# fetch {url} ...')
        req = FetchDistroCore.request(url, False)
        if req is None:
            raise RuntimeError(f'http request for {path} failed')
        return json.loads(req.text)

    def find_stream_download(self):
        stream = self.fetch_json('streams/v1/index.json')
        for name, desc in stream.get('index').items():
            if desc.get('datatype') != 'image-downloads':
                continue
            return desc.get('path')
        return None

    def find_stream_diskimage(self, download, release):
        stream = self.fetch_json(download)
        for name, desc in stream.get('products').items():
            if desc.get("arch") != self.ubuntu_arch:
                continue
            names = desc.get("aliases", desc.get("version")).split(',')
            if not release in names:
                continue
            latest = sorted(desc.get('versions').keys(), reverse = True)[0]
            disk1 = desc.get('versions').get(latest).get('items').get('disk1.img')
            image = {
                'os'      : desc.get('os'),
                'release' : desc.get('release'),
                'version' : desc.get('version'),
                'date'    : latest,
                'path'    : disk1.get('path'),
                'sha256'  : disk1.get('sha256'),
            }
            return image
        return None

    def list_images(self):
        download = self.find_stream_download()
        stream = self.fetch_json(download)
        for name, desc in stream.get('products').items():
            if desc.get("arch") != self.ubuntu_arch:
                continue
            print(f'# {desc.get("os")} {desc.get("version")} ({desc.get("release")})')

    def find_image(self, release):
        download = self.find_stream_download()
        return self.find_stream_diskimage(download, release)

    def fetch_image(self, release = 'lts', imgdir = 'distro'):
        image = self.find_image(release)
        if image is None:
            raise RuntimeError('image not found')
        filename = "-".join([
            image.get("os"),
            image.get("release"),
            image.get("version"),
            image.get("date"),
            self.arch,
        ])
        print(f'# ----- {filename} -----')
        url = f'{self.releases}/{image.get("path")}'
        out = f'{imgdir}/{filename}.qcow2'
        os.makedirs(imgdir, exist_ok = True)
        self.fetch_image_checksum(url, out, image.get("sha256"))
        self.index_update(imgdir, self.image_id(image.get("version")),
                          f'{filename}.qcow2', image.get("sha256"),
                          url = url)

if __name__ == "__main__":
    f = FetchUbuntuCloud()
    f.list_images()
    f.fetch_image()
