/*
 *  Copyright © 2008-2015 Amichai Rothman
 *
 *  This file is part of JNRG - the Java NRG Reader utility.
 *
 *  JNRG is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  JNRG is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with JNRG.  If not, see <http://www.gnu.org/licenses/>.
 *
 *  For additional info see http://www.freeutils.net/source/jnrg/
 */

package net.freeutils.nrg;

import java.io.*;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.util.*;

/**
 * The {@code NRGReader} class reads the NRG disc image file format, and
 * allows the extraction of parts of it to external files.
 * <p>
 * Note: The NRG file format is not well documented at this time (to the best
 * of my knowledge), and so this implementation is based on partial data,
 * a few sample NRG files, and guesswork.
 * It may be wrong, and is certainly incomplete.
 *
 * @author Amichai Rothman
 * @since  2008-11-05
 */
public class NRGReader {

    /**
     * The {@code IFFFile} class represents an Interchange File Format file.
     * It extends RandomAccessFile with additional convenience read methods.
     */
    static class IFFFile extends RandomAccessFile {

        protected byte[] buf = new byte[4];

        public IFFFile(String name, String mode) throws FileNotFoundException {
            super(name, mode);
        }

        public String readChunkID() throws IOException {
            readFully(buf, 0, 4);
            return new String(buf, 0, 4, "ISO8859_1");
        }

        public long readUnsignedInt() throws IOException {
            return readInt() & 0xFFFFFFFFL;
        }

        public void skipFully(int n) throws IOException {
            for (; n > 0; n -=4)
                readFully(buf, 0, n > 4 ? 4 : n);
        }
    }

    /**
     * The {@code ChunkID} annotation decorates Chunk subclasses with
     * the ID(s) which they can handle.
     */
    @Retention(RetentionPolicy.RUNTIME)
    protected static @interface ChunkID {
        String[] value();
    }

    /**
     * The {@code Header} class represents an NRG file header.
     * <p>
     * Unfortunately, the header (which is actually a footer, as it appears
     * at the very end of the file) is not a valid IFF chunk, although
     * it is quite similar.
     */
    @ChunkID({"NERO", "NER5"})
    public static class Header {
        String ID;
        long firstChunkOffset;

        /**
         * Reads the NRG header data.
         *
         * @param file the file to read the header from
         * @throws IOException if an error occurs
         */
        protected void readHeader(IFFFile file) throws IOException {
            long len = file.length();
            if (len < 12)
                throw new IOException("invalid NRG file: too small");
            file.seek(len - 12);
            ID = file.readChunkID();
            if (ID.equals("NER5")) {
                firstChunkOffset = file.readLong();
            } else {
                ID = file.readChunkID();
                if (!ID.equals("NERO"))
                    throw new IOException("invalid NRG file: header not found");
                firstChunkOffset = file.readUnsignedInt();
            }
        }
    }

    /**
     * The {@code Chunk} class represents a single IFF chunk.
     */
    public static class Chunk {
        String ID;   // always exactly 4 characters
        long offset; // offset of data within file
        long size;   // size of data (in bytes)

        /**
         * Reads the chunk's header data starting at the current file position.
         *
         * @param file the file to read the chunk from
         * @throws IOException if an error occurs
         */
        public void readHeader(IFFFile file) throws IOException {
            ID = file.readChunkID();
            size = file.readUnsignedInt();
            offset = file.getFilePointer();
        }

        /**
         * Reads the chunk's data starting at the current file position.
         * <p>
         * Subclasses should override this method with chunk-specific parsing.
         *
         * @param file the file to read the chunk from
         * @throws IOException if an error occurs
         */
        public void readData(IFFFile file) throws IOException {}

        @Override
        public String toString() {
            return String.format("chunk %s (offset %d, size %d)",
                ID, offset, size);
        }
    }

    @ChunkID({"CUES", "CUEX"})
    public static class CueSheetChunk extends Chunk {

        public static final int
            TYPE_AUDIO = 0x01,
            TYPE_AUDIO2 = 0x21,
            TYPE_DATA_MODE_1 = 0x41;

        public static class Index {
            int type;
            int trackNumber;
            int indexNumber;
            int unknown;
            int block; // signed - can be negative for gap before first track
            int blockCount;

            public String getTimeString() {
                int frames = blockCount % 75;
                int secs = blockCount / 75;
                int mins = secs / 60;
                secs = secs % 60;
                return String.format("%02d:%02d:%02d", mins, secs, frames);
            }

            public String getTypeDesc() {
                switch (type) {
                    case TYPE_AUDIO:
                    case TYPE_AUDIO2:
                        return "AUDIO";
                    case TYPE_DATA_MODE_1:
                        return "DATA";
                    default:
                        return "UNKNOWN";
                }
            }
        }

        Index[] indices;

        public void readData(IFFFile file) throws IOException {
            List<Index> indices = new ArrayList<Index>();
            long count = size / 8;
            for (int i = 0; i < count; i++) {
                Index ind = new Index();
                ind.type = file.readUnsignedByte();
                ind.trackNumber = file.readUnsignedByte();
                ind.indexNumber = file.readUnsignedByte();
                ind.unknown = file.readUnsignedByte();
                ind.block = file.readInt();
                indices.add(ind);
            }
            this.indices = indices.toArray(new Index[indices.size()]);
            for (int i = 0; i < this.indices.length - 1; i++)
                this.indices[i].blockCount =
                    this.indices[i + 1].block - this.indices[i].block;
        }
    }

    @ChunkID({"DAOI", "DAOX"})
    public static class DAOInfoChunk extends Chunk {

        public static class TrackInfo {
            int sectorSize;
            int mode;
            long offsetIndex0;
            long offsetIndex1;
            long offsetTrackEnd; // exclusive (one after the last byte)
        }

        long size2;
        byte[] unknown;
        int tocType;
        int firstTrack;
        int lastTrack;
        TrackInfo[] tracks;

        public void readData(IFFFile file) throws IOException {
            List<TrackInfo> tracks = new ArrayList<TrackInfo>();
            boolean isVersion1 = ID.equals("DAOI");
            file.readUnsignedInt();
            unknown = new byte[14];
            file.readFully(unknown);
            tocType = file.readUnsignedShort();
            firstTrack = file.readUnsignedByte();
            lastTrack = file.readUnsignedByte();
            long count = (size - 22) / (isVersion1 ? 30 : 42);
            for (long i = 0; i < count; i++) {
                TrackInfo t = new TrackInfo();
                file.skipFully(10);
                t.sectorSize = file.readInt();
                t.mode = file.readInt();
                if (isVersion1) {
                    t.offsetIndex0 = file.readUnsignedInt();
                    t.offsetIndex1 = file.readUnsignedInt();
                    t.offsetTrackEnd = file.readUnsignedInt();
                } else {
                    t.offsetIndex0 = file.readLong();
                    t.offsetIndex1 = file.readLong();
                    t.offsetTrackEnd = file.readLong();
                }
                tracks.add(t);
            }
            this.tracks = tracks.toArray(new TrackInfo[tracks.size()]);
        }
    }

    @ChunkID("CDTX")
    public static class CDTextChunk extends Chunk {

        public static class Pack {
            byte[] data;
        }

        Pack[] packs;

        public void readData(IFFFile file) throws IOException {
            List<Pack> packs = new ArrayList<Pack>();
            long count = size / 18;
            for (int i = 0; i < count; i++) {
                Pack p = new Pack();
                p.data = new byte[18];
                file.readFully(p.data);
                packs.add(p);
            }
            this.packs = packs.toArray(new Pack[packs.size()]);
        }
    }

    @ChunkID({"ETNF", "ETN2"})
    public static class ExtendedTrackInfoChunk extends Chunk {

        public static class TrackInfo {
            long offset;
            long size;
            int mode;
            int block;
            int unknown;
        }

        TrackInfo[] tracks;

        public void readData(IFFFile file) throws IOException {
            List<TrackInfo> tracks = new ArrayList<TrackInfo>();
            boolean isVersion1 = ID.equals("ETNF");
            long count = size / (isVersion1 ? 20 : 28);
            for (long i = 0; i < count; i++) {
                TrackInfo t = new TrackInfo();
                if (isVersion1) {
                    t.offset = file.readUnsignedInt();
                    t.size = file.readUnsignedInt();
                } else {
                    t.offset = file.readLong();
                    t.size = file.readLong();
                }
                t.mode = file.readInt();
                t.block = file.readInt();
                t.unknown = file.readInt();
                tracks.add(t);
            }
            this.tracks = tracks.toArray(new TrackInfo[tracks.size()]);
        }
    }

    @ChunkID("SINF")
    public static class SessionInfoChunk extends Chunk {
        int tracksInSession;

        public void readData(IFFFile file) throws IOException {
            tracksInSession = file.readInt();
        }
    }

    @ChunkID("MTYP")
    public static class MediaTypeChunk extends Chunk {
        public static final int CD = 0x01, DVD = 0x1C;
        int mediaType;

        public void readData(IFFFile file) throws IOException {
            mediaType = file.readInt();
        }
    }

    @ChunkID("RELO")
    public static class RELOChunk extends Chunk {
        int unknown;

        public void readData(IFFFile file) throws IOException {
            unknown = file.readInt();
        }
    }

    @ChunkID("DINF")
    public static class DiscInfoChunk extends Chunk {
        public static final int FINALIZED = 0, NOT_FINALIZED = 1;
        int finalization;

        public void readData(IFFFile file) throws IOException {
            finalization = file.readInt();
        }
    }

    @ChunkID("TOCT")
    public static class TOCTChunk extends Chunk {
        int unknown;

        public void readData(IFFFile file) throws IOException {
            unknown = file.readUnsignedShort();
        }
    }

    @ChunkID("END!")
    public static class EndChunk extends Chunk {} // just a marker - no data

    protected static final Map<String, Class<? extends Chunk>> chunkTypes;
    static {
        // fill lookup map of Chunk implementations by chunk IDs
        chunkTypes = new HashMap<String, Class<? extends Chunk>>();
        for (Class<?> cls : NRGReader.class.getDeclaredClasses()) {
            if (Chunk.class.isAssignableFrom(cls)) {
                @SuppressWarnings("unchecked")
                Class<? extends Chunk> chunkCls = (Class<? extends Chunk>)cls;
                ChunkID IDs = chunkCls.getAnnotation(ChunkID.class);
                if (IDs != null)
                    for (String ID: IDs.value())
                        chunkTypes.put(ID, chunkCls);
            }
        }
    }

    protected String filename;
    protected IFFFile file;
    protected Header header;

    /**
     * Opens the given NRG file.
     *
     * @param filename the path to the NRG file
     * @throws IOException if an error occurs, or the given file is not
     *         a valid NRG file
     */
    public void open(String filename) throws IOException {
        if (header != null)
            throw new IOException("previous file has not been closed");
        this.filename = filename;
        file = new IFFFile(filename, "r");
        header = new Header();
        header.readHeader(file);
    }

    /**
     * Closes the NRG file (if it was previously open).
     *
     * @throws IOException if an error occurs
     */
    public void close() throws IOException {
        if (file != null)
            file.close();
        filename = null;
        file = null;
        header = null;
    }

    /**
     * Reads a chunk's header data starting at the current file position.
     *
     * @return the read chunk
     * @throws IOException if an error occurs
     */
    protected Chunk readCurrentChunkHeader() throws IOException {
        String ID = file.readChunkID();
        file.seek(file.getFilePointer() - 4);
        Chunk c;
        Class<? extends Chunk> cls = chunkTypes.get(ID);
        if (cls == null)
            cls = Chunk.class; // generic chunk
        try {
            c = cls.newInstance();
        } catch (Exception e) {
            throw new IOException("can't instantiate chunk: " + e);
        }
        c.readHeader(file);
        return c;
    }

    /**
     * Reads all chunks whose ID is among the given IDs, or all available
     * chunks if no IDs are specified.
     *
     * @param ids the IDs of the chunks to return. If empty, all available
     *        chunks are returned.
     * @return the requested chunks (in their order of appearance)
     * @throws IOException if an error occurs
     */
    public List<Chunk> readChunks(String... ids) throws IOException {
        if (header == null)
            throw new IOException("no open NRG file");
        List<String> IDList = Arrays.asList(ids);
        List<Chunk> chunks = new ArrayList<Chunk>();
        file.seek(header.firstChunkOffset);
        Chunk c;
        do {
            c = readCurrentChunkHeader();
            if (ids.length == 0 || IDList.contains(c.ID)) {
                c.readData(file);
                chunks.add(c);
            }
            file.seek(c.offset + c.size); // seek next chunk
        } while (!c.ID.equals("END!"));
        return chunks;
    }

    /**
     * Checks whether there is a valid ISO image starting at the given offset.
     *
     * @param offset the offset where the ISO image should start
     * @return true if there is a valid ISO image at the given offset
     * @throws IOException if an error occurs
     */
    protected boolean isISOImage(long offset) throws IOException {
        file.seek(offset + 16 * 2048); // first 16 blocks are unused
        byte[] buf = new byte[6]; // read first few bytes of image header
        file.readFully(buf);
        String isoID = new String(buf, 1, 5, "ISO8859_1");
        return isoID.equals("CD001");
    }

    /**
     * Extracts the specified part of the NRG file to the given file.
     *
     * @param outfilename the output filename
     * @param offset the offset within the NRG file to start copying from
     * @param len the number of bytes to copy
     * @throws IOException if an error occurs
     */
    protected void extract(String outfilename, long offset, long len) throws IOException {
        RandomAccessFile outfile = new RandomAccessFile(outfilename, "rw");
        while (len > 0) {
            long transferred = file.getChannel().transferTo(
                offset, len, outfile.getChannel());
            offset += transferred;
            len -= transferred;
        }
        outfile.close();
    }

    /**
     * Extracts the specified track from the NRG file to the given file.
     *
     * @param outfilename the output filename
     * @param track the track number to extract (1-based)
     * @throws IOException if an error occurs
     */
    public void extract(String outfilename, short track) throws IOException {
        List<Chunk> chunks = readChunks("DAOI", "DAOX", "ETNF", "ETN2");
        Chunk c = chunks.isEmpty() || track < 1 ? null : chunks.get(0);
        long offset = 0;
        long len = -1;
        if (c instanceof ExtendedTrackInfoChunk) {
            ExtendedTrackInfoChunk cc = (ExtendedTrackInfoChunk)c;
            if (cc.tracks.length >= track) {
                offset = cc.tracks[track - 1].offset;
                len = cc.tracks[track - 1].size;
            }
        } else if (c instanceof DAOInfoChunk) {
            DAOInfoChunk cc = (DAOInfoChunk)c;
            if (cc.tracks.length >= track) {
                offset = cc.tracks[track - 1].offsetIndex1;
                len = cc.tracks[track - 1].offsetTrackEnd - offset;
            }
        }
        if (len < 0)
            throw new IOException("can't find track #" + track);
        extract(outfilename, offset, len);
    }

    /**
     * Extracts a CUE file from the NRG file to the given file.
     *
     * @param imgfilename the image filename corresponding to the CUE file
     * @param cuefilename the output CUE filename
     * @throws IOException if an error occurs
     */
    public void extractCue(String imgfilename, String cuefilename) throws IOException {
        List<Chunk> chunks = readChunks("CUES", "CUEX");
        if (chunks.isEmpty())
            throw new IOException("no CUE info found");
        CueSheetChunk c = (CueSheetChunk)chunks.get(0);
        PrintWriter out = new PrintWriter(cuefilename, "ISO8859_1");
        out.printf("FILE \"%s\" BINARY\r\n", imgfilename);
        int track = -1;
        for (CueSheetChunk.Index ind : c.indices) {
            if (ind.trackNumber == 0 || ind.trackNumber == 0xAA)
                continue; // ignore leading and final index placeholders
            if (track != ind.trackNumber) {
                track = ind.trackNumber;
                out.printf("  TRACK %02d %s\r\n", track, ind.getTypeDesc());
            }
            out.printf("    INDEX %02d %s\r\n", ind.indexNumber, ind.getTimeString());
        }
        out.close();
    }

    /**
     * Prints information about the disc image in the NRG file.
     *
     * @throws IOException if an error occurs
     */
    public void printInfo() throws IOException {
        PrintStream out = System.out;
        boolean version1 = header.ID.equals("NERO");
        out.printf("File:\t\t %s%n", filename);
        out.printf("File size:\t %d bytes%n", file.length());
        out.printf("NRG version:\t %d%n", version1 ? 1 : 2);
        int sessions = 1;
        for (Chunk c : readChunks()) {
            if (c instanceof MediaTypeChunk) {
                MediaTypeChunk cc = (MediaTypeChunk)c;
                out.printf("Media type:\t %s%n",
                    cc.mediaType == MediaTypeChunk.CD ? "CD" : "DVD");
            } else if (c instanceof ExtendedTrackInfoChunk) {
                ExtendedTrackInfoChunk cc = (ExtendedTrackInfoChunk)c;
                out.printf("Recording type:\t Track-At-Once (TAO)%n");
                int i = 1;
                for (ExtendedTrackInfoChunk.TrackInfo t : cc.tracks) {
                    out.printf("  Track #%d:\t mode %d, block %d, %n" +
                        "\t\t offset %d, size %d%n",
                        i++, t.mode, t.block, t.offset, t.size);
                    if (isISOImage(t.offset))
                        out.printf("\t\t Track contains valid ISO image");
                }
            } else if (c instanceof DAOInfoChunk) {
                DAOInfoChunk cc = (DAOInfoChunk)c;
                out.printf("Recording type:\t Disc-At-Once (DAO)%n");
                int i = 1;
                for (DAOInfoChunk.TrackInfo t : cc.tracks) {
                    out.printf("  Track #%d:\t mode %d, sector size: %d, %n" +
                        "\t\t index #0 offset: %d, index #1 offset: %d, %n" +
                        "\t\t size: %d%n",
                        i++, t.mode, t.sectorSize, t.offsetIndex0, t.offsetIndex1,
                        t.offsetTrackEnd - t.offsetIndex0);
                    if (isISOImage(t.offsetIndex1))
                        out.printf("\t\t Track contains valid ISO image%n");
                }
            } else if (c instanceof CDTextChunk) {
                CDTextChunk cc = (CDTextChunk)c;
                out.printf("CDText:\t %d packs%n", cc.packs.length);
            } else if (c instanceof CueSheetChunk) {
                CueSheetChunk cc = (CueSheetChunk)c;
                out.printf("Cue Sheet:%n");
                for (CueSheetChunk.Index ind : cc.indices) {
                    if (ind.trackNumber == 0 || ind.trackNumber == 0xAA)
                        continue; // ignore leading and final index placeholders
                    out.printf("  Track #%d Index #%d:\t type: %s (0x%02X), " +
                        "block: %6d, %n" +
                        "\t\t\t block count: %6d, duration: %s%n",
                        ind.trackNumber, ind.indexNumber,
                        ind.getTypeDesc() , ind.type,
                        ind.block, ind.blockCount, ind.getTimeString());
                }
            } else if (c instanceof SessionInfoChunk) {
                SessionInfoChunk cc = (SessionInfoChunk)c;
                out.printf("Session #%d:\t total of %d tracks%n",
                    sessions++, cc.tracksInSession);
            } else if (c instanceof DiscInfoChunk) {
                DiscInfoChunk cc = (DiscInfoChunk)c;
                out.printf("Finalized:\t %s%n",
                    cc.finalization == DiscInfoChunk.FINALIZED ? "yes" : "no");
            }
        }
        System.out.println();
    }

    /**
     * The main entry point of the command line utility.
     *
     * @param args the command line arguments
     */
    public static void main(String[] args) {
        String infilename = null;
        String outfilename = null;
        String cuefilename = null;
        short track = 1;
        // parse params
        for (int i = 0; i < args.length - 1; i++) {
            String a = args[i];
            String v = args[++i];
            if (a.equals("-i"))
                infilename = v;
            else if (a.equals("-o"))
                outfilename = v;
            else if (a.equals("-cue"))
                cuefilename = v;
            else if (a.equals("-t"))
                track = Short.parseShort(v);
            else {
                // unknown option - show usage
                System.out.printf("error: unknown option '%s'%n", a);
                infilename = null;
                break;
            }
        }
        // show usage
        if (infilename == null) {
            System.out.println();
            System.out.println("usage: NRGReader <options>");
            System.out.println("  options are:");
            System.out.println("  -i <filename>\t\tthe NRG input file");
            System.out.println("  -o <filename>\t\tthe output (ISO/BIN) filename to extract");
            System.out.println("  -cue <file>\t\tthe output CUE filename to extract");
            System.out.println("  -t <track>\t\tthe track number to extract (default is 1)");
            return;
        }
        // read and extract
        NRGReader r = new NRGReader();
        try {
            System.out.printf("opening %s...%n", infilename);
            r.open(infilename);
            r.printInfo();
            if (outfilename != null) {
                System.out.printf("extracting track #%d to %s...%n", track, outfilename);
                r.extract(outfilename, track);
            }
            if (cuefilename != null) {
                System.out.printf("extracting CUE info to %s...%n", cuefilename);
                r.extractCue(outfilename, cuefilename);
            }
        } catch (Throwable t) {
            System.err.println("error: " + t);
        } finally {
            try {
                r.close();
            } catch (IOException ignore) {}
        }
    }
}