/*
 * Copyright (C) 2007-2009 KenD00
 * 
 * This file is part of DumpHD.
 * 
 * DumpHD 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 3 of the License, or
 * (at your option) any later version.
 * 
 * This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package dumphd.core;

import java.io.File;
import java.io.IOException;
import java.util.Iterator;
import java.util.LinkedList;

import dumphd.util.ByteArray;
import dumphd.util.ByteSource;
import dumphd.util.FileSource;
import dumphd.util.PackSource;
import dumphd.util.PESPack;
import dumphd.util.PESParserException;

/**
 * A MultiplexedArf is used to access an ARF multiplexed into an EVO. This class is used to build a list with the locations of the packets for one ARF.
 * Then a reader and a writer can be requested from this class which allow reading and writing to the ARF through the ByteSource interface. It is possible to
 * redirect the writes to demux the ARF while processing it. If the ARF gets not demuxed, the reader and writer use the same underlying ByteSource to access the ARF.
 * This class automatically pads the resulting ARF in the EVO (if not demuxed) when closed.
 * 
 * TODO: Don't throw IllegalStateExceptions, or at least write that into the javadoc? 
 * 
 * @author KenD00
 */
public final class MultiplexedArf {

   public final static int START_PACKET = 0x10;
   public final static int INTERMEDIATE_PACKET = 0x00;
   public final static int END_PACKET = 0x20;

   /**
    * This is the file that gets returned by the getFile() call from the reader and writer
    */
   private File arfFile = null;
   /**
    * The underlying ByteSource used to access the EVO
    */
   private ByteSource bs = null;
   /**
    * If true, the ARF gets demuxed, otherwise it is rewritten in place
    */
   private boolean demux = false;

   /**
    * The list with the records to the packs
    */
   private LinkedList<PackSource.PackRecord> packRecords = new LinkedList<PackSource.PackRecord>(); 
   /**
    * Current dataOffset, used when creating the packRecord list
    */
   private long currentDataOffset = 0;
   /**
    * If the MultiplexedArf if completed, no more data can be added and the reader and writer can get returned
    */
   private boolean complete = false;

   /**
    * The created reader returned by this object
    */
   private ByteSource reader = null;
   /**
    * The created writer returned by this object
    */
   private ByteSource writer = null;


   /**
    * Creates a new MultiplexedArf which is contained in the given EVO.
    * 
    * @param arfName The filename of the ARF, this get added behind the path of bs
    * @param bs The ByteSource of the EVO to use
    * @param demux If true, the writer demuxed the ARF, otherwise it rewrites it into the EVO
    */
   public MultiplexedArf(String arfName, ByteSource bs, boolean demux) {
      this.arfFile = new File(bs.getFile().getParentFile(), arfName);
      this.bs = bs;
      this.demux = demux;
   }

   /**
    * Adds a new data packet to the list. This method can only be called when the ARF is not completed.
    * 
    * @param packOffset Offset to the pack
    * @param payloadOffset Relative offset from the pack start to the payload start
    * @param payloadSize The size of the payload
    */
   public void addData(long packOffset, int payloadOffset, int payloadSize) {
      if (!complete) {
         packRecords.add(new PackSource.PackRecord(currentDataOffset, packOffset, payloadOffset, payloadSize));
         currentDataOffset += (long)payloadSize;
      } else throw new IllegalStateException("Cannot add additional data, MultiplexedArf already completed");
   }

   /**
    * Marks this MultiplexedArf as complete.
    */
   public void markComplete() {
      if (!complete) {
         complete = true;
         // If there are no entries in the list, add a zero-entry, or the PackSource will throw exceptions
         if (packRecords.size() == 0) {
            packRecords.add(new PackSource.PackRecord());
         }
         //Iterator<PackSource.PackRecord> it = packRecords.iterator();
         //while (it.hasNext()) {
         //    System.out.println(it.next().toString());
         //}
      }
   }

   /**
    * If not already done, creates a reader. The reader then gets returned.
    * This method may only be called after this MultiplexedArf has been marked as complete.
    * 
    * @return The reader to read the ARF
    * @throws IOException An I/O error during the creation of the reader occurred
    */
   public ByteSource getReader() throws IOException {
      if (complete) {
         if (reader == null) {
            reader = new PackSource(bs, ByteSource.R_MODE, arfFile, packRecords);
         }
         return reader;
      } else throw new IllegalStateException("Cannot create reader, MultiplexedArf not completed");
   }

   /**
    * If not already done, creates a writer. The writer then gets returned.
    * This method may only be called after this MultiplexedArf has been marked as complete.
    * 
    * @return The writer to write the ARF. Depending on the creation of this MultiplexedArf, the ARF gets demuxed or rewritten in place
    * @throws IOException An I/O error during the creation of the writer occurred
    */
   public ByteSource getWriter() throws IOException {
      if (complete) {
         if (writer == null) {
            if (demux) {
               writer = new FileSource(arfFile, ByteSource.RW_MODE);
            } else {
               writer = new PackSource(bs, ByteSource.W_MODE, arfFile, packRecords);
            }
         }
         return writer;
      } else throw new IllegalStateException("Canot create writer, MultiplexedArf not completed");
   }

   /**
    * Closes the created reader and writer and pads the EVO if necessary.
    * 
    * @throws IOException An I/O error occurred
    */
   public void close() throws IOException {
      StringBuffer exceptionMessage = new StringBuffer(128);
      if (writer != null) {
         if (!demux) {
            try {
               // Fix the packet structure of the underlying source
               // The PESPack to parse the pack data
               PESPack pack = new PESPack();
               // The ARF_PCK that gets modified
               PESPack.PESPacket packet = null;
               // Iterator over all PackRecords
               Iterator<PackSource.PackRecord> it = packRecords.iterator();
               // The current PackRecord
               PackSource.PackRecord packRecord = it.next();
               // The byte array to read the pack data in
               byte[] packBuffer = new byte[PESPack.PACK_LENGTH];
               // The maximum written position
               long maxWritten = ((PackSource)writer).maxWrittenPosition();
               // Check if the ARF is more than one packet big
               if (maxWritten <= (long)packRecord.payloadSize) {
                  // The multiplexed ARF occupies only one packet
                  // FIXME: Temporary message until i know how to handle only one packet ARFs
                  System.out.println(String.format("0x%1$010X Warning! Multiplexed ARF is only one packet big!", packRecord.packetOffset));
                  // Read the pack
                  bs.setPosition(packRecord.packetOffset);
                  bs.read(packBuffer, 0, packBuffer.length);
                  try {
                     // Parse the pack and set the packet type to START_PACKET
                     // FIXME: Is this the correct packet type?? If not change it to the correct one!
                     pack.parse(packBuffer, 0, false);
                     packet = pack.getPacket(0);
                     ByteArray.setByte(packBuffer, packet.payloadOffset(), MultiplexedArf.START_PACKET);
                     // Adjust the payload length and pad the pack. Need to add the size of the header and the filename field
                     packet.setPayloadLength((int)maxWritten + 258);
                     pack.pad(0);
                     // Write the pack back
                     bs.setPosition(packRecord.packetOffset);
                     bs.write(packBuffer, 0, packBuffer.length);
                  }
                  catch (PESParserException e) {
                     if (exceptionMessage.length() > 0) {
                        exceptionMessage.append(", ");
                     }
                     exceptionMessage.append(String.format("0x%1$010X Error padding stream, PES Parser Exception: ", packRecord.packetOffset));
                     //exceptionMessage.append("Error padding stream, PES Parser Exception: ");
                     exceptionMessage.append(e.getMessage());
                  }
               } else {
                  // The multiplexed ARF occupies multiple packets
                  // Search the end packet. Because the decrypted ARF is always smaller than the original one we don't need to adjust the
                  // INTERMEDIATE_PACKETs, we only need to set the new END_PACKET which will be before the original one.
                  //System.out.println("ARF occupies more than one packet, searching end packet");
                  while (maxWritten > packRecord.dataOffset + (long)packRecord.payloadSize) {
                     if (it.hasNext()) {
                        packRecord = it.next();
                     } else throw new IOException("Unexpected EOF");
                  }
                  // Found the new END_PACKET, read it in
                  //System.out.println(String.format("End packet found at position: 0x%1$010X, maxWritten: 0x%2$010X", packRecord.packetOffset, maxWritten));
                  bs.setPosition(packRecord.packetOffset);
                  bs.read(packBuffer, 0, packBuffer.length);
                  try {
                     // Parse the pack and set the packet type to END_PACKET
                     pack.parse(packBuffer, 0, false);
                     packet = pack.getPacket(0);
                     ByteArray.setByte(packBuffer, packet.payloadOffset(), MultiplexedArf.END_PACKET);
                     // Adjust the payload length and pad the pack. Need to add the size of the header
                     packet.setPayloadLength(3 + (int)(maxWritten - packRecord.dataOffset));
                     pack.pad(0);
                     // Write the pack back
                     bs.setPosition(packRecord.packetOffset);
                     bs.write(packBuffer, 0, packBuffer.length);
                  }
                  catch (PESParserException e) {
                     if (exceptionMessage.length() > 0) {
                        exceptionMessage.append(", ");
                     }
                     exceptionMessage.append(String.format("0x%1$010X Error padding stream, PES Parser Exception: ", packRecord.packetOffset));
                     //exceptionMessage.append("Error padding stream, PES Parser Exception: ");
                     exceptionMessage.append(e.getMessage());
                  }
               }
               // Pad the whole packs which follow the new END_PACKET 
               while (it.hasNext()) {
                  packRecord = it.next();
                  // Read the pack in
                  bs.setPosition(packRecord.packetOffset);
                  bs.read(packBuffer, 0, packBuffer.length);
                  try {
                     // Parse and pad the whole pack
                     pack.parse(packBuffer, 0, false);
                     pack.pad(-1);
                     // Write the pack back
                     bs.setPosition(packRecord.packetOffset);
                     bs.write(packBuffer, 0, packBuffer.length);
                  }
                  catch (PESParserException e) {
                     if (exceptionMessage.length() > 0) {
                        exceptionMessage.append(", ");
                     }
                     exceptionMessage.append(String.format("0x%1$010X Error padding stream, PES Parser Exception: ", packRecord.packetOffset));
                     //exceptionMessage.append("Error padding stream, PES Parser Exception: ");
                     exceptionMessage.append(e.getMessage());
                  }
               }
            }
            catch (IOException e) {
               if (exceptionMessage.length() > 0) {
                  exceptionMessage.append(", ");
               }
               exceptionMessage.append("Error padding stream: ");
               exceptionMessage.append(e.getMessage());
            }
         }
         try {
            writer.close();
         }
         catch (IOException e) {
            if (exceptionMessage.length() > 0) {
               exceptionMessage.append(", ");
            }
            exceptionMessage.append("Error closing writer: ");
            exceptionMessage.append(e.getMessage());
         }
      }
      if (reader != null) {
         try {
            reader.close();
         }
         catch (IOException e) {
            if (exceptionMessage.length() > 0) {
               exceptionMessage.append(", ");
            }
            exceptionMessage.append("Error closing reader: ");
            exceptionMessage.append(e.getMessage());
         }
      }
      // Release resources
      reader = null;
      writer = null;
      packRecords = null;
      bs = null;
      arfFile = null;
      // If there was an exception, now throw it
      if (exceptionMessage.length() > 0) {
         throw new IOException(exceptionMessage.toString());
      }
   }

}
