#!/usr/bin/python3

# Copyright 2023 Acronis
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# This file is a barebones file needed to file a gap until Ansible 2.0. No
# error checking, no deletions, no updates. Idempotent creation only.

# If you look closely, you will see we arent _really_ using the shade module
# we just use it to slightly abstract the authentication model. As patches land
# in upstream shade we will be able to use more of the shade module. Until then
# if we want to be 'stable' we really need to be using it as a passthrough

import argparse
import atexit
import datetime
import json
import logging
import os
import shutil
import socket
import subprocess
import sys
import time
import uuid

from glanceclient.client import Client as GlanceClient
from keystoneauth1.identity import v3
from keystoneauth1 import session

LOG = logging.getLogger()

INVENTORY_CMD = ['/etc/kolla/vstorage-backend.py', '--list']
DEFAULT_DOMAIN = "default"
IMAGE_IMPORT_DIR = '/mnt/vstorage/vols/datastores/glance'

HOSTNAME = socket.gethostname()


def parse_arguments():
    parser = argparse.ArgumentParser(description="Image importing tool.")
    parser.add_argument("--project", help="Project ID.")
    parser.add_argument("--min-disk", type=int, help="Minimum disk size required to boot from image, in gigabytes.")
    parser.add_argument("--min-ram", type=int, help="Minimum RAM size required to boot from image, in megabytes.")
    parser.add_argument("--os-type", help="OS type.")
    parser.add_argument("--os-distro", help="OS distribution.")
    parser.add_argument("--protected", action="store_true", help="Protect image from deletion.")
    parser.add_argument("--public", action="store_true", help="Make image accessible to all users.")
    parser.add_argument("--tags", help="A comma-separated list of tags.")
    parser.add_argument("--verify", action="store_true", help="Verify the checksum of the uploaded image.")
    parser.add_argument("--uefi", action="store_true", help="Create image with UEFI.")
    parser.add_argument("--debug", action="store_true", help="Debug logging verbosity.")
    parser.add_argument("filepath", help=f"Create image from a local file. File will be copied to {IMAGE_IMPORT_DIR} for import.")
    parser.add_argument("new_image_name", help="Image name.")

    return parser.parse_args()


if __name__ == "__main__":
    # Run the command and capture both stdout and stderr
    result = subprocess.run(
        INVENTORY_CMD, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

    if result.returncode != 0:
        LOG.error(result.stderr)
        sys.exit(1)

    host_inventory = {
        params['ansible_host']: params for params in
        json.loads(result.stdout)['_meta']['hostvars'].values()}
    # sys.stdout.write(json.dumps(host_inventory[HOSTNAME], indent=4))

    params = host_inventory[HOSTNAME]

    auth = v3.Password(auth_url=f"https://{params['internal_vip_address']}:35357",
                       username=params["keystone_admin_user"],
                       password=params["keystone_admin_password"],
                       project_name=params["keystone_admin_project"],
                       user_domain_id=DEFAULT_DOMAIN,
                       project_domain_id=DEFAULT_DOMAIN)
    session = session.Session(auth=auth, verify=False)

    client = GlanceClient('2', interface='admin', session=session, endpoint_override="http://127.0.0.1:9294")

    parsed_args = parse_arguments()
    logging.basicConfig(
        level=logging.DEBUG if parsed_args.debug else logging.INFO,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        handlers=[
            logging.StreamHandler(),
        ]
    )

    image_params = {}

    for arg in ['min_disk', 'min_ram','os_type', 'os_distro', 'protected']:
        if getattr(parsed_args, arg, None):
            image_params[arg] = getattr(parsed_args, arg)
    if parsed_args.project:
        image_params['owner'] = parsed_args.project
    if parsed_args.tags:
        image_params['tags'] = parsed_args.tags.split(',')
    if parsed_args.uefi:
        image_params['hw_firmware_type'] = 'uefi'
    image_params['visibility'] = 'public' if parsed_args.public else 'shared'
    image_params['name'] = parsed_args.new_image_name
    filepath = os.path.abspath(parsed_args.filepath)

    import_filepath = filepath
    file_is_in_import_dir = os.path.commonprefix([IMAGE_IMPORT_DIR, import_filepath]) == IMAGE_IMPORT_DIR
    if not file_is_in_import_dir:
        import_filepath = os.path.join(IMAGE_IMPORT_DIR, f"{uuid.uuid4()}.import")
        if os.path.exists(import_filepath):
            LOG.error("unexpectedly, %s file already exists", import_filepath)
            sys.exit(1)
        atexit.register(lambda: os.path.exists(import_filepath) and os.remove(import_filepath))
        shutil.copyfile(filepath, import_filepath)

    req = {
        'import_from': f"file://{import_filepath}",
        'import_from_format': '',
        'image_properties': {
            'hw_qemu_guest_agent': 'yes',
            'disk_format': 'detect',
            'container_format': 'bare',
        }
    }
    req['image_properties'].update(image_params)

    task = client.tasks.create(type='import', input=req)

    count = 0
    while task.get('status') in ('processing', 'pending'):
        time.sleep(1)
        count += 1
        ct = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        print(f"{ct} glance task {task['id']} {task['status']}        ", end="\r")
        sys.stdout.flush()
        if count % 3 == 0:
            task = client.tasks.get(task.id)
    print(f"{ct} glance task {task['id']} {task['status']} {task['message']}        ")
