#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import argparse
import socket
from pathlib import Path
from queue import Queue

from p1st.exec_consume_chunks import execute_consume_chunks
from p1st.unix_sock_input import accept_loop_until_message


def main():
    args = parse_args()

    command = ['btrfs', 'receive', str(args.target_path.parent)]

    def handle_chunks(queue_put: Queue.put):
        chunk_dir: Path = args.chunk_dir
        chunk_dir.mkdir(parents=True, exist_ok=True)

        socket_file: Path = args.target_path.parent.joinpath(f'{args.target_path.name}.SOCKET')
        socket_file.parent.mkdir(parents=True, exist_ok=True)

        print(f'Listening on socket {socket_file}')
        sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        sock.bind(str(socket_file))
        sock.listen(1)

        chunk_no = 1
        messages = [b'OK\n', b'EOF\n']
        while True:
            msg = accept_loop_until_message(sock, messages)

            last_chunk = msg == b'EOF\n'
            chunk_path = chunk_dir.joinpath(f'{args.target_path.name}.CHUNK.{chunk_no}')
            queue_put((chunk_path, last_chunk))

            if last_chunk:
                break
            chunk_no += 1

        print(f'Closing socket {socket_file}')
        sock.close()
        socket_file.unlink(missing_ok=False)

    returncode = execute_consume_chunks(
        command=command,
        handle_chunks=handle_chunks,
    )
    exit(returncode)


def parse_args():
    parser = argparse.ArgumentParser(prog='btrfs-receive-chunks')

    parser.add_argument('--chunk-dir',
                        help='Chunks are saved in this directory. '
                             'Defaults to parent directory of SUBVOLUME.',
                        dest='chunk_dir',
                        type=Path,
                        metavar='CHUNK_DIR',
                        default=None,
                        )

    parser.add_argument('target_path',
                        help='Path where the subvolume will be created. '
                             'The last component of the path '
                             'must be equal to the name of the subvolume on the sending side.',
                        type=Path,
                        metavar='SUBVOLUME'
                        )

    args = parser.parse_args()

    # Make all paths absolute. Set default values.
    args.target_path = args.target_path.absolute()
    if args.chunk_dir:
        args.chunk_dir = args.chunk_dir.absolute()
    else:
        args.chunk_dir = args.target_path.absolute().parent

    return args


if __name__ == '__main__':
    main()