#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# vim: set ts=4
#
# Copyright 2024-present Linaro Limited
#
# SPDX-License-Identifier: MIT


import argparse
import json
import logging
import os
import shutil
import sys
import subprocess as sp
import tarfile
from datetime import datetime

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def run_proc(*args):
    cmd = list(args)
    logger.debug(f"Running {cmd}")
    proc = sp.Popen(cmd, stdout=sp.PIPE, stderr=sp.PIPE, env=os.environ)
    stdout, stderr = proc.communicate()
    proc.ok = proc.returncode == 0

    proc.out = stdout.decode()
    proc.err = stderr.decode()
    return proc


def tuxsuite_download_artefacts(build_id, directory):
    cmd = [
        "tuxsuite",
        "build",
        "get",
        build_id,
        "--download-artifacts",
        directory,
    ]
    ret = run_proc(*cmd)
    return ret


def move_files(fromdir, todir):
    allfiles = os.listdir(fromdir)

    for f in allfiles:
        src_path = os.path.join(fromdir, f)
        dst_path = os.path.join(todir, f)
        shutil.move(src_path, dst_path)


def extract_tar_xz(tar_file):
    if not os.path.isfile(tar_file):
        logger.debug(f"Error: File {tar_file} does not exist.")
        return

    extract_dir = os.path.dirname(tar_file)

    try:
        # Open and extract the tar.xz file
        with tarfile.open(tar_file, mode='r:xz') as tar:
            tar.extractall(path=extract_dir)
            logger.debug(f"Extracted {tar_file} to {extract_dir}")
    except Exception as e:
        logger.error(f"Failed to extract {tar_file}: {e}")


def extract_build_info(json_file, base_dir):
    with open(json_file, 'r') as file:
        data = json.load(file)

    now = datetime.now()
    for build_key, build_info in data['builds'].items():
        target_arch = build_info.get('target_arch', None)
        kernel_version = build_info.get('kernel_version', None)
        build_name = build_info.get('build_name', None)
        logger.debug(f"{build_key} {target_arch} {kernel_version}")
        date = now.strftime("%Y%m%d_%H%M%S")
        directory_path = os.path.join(base_dir, date, kernel_version, build_name or target_arch)
        os.makedirs(directory_path, exist_ok=True)
        print(f"Created: {directory_path}")
        tuxsuite_download_artefacts(build_key, directory_path)
        fromdir = os.path.join(directory_path, build_key)
        move_files(fromdir, directory_path)
        shutil.rmtree(fromdir)
        extract_tar_xz(os.path.join(directory_path, "dtbs.tar.xz"))


def parse_args():
    parser = argparse.ArgumentParser(description="Extract build keys, target architecture, and kernel version from a JSON file.")
    parser.add_argument(
        "--json-file",
        required=True,
        help="Path to the JSON file containing build information.",
    )

    parser.add_argument(
        "--base-dir",
        default="./",
        help="Base directory where the directories will be created.",
    )

    parser.add_argument(
        "--debug",
        action="store_true",
        default=False,
        help="Display debug messages",
    )

    return parser.parse_args()


def run():

    args = parse_args()

    extract_build_info(args.json_file, args.base_dir)


if __name__ == "__main__":
    sys.exit(run())