/*
Copying and distribution of this file, with or without modification,
are permitted in any medium without royalty.  This file is offered
as-is, without any warranty.
*/

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.math.BigInteger;
import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.MappedByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.channels.FileChannel.MapMode;
import java.util.ArrayList;
import java.util.Arrays;


public class ZfsReader {
    public static void main(String[] argv) {
        if (argv.length != 1) {
            System.err.println("Usage <app> <vdev file>");
            System.exit(1);
            return;
        }

        FileInputStream in;
        try {
            in = new FileInputStream(argv[0]);
        } catch (FileNotFoundException e) {
            System.err.println(e.getLocalizedMessage());
            System.exit(1);
            return;
        }

        FileChannel channel = in.getChannel();
        try {
            try {
                for (int i = 0; i < 4; i++) {
                    VdevLabel vdevLabel = new VdevLabel(channel, i);

                    Uberblock mostRecent = null;
                    for (Uberblock uberblock : vdevLabel.getUberblocks()) {
                        if (uberblock != null &&
                                (mostRecent == null || uberblock.getTransactionGroup().compareTo(mostRecent.getTransactionGroup()) > 0)) {
                            mostRecent = uberblock;
                        }
                    }

                    System.out.printf("vdev base: 0x%x\n", vdevLabel.getBase());
                    if (mostRecent != null) {
                        System.out.printf("  uberblock offset: 0x%x, txn group = %d, sum = %d\n", mostRecent.getOffset(), mostRecent.getTransactionGroup(), mostRecent.getGuidSum());
                    } else {
                        System.out.printf("  no valid uberblocks found\n");
                    }
                }

            } finally {
                channel.close();
            }
        } catch (IOException e) {

        }
    }

    public static class BlockPointer extends ZfsStruct {

        //dva_t           blk_dva[3];     /* 128-bit Data Virtual Address */
        //uint64_t        blk_prop;       /* size, compression, type, etc */
        //uint64_t        blk_pad[3];     /* Extra space for the future   */
        //uint64_t        blk_birth;      /* transaction group at birth   */
        //uint64_t        blk_fill;       /* fill count                   */
        //zio_cksum_t     blk_cksum;      /* 256-bit checksum             */

        private final BigInteger[] dataVirtualAddress = new BigInteger[3];
        private final BigInteger blockProperties;
        private final long[] pad = new long[3];
        private final BigInteger blockBirth;
        private final BigInteger blockFillCount;
        private final byte[] blockChecksum = new byte[32];

        public BlockPointer(ByteBuffer buffer) {
            dataVirtualAddress[0] = getUnsignedBigInteger(buffer, 16);
            dataVirtualAddress[1] = getUnsignedBigInteger(buffer, 16);
            dataVirtualAddress[2] = getUnsignedBigInteger(buffer, 16);

            blockProperties = getUnsignedBigInteger(buffer, 8);

            pad[0] = buffer.getLong();
            pad[1] = buffer.getLong();
            pad[2] = buffer.getLong();

            blockBirth = getUnsignedBigInteger(buffer, 8);
            blockFillCount = getUnsignedBigInteger(buffer, 8);

            buffer.get(blockChecksum);
        }

        public BigInteger[] getDataVirtualAddress() {
            return dataVirtualAddress;
        }

        public BigInteger getBlockProperties() {
            return blockProperties;
        }

        public long[] getPad() {
            return pad;
        }

        public BigInteger getBlockBirth() {
            return blockBirth;
        }

        public BigInteger getBlockFillCount() {
            return blockFillCount;
        }

        public byte[] getBlockChecksum() {
            return blockChecksum;
        }
    }


    public static class InvalidBlockException extends Exception {

        private static final long serialVersionUID = 1286497411777506847L;

        public InvalidBlockException() {

        }

        public InvalidBlockException(String message) {
            super(message);
        }
    }

    public static class Uberblock extends ZfsStruct {
        private static final byte[] BIG_ENDIAN_MAGIC = new byte[] {
            (byte)0x00, (byte)0x00, (byte)0x00, (byte)0x00,
            (byte)0x00, (byte)0xba, (byte)0xb1, (byte)0x0c
        };

        private static final byte[] LITTLE_ENDIAN_MAGIC = new byte[] {
            (byte)0x0c, (byte)0xb1, (byte)0xba, (byte)0x00,
            (byte)0x00, (byte)0x00, (byte)0x00, (byte)0x00
        };

        private static final int UBERBLOCK_SIZE = 1024;

        private final int offset;
        private long uberblockVersion;
        private BigInteger transactionGroup;
        private BigInteger guidSum;
        private long timestamp;
        private BlockPointer blockPointer;

        public Uberblock(ByteBuffer buffer) throws InvalidBlockException {
            offset = buffer.position();

            byte[] magic = new byte[8];
            buffer.get(magic);

            ByteOrder byteOrder = null;

            if (Arrays.equals(BIG_ENDIAN_MAGIC, magic)) {
                byteOrder = ByteOrder.BIG_ENDIAN;
                buffer.order(byteOrder);
            } else if (Arrays.equals(LITTLE_ENDIAN_MAGIC, magic)) {
                byteOrder = ByteOrder.LITTLE_ENDIAN;
                buffer.order(byteOrder);
            }

            uberblockVersion = buffer.getLong();
            transactionGroup = getUnsignedBigInteger(buffer, 8);
            guidSum = getUnsignedBigInteger(buffer, 8);
            timestamp = buffer.getLong();
            blockPointer = new BlockPointer(buffer);

            buffer.position(offset + UBERBLOCK_SIZE);

            if (byteOrder == null) {
                throw new InvalidBlockException("No ZFS Magic");
            }
        }

        public long getUberblockVersion() {
            return uberblockVersion;
        }

        public BigInteger getTransactionGroup() {
            return transactionGroup;
        }

        public BigInteger getGuidSum() {
            return guidSum;
        }

        public long getTimestamp() {
            return timestamp;
        }

        public BlockPointer getBlockPointer() {
            return blockPointer;
        }

        public int getOffset() {
            return offset;
        }
    }

    public static class VdevLabel {
        private static final int VDEV_LABEL_SIZE = 256 * 1024;
        private static final int FIRST_UBERBLOCK = 128 * 1024;
        private static final int LAST_UBERBLOCK = 256 * 1024;

        private final Uberblock[] uberblocks;

        private final long base;

        // 0x00 - 0x1FFF: blank
        // 0x2000 - 0x3FFF: Boot Header (reserved)
        // 0x4000 - 0x1FFFF: Name/Value pairs
        // 0x20000 - 0x3ffff: Uberblock array


        public VdevLabel(FileChannel channel, int labelIndex) throws IOException {
            if (labelIndex >= 4 || labelIndex < 0) {
                throw new IllegalArgumentException();
            }

            if (channel.size() < VDEV_LABEL_SIZE * 4) {
                throw new IOException("Vdev not big enough");
            }

            MappedByteBuffer buffer;

            if (labelIndex < 2) {
                base = VDEV_LABEL_SIZE * labelIndex;
            } else {
                base = channel.size() - VDEV_LABEL_SIZE * (4 - labelIndex);
            }

            buffer = channel.map(MapMode.READ_ONLY, base, VDEV_LABEL_SIZE);

            ArrayList<Uberblock> uberblockList = new ArrayList<Uberblock>(128);

            for (buffer.position(FIRST_UBERBLOCK); buffer.position() < LAST_UBERBLOCK;) {
                try {
                    uberblockList.add(new Uberblock(buffer));
                } catch (InvalidBlockException e) {
                    uberblockList.add(null);
                } catch (BufferUnderflowException e) {
                    buffer.position(LAST_UBERBLOCK);
                }
            }

            uberblocks = uberblockList.toArray(new Uberblock[0]);
        }

        public Uberblock[] getUberblocks() {
            return uberblocks;
        }

        public long getBase() {
            return base;
        }
    }

    public static abstract class ZfsStruct {
        protected BigInteger getUnsignedBigInteger(ByteBuffer buffer, int bytes) {
            if (bytes < 1) {
                throw new IllegalArgumentException();
            }

            // Pad the array with a leading 0x00 to ensure the result is positive
            byte[] temp = new byte[bytes + 1];

            if (buffer.order() == ByteOrder.BIG_ENDIAN) {
                buffer.get(temp, 1, bytes);
            } else {
                // Read the bytes, then flip the order
                buffer.get(temp, 0, bytes);
                for (int i = 0, j = temp.length - 1; i < temp.length / 2; i++, j--) {
                    byte temp2 = temp[i];
                    temp[i] = temp[j];
                    temp[j] = temp2;
                }
            }

            return new BigInteger(temp);
        }
    }
}