/*
 * 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.util.Collections;
import java.util.Date;
import java.util.Iterator;
import java.util.Set;
import java.util.Map;
import java.util.TreeMap;
import java.util.ArrayList;

import dumphd.util.Utils;

/**
 * A class holding all keys to decrypt a disc.
 * 
 * The constructor and the set methodes actually make a copy of their given arguments, therefore it is safe to modify them after a method call.
 * The get methodes return a reference to the internal members, so don't modify these! All getters may return null if not noted otherwise.
 * 
 * @author KenD00 
 */
public class KeyData {

   public final static int DATA_TITLE = 2;
   public final static int DATA_DATE = 4;
   public final static int DATA_MEK = 8;
   public final static int DATA_VIDBN = 16;
   public final static int DATA_VUKPAK = 32;
   public final static int DATA_TUK = 64;
   public final static int DATA_UM = 128;
   public final static int DATA_COMMENT = 256;
   public final static int DATA_ALL = 0xFFFFFFFF;

   /**
    * The DiscID
    */
   private byte[] discId = null;
   /**
    * The Title
    */
   private String title = null;
   /**
    * The Date
    */
   private Date date = null;
   /**
    * An optional comment
    */
   private String comment = null;

   /**
    * The Media Key
    */
   private byte[] mek = null;
   /**
    * The VolumeID or Binding Nonces.
    * A Volume ID has the Key 0.
    */
   private TreeMap<Integer, byte[]> vids = new TreeMap<Integer, byte[]>();
   /**
    * The Volume Unique Key or Protected Area Keys.
    * A Volume Unique Key has the Key 0.
    */
   private TreeMap<Integer, byte[]> vuks = new TreeMap<Integer, byte[]>();
   /**
    * The Title Keys / CPS Unit Keys. There is no difference made between Title Keys and CPS Unit Keys.
    * Key is the key number.
    */
   private TreeMap<Integer, byte[]> tuks = new TreeMap<Integer, byte[]>();

   /**
    * Head / Title/Clip Unit mappings for BD directories
    */
   private ArrayList<UnitMappingSet> umss = new ArrayList<UnitMappingSet>(5); 


   /**
    * Constructs a new KeyData Object with the given DiscID.
    * 
    * @param discId The DiscID
    * @param offset Offset into the byte array to start reading from. 20 Bytes get read.
    */
   public KeyData(byte[] discId, int offset) {
      this.discId = new byte[20];
      System.arraycopy(discId, offset, this.discId, 0, this.discId.length);
   }


   /**
    * Returns the DiscID.
    * 
    * @return The DiscID, this is never null
    */
   public byte[] getDiscId() {
      return discId;
   }

   /**
    * Returns the Title.
    * 
    * @return The Title
    */
   public String getTitle() {
      return title;
   }

   /**
    * Returns the Date.
    * 
    * @return The Date
    */
   public Date getDate() {
      return date;
   }
   
   /**
    * Returns the Comment.
    * 
    * @return The Comment
    */
   public String getComment() {
      return comment;
   }

   /**
    * Returns the Media Key.
    * 
    * @return The Media Key
    */
   public byte[] getMek() {
      return mek;
   }

   /**
    * Returns the Volume ID.
    * 
    * This method is the same as getBn(0).
    * 
    * @return The Volume ID
    */
   public byte[] getVid() {
      return vids.get(0);
   }

   /**
    * Returns a Binding Nonce.
    *  
    * @param index The index of the Binding Nonce, counting starts at 0
    * @return The Binding Nonce with the given index
    */
   public byte[] getBn(int index) {
      return vids.get(index);
   }

   /**
    * Returns the Volume Unique Key.
    * 
    * This method is the same as getPak(0).
    * 
    * @return The Volume Unique Key
    */
   public byte[] getVuk() {
      return vuks.get(0);
   }

   /**
    * Returns a Protected Area Key.
    *  
    * @param index The index of the Protected Area Key, counting starts at 0
    * @return The Protected Area Key with the given index
    */
   public byte[] getPak(int index) {
      return vuks.get(index);
   }

   /**
    * Returns a Title Key / CPS Unit Key. There is no difference made between Title Keys and CPS Unit Keys.
    * 
    * @param index The index of the key, counting starts at 1
    * @return The Key with the given index
    */
   public byte[] getTuk(int index) {
      return tuks.get(index);
   }

   /**
    * Returns the CPS Unit Number for the given head element.
    *  
    * @param unitMappingSet Number of the unitMappingSet to query. For BDMV this is always 0, for BDAV this is the number of the BDAV directory,
    *                       the Basic BDAV directory has the number 0
    * @param henr BD-ROM: 0 = First Playback, 1 = Top Menu
    *             BR-R  : 0 = Menu file, 1 = Mark file
    * @return CPS Unit number, 0 if no mapping present
    */
   public int getHeadUnit(int unitMappingSet, int henr) {
      if (unitMappingSet < umss.size())  {
         UnitMappingSet ums = umss.get(unitMappingSet);
         if (henr < ums.hm.size()) {
            return ums.hm.get(henr);
         }
      }
      return 0;
   }

   /**
    * Returns the CPS Unit Number for the Title / Clip number.
    *  
    * @param unitMappingSet Number of the unitMappingSet to query. For BDMV this is always 0, for BDAV this is the number of the BDAV directory,
    *                    the Basic BDAV directory has the number 0
    * @param tcnr Title number (BD-ROM) / Clip number (BD-R)
    * @return CPS Unit number, 0 if no mapping present
    */
   public int getTcUnit(int unitMappingSet, int tcnr) {
      // The ArrayList starts at 0, the Title/Clip number at 1
      tcnr--;
      if (unitMappingSet < umss.size()) {
         UnitMappingSet kis = umss.get(unitMappingSet);
         if (tcnr < kis.tcm.size()) {
            return kis.tcm.get(tcnr);
         }
      }
      return 0;
   }

   /**
    * Sets the Title.
    * 
    * @param title The Title, null clears the Title
    */
   public void setTitle(String title) {
      this.title = title;
   }

   /**
    * Sets the Date.
    * 
    * @param date The Date, null clears the Date
    */
   public void setDate(Date date) {
      if (date != null) {
         this.date = (Date)date.clone();
      } else {
         this.date = null;
      }
   }
   
   /**
    * Sets the Comment.
    * 
    * @param comment The Comment, null clears the Comment
    */
   public void setComment(String comment) {
      this.comment = comment;
   }

   /**
    * Sets the Media Key.
    * 
    * @param mek The Media Key, null clears the Media Key.
    * @param offset Offset into the byte array to start reading from. 16 Bytes get read.
    */
   public void setMek(byte[] mek, int offset) {
      if (mek != null) {
         this.mek = new byte[16];
         System.arraycopy(mek, offset, this.mek, 0, this.mek.length);
      } else {
         this.mek = null;
      }
   }

   /**
    * Sets the Volume ID.
    * 
    * This method is the same as setBn(0, vid, offset).
    * 
    * @param vid The Volume ID, null clears the Volume ID
    * @param offset Offset into the byte array to start reading from. 16 Bytes get read. 
    */
   public void setVid(byte[] vid, int offset) {
      setKey(vids, 0, vid, offset);
   }

   /**
    * Sets a Binding Nonce. Do either set Binding Nonces or a Volume ID but not both!
    * 
    * @param index The index of the Binding Nonce
    * @param bn The Binding Nonce, null clears the Binding Nonce
    * @param offset Offset into the byte array to start reading from. 16 Bytes get read. 
    */
   public void setBn(int index, byte[] bn, int offset) {
      setKey(vids, index, bn, offset);
   }

   /**
    * Sets the Volume Unique Key.
    * 
    * This method is the same as setPak(0, pak, offset).
    * 
    * @param vuk The Volume Unique Key, null clears the Volume Unique Key
    * @param offset Offset into the byte array to start reading from. 16 Bytes get read.
    */
   public void setVuk(byte[] vuk, int offset) {
      setKey(vuks, 0, vuk, offset);
   }

   /**
    * Sets a Protected Area Key. Do either set Protected Area Keys or a Volume Unique Key but not both!
    * 
    * @param index The index of the Protected Area Key
    * @param pak The Protected Area Key, null clears the Protected Area Key
    * @param offset Offset into the byte array to start reading from. 16 Bytes get read. 
    */
   public void setPak(int index, byte[] pak, int offset) {
      setKey(vuks, index, pak, offset);
   }

   /**
    * Sets a Title Key / CPS Unit Key. There is no difference made between Title Keys and CPS Unit Keys.
    * 
    * @param index The index of the key, must be > 0
    * @param tuk The Title Key / CPS Unit Key, null removes the Title Key / CPS Unit Key
    * @param offset Offset into the byte array to start reading from. 16 Bytes get read.
    */
   public void setTuk(int index, byte[] tuk, int offset) {
      setKey(tuks, index, tuk, offset);
   }

   /**
    * Sets a Head unit.
    * 
    * TODO: There is no way to remove such a mapping
    * 
    * @param unitMappingSet unitMappingSet to use
    * @param henr Head element number
    * @param unit Number of the CPS Unit
    */
   public void setHeadUnit(int unitMappingSet, int henr, int unit) {
      // Add empty unitMappingSets until we have enough to access the requested one
      while (umss.size() <= unitMappingSet) {
         umss.add(new UnitMappingSet());
      }
      // Now our unitMappingSet is present or unitMappingSet is negative
      UnitMappingSet ums = umss.get(unitMappingSet);
      // Add empty mappings until we have one less than the requested one
      while (ums.hm.size() < henr) {
         ums.hm.add(0);
      }
      // This checks implicit if henr is negative
      if (ums.hm.size() == henr) {
         ums.hm.add(unit);
      } else {
         ums.hm.set(henr, unit);
      }
   }

   /**
    * Sets a Title / Clip unit.
    * 
    * TODO: There is no way to remove such a mapping
    * 
    * @param unitMappingSet unitMappingSet to use
    * @param tcnr Title / Clip number
    * @param unit Number of the CPS Unit
    */
   public void setTcUnit(int unitMappingSet, int tcnr, int unit) {
      // The ArrayList starts at 0, the Title/Clip number at 1
      tcnr--;
      // Add empty unitMappingSets until we have enough to access the requested one
      while (umss.size() <= unitMappingSet) {
         umss.add(new UnitMappingSet());
      }
      // Now our unitMappingSet is present or unitMappingSet is negative
      UnitMappingSet ums = umss.get(unitMappingSet);
      // Add empty mappings until we have one less than the requested one
      while (ums.tcm.size() < tcnr) {
         ums.tcm.add(0);
      }
      // This checks implicit if tcnr is negative
      if (ums.tcm.size() == tcnr) {
         ums.tcm.add(unit);
      } else {
         ums.tcm.set(tcnr, unit);
      }
   }

   /**
    * Returns a bitmask of all present data entries. This is a binary OR of all present DATA_* values.
    * 
    * @return Bitmask of all present data entries
    */
   public int dataMask() {
      int mask = 0;
      if (title != null) {
         mask |= DATA_TITLE;
      }
      if (date != null) {
         mask |= DATA_DATE;
      }
      if (comment != null) {
         mask |= DATA_COMMENT;
      }
      if (mek != null) {
         mask |= DATA_MEK;
      }
      if (vids.size() > 0) {
         mask |= DATA_VIDBN;
      }
      if (vuks.size() > 0) {
         mask |= DATA_VUKPAK;
      }
      if (tuks.size() > 0) {
         mask |= DATA_TUK;
      }
      if (umss.size() > 0) {
         mask |= DATA_UM;
      }
      return mask;
   }

   /**
    * Returns the number of present Binding Nonces.
    * 
    * @return The number of present Binding Nonces, 1 if a VID is set
    */
   public int bnCount() {
      return vids.size();
   }

   /**
    * Returns the number of present Protected Area Keys.
    * 
    * @return The number of present Protected Area Keys, 1 if a VUK is set
    */
   public int pakCount() {
      return vuks.size();
   }

   /**
    * Returns the number of present Title Keys / CPS Unit Keys.
    * 
    * @return The number of present Title Keys / CPS Unit Keys
    */
   public int tukCount() {
      return tuks.size();
   }

   /**
    * Returns an unmodifiable set view of the indices from the present Binding Nonces.
    * 
    * @return The indices of the present Binding Nonces
    */
   public Set<Integer> bnIdx() {
      return Collections.unmodifiableSet(vids.keySet());
   }

   /**
    * Returns an unmodifiable set view of the indices from the present Protected Area Keys.
    * 
    * @return The indices of the present Protected Area Keys
    */
   public Set<Integer> pakIdx() {
      return Collections.unmodifiableSet(vuks.keySet());
   }

   /**
    * Returns an unmodifiable set view of the indices from the present Title Keys / CPS Unit Keys.
    * 
    * @return The indices of the present Title Keys / CPS Unit Keys
    */
   public Set<Integer> tukIdx() {
      return Collections.unmodifiableSet(tuks.keySet());
   }


   public String toString() {
      StringBuffer out = new StringBuffer(8 * 1024);
      out.append("DiscID      : ");
      out.append(Utils.toHexString(discId, 0, discId.length));
      out.append("\nTitle       : ");
      if (title != null) {
         out.append(title);
      } else {
         out.append("N/A");
      }
      out.append("\nDate        : ");
      if (date != null) {
         out.append(String.format("%1$tY-%1$tm-%1$td", date));
      } else {
         out.append("N/A");
      }
      out.append("\nComment     : ");
      if (comment != null) {
         out.append(comment);
      } else {
         out.append("N/A");
      }
      out.append("\nMEK         : ");
      if (mek != null) {
         out.append(Utils.toHexString(mek, 0, mek.length));
      } else {
         out.append("N/A");
      }
      out.append("\nVID / BN's  : ");
      printKeys(out, vids);
      out.append("\nVUK / PAK's : ");
      printKeys(out, vuks);
      out.append("\nTUK's       : ");
      printKeys(out, tuks);
      if (umss.size() > 0) {
         out.append("\nUM's        : ");
         out.append(umss.size());
         for (int set = 0; set < umss.size(); set++) {
            out.append(String.format("\n              Set %1$d", set));
            UnitMappingSet ums = umss.get(set);
            for (int i = 0; i < ums.hm.size(); i++) {
               out.append(String.format("\n%1$12dH-%2$d", i, ums.hm.get(i)));
            }
            for (int i = 0; i < ums.tcm.size(); i++) {
               out.append(String.format("\n%1$13d-%2$d", i + 1, ums.tcm.get(i)));
            }
         }
      }
      return out.toString();
   }

   /**
    * Helper Method, puts / removes the given key from the given TreeMap.
    * 
    * @param tm The TreeMap to use
    * @param index The index of the key
    * @param key The key. If null the key at index gets removed
    * @param offset Offset into key where the actual key starts. 16 bytes get read.
    */
   private void setKey(TreeMap<Integer, byte[]> tm, int index, byte[] key, int offset) {
      if (key != null) {
         // Its more common that a new key gets inserted than an old one updated, therefore don't check for an old entry
         byte[] tmp = new byte[16];
         System.arraycopy(key, offset, tmp, 0, tmp.length);
         tm.put(index, tmp);
      } else {
         tm.remove(index);
      }
   }

   /**
    * Helper Method, prints the keys from the given TreeMap to the given StringBuffer.
    * 
    * @param out StringBuffer to write to
    * @param keyMap TreeMap to read from
    */
   private void printKeys(StringBuffer out, TreeMap<Integer, byte[]> keyMap) {
      int keyCount = keyMap.size();
      if (keyCount > 0) {
         out.append(keyCount);
         Iterator<Map.Entry<Integer, byte[]>> keyIt = keyMap.entrySet().iterator();
         while (keyIt.hasNext()) {
            Map.Entry<Integer, byte[]> keyEntry = keyIt.next();
            out.append(String.format("\n%1$13d-%2$s", keyEntry.getKey(), Utils.toHexString(keyEntry.getValue(), 0, keyEntry.getValue().length)));
         }
      }
      else {
         out.append("N/A");
      }
   }



   /**
    * Small structure to hold Head / Title/Clip mappings for a BD directory
    */
   private class UnitMappingSet {
      /**
       * Head Mappings: CPS Unit Number for First Playback and Top Menu (BD-ROM) or Menu and Mark file (BD-R)  
       */
      public ArrayList<Integer> hm = new ArrayList<Integer>(2);
      /**
       * Title/Clip Mappings: CPS Unit Number for title number (BD-ROM) or clip number (BD-R)
       */
      public ArrayList<Integer> tcm = new ArrayList<Integer>(200);
   }

}
