#!/usr/bin/env python3
# Copyright (C) 2023 Leandro Lisboa Penz <lpenz@lpenz.org>
# This file is subject to the terms and conditions defined in
# file 'LICENSE', which is part of this source code package.

import argparse
import logging
import os
import shutil

from diskimgtool import chdir, chroot, cmd, image_fully_mounted, root_mounts


def log():
    if not hasattr(log, "logger"):
        log.logger = logging.getLogger(os.path.basename("diskimgtool"))
    return log.logger


def cmd_list(args):
    with image_fully_mounted(args.image) as rootdir:
        with chdir(rootdir):
            cmd(["find", ".", "-type", "f", "-ls"])


def cmd_enter(args):
    with image_fully_mounted(args.image) as rootdir:
        with root_mounts(rootdir):
            chroot(rootdir, ["/bin/bash", "-i"])


def cmd_get(args):
    srcrel = args.source.lstrip("/")
    with image_fully_mounted(args.image) as rootdir:
        srcfull = os.path.join(rootdir, srcrel)
        print(f"{srcfull}, {args.destination.name}")
        with open(srcfull, "rb") as fd:
            log().info(f"Getting /{srcrel} from image as {args.destination.name}")
            shutil.copyfileobj(fd, args.destination)


def cmd_put(args):
    dstrel = args.destination.lstrip("/")
    with image_fully_mounted(args.image) as rootdir:
        dstfull = os.path.join(rootdir, dstrel)
        if os.path.isdir(dstfull):
            dstrel = os.path.join(dstrel, os.path.basename(args.source.name))
            dstfull = os.path.join(rootdir, dstrel)
        with open(dstfull, "wb") as fd:
            log().info(f"Putting {args.source.name} in image as /{dstrel}")
            shutil.copyfileobj(args.source, fd)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("-v", "--verbose", action="store_true", help="Verbose logging")
    parser.add_argument("image", help="Image to change")
    subparser = parser.add_subparsers(dest="command", required=True)
    # List:
    parser_list = subparser.add_parser("list", help="List all files in image")
    parser_list.set_defaults(func=cmd_list)
    # Enter:
    parser_enter = subparser.add_parser("enter", help="Start a shell inside the image")
    parser_enter.set_defaults(func=cmd_enter)
    # Get:
    parser_get = subparser.add_parser("get", help="Get a file from the image")
    parser_get.add_argument(
        "source",
        help="Source file in image",
    )
    parser_get.add_argument(
        "destination",
        type=argparse.FileType("wb"),
        help="Path to the destination to write",
    )
    parser_get.set_defaults(func=cmd_get)
    # Put:
    parser_put = subparser.add_parser("put", help="Put a file in the image")
    parser_put.add_argument(
        "source",
        type=argparse.FileType("rb"),
        help="Source file",
    )
    parser_put.add_argument(
        "destination",
        help="Path of the destination, relative to the image",
    )
    parser_put.set_defaults(func=cmd_put)
    # Configure and call subcommand function:
    args = parser.parse_args()
    if args.verbose:
        level = logging.INFO
    else:
        level = logging.WARN
    logging.basicConfig(format="%(name)s: %(message)s", level=level)
    args.func(args)


if __name__ == "__main__":
    main()
