context: /*
   * Copyright (c) 2007, 2008 dmz
   * 
   * This file is part of TsRemux.
   * 
   * TsRemux 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.
   * 
   * TsRemux 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 TsRemux.  If not, see <http://www.gnu.org/licenses/>.
   * 
   */
  
  using System;
  using System.Collections.Generic;
  using System.Text;
  using System.IO;
  using System.ComponentModel;
  using System.Collections;
  
  namespace TsRemux
  {
      public struct EpElement
      {
          private Int64 pts;
          private UInt32 spn;
  
          public EpElement(Int64 pts, UInt32 spn)
          {
              this.pts = pts;
              this.spn = spn;
          }
  
          public Int64 PTS
          {
              get { return pts; }
          }
  
          public UInt32 SPN
          {
              get { return spn; }
          }
      }
  
      // Write stream to MKV directly, just experimenting atm.
      public class MkvMux : Muxer
      {
          private FileStream _FileStream;
          private StreamWriter _StreamWriter;
          private List<StreamInfo> _StreamsToKeep;
  
          public MkvMux(string fileName, List<StreamInfo> StreamsToKeep, bool fAsync, bool fProcessAudio)
          {
              _FileStream = new FileStream(fileName, FileMode.Create, FileAccess.Write, FileShare.None, Constants.DISK_BUFFER, fAsync);
              _StreamWriter = new StreamWriter(_FileStream);
              _StreamsToKeep = StreamsToKeep;
          }
  
          public override void MuxPacket(PesPacket pp)
          {
              // TODO: Fix, just QnD ...
              bool bKeep = false;
              foreach (StreamInfo si in _StreamsToKeep)
              {
                  if (si.ElementaryPID == pp.PID)
                      bKeep = true;
              }
  
              if (bKeep)
              {
                  byte[] data = pp.GetData();
                  _FileStream.Write(data, 0, data.Length);
              }
          }
  
          public override void Close()
          {
              _StreamWriter.Close();
              _FileStream.Close();
              _FileStream = null;
          }
  
          public override void PcrChanged(Int64 pcr)
          {
          }
  
          public override EpElement[] GetEpData()
          {
              throw new NotSupportedException("Not MKV compatible.");
          }
  
          public override StreamInfo[] GetPsi()
          {
              throw new NotSupportedException("Not MKV compatible.");
          }
  
          public override UInt32 GetCurrentPacketNumber()
          {
              throw new NotSupportedException("Not MKV compatible.");
          }
      }
  
      // Simple stream demuxer, writes every stream to one file.
      public class DeMux : Muxer
      {
          // Represents one used stream.
          private struct DeMuxStream
          {
              private FileStream _FileStream;
              private int _PID;
  
              public FileStream FileStream
              {
                  get { return _FileStream; }
              }
  
              public int PID
              {
                  get { return _PID;  }
              }
  
              public DeMuxStream(FileStream FS, int PID)
              {
                  _FileStream = FS;
                  _PID = PID;
              }
          }
  
          // Used streams.
          private List<DeMuxStream> _Streams;
  
          // CTor, create list of stream including their stream write instances.
          public DeMux(string fileName, List<StreamInfo> StreamsToKeep, bool fAsync, bool fProcessAudio)
          {
              _Streams = new List<DeMuxStream>();
  
              foreach (StreamInfo si in StreamsToKeep)
              {
                  // Give stream a useful extension.
                  string sExt = "unknown";
                  switch(si.StreamType)
                  {
                      case ElementaryStreamTypes.AUDIO_STREAM_AC3:
                      case ElementaryStreamTypes.AUDIO_STREAM_AC3_PLUS:
                      case ElementaryStreamTypes.AUDIO_STREAM_AC3_TRUE_HD:
                      case ElementaryStreamTypes.SECONDARY_AUDIO_AC3_PLUS:
                          sExt = "ac3";
                          break;
  
                      case ElementaryStreamTypes.AUDIO_STREAM_DTS:
                      case ElementaryStreamTypes.AUDIO_STREAM_DTS_HD:
                      case ElementaryStreamTypes.AUDIO_STREAM_DTS_HD_MASTER_AUDIO:
                      case ElementaryStreamTypes.SECONDARY_AUDIO_DTS_HD:
                          sExt = "dts";
                          break;
  
                      case ElementaryStreamTypes.AUDIO_STREAM_LPCM:
                          sExt = "pcm";
                          break;
  
                      case ElementaryStreamTypes.AUDIO_STREAM_MPEG1:
                      case ElementaryStreamTypes.AUDIO_STREAM_MPEG2:
                          sExt = "mpa";
                          break;
  
                      case ElementaryStreamTypes.VIDEO_STREAM_H264:
                          sExt = "264";
                          break;
  
                      case ElementaryStreamTypes.VIDEO_STREAM_MPEG1:
                          sExt = "m1v";
                          break;
  
                      case ElementaryStreamTypes.VIDEO_STREAM_MPEG2:
                          sExt = "m2v";
                          break;
  
                      case ElementaryStreamTypes.VIDEO_STREAM_VC1:
                          sExt = "vc1";
                          break;
                  }
  
                  string sFileName = string.Format("{0}_pid{1}.{2}", fileName, si.ElementaryPID, sExt);
  
                  FileStream FS = new FileStream(sFileName, FileMode.Create, FileAccess.Write, FileShare.None, Constants.DISK_BUFFER, fAsync);
                  _Streams.Add(new DeMuxStream(FS, si.ElementaryPID));
              }
          }
  
          // If PesPacket contains valid stream data, store it
          // in the respective output stream.
          public override void MuxPacket(PesPacket pp)
          {
              foreach (DeMuxStream S in _Streams)
              {
                  if (pp.PID == S.PID)
                  {
                      byte[] data = pp.GetData();
                      S.FileStream.Write(data, 0, data.Length);
                  }
              }
          }
  
          // Finish (De)Muxing.
          public override void Close()
          {
              foreach (DeMuxStream S in _Streams)
                  S.FileStream.Close();
              _Streams.Clear();
          }
  
          // Ignore atm.
          public override void PcrChanged(Int64 pcr)
          {
          }
  
          public override EpElement[] GetEpData()
          {
              throw new NotSupportedException("Not MKV compatible.");
          }
  
          public override StreamInfo[] GetPsi()
          {
              throw new NotSupportedException("Not MKV compatible.");
          }
  
          public override UInt32 GetCurrentPacketNumber()
          {
              throw new NotSupportedException("Not MKV compatible.");
          }
      }
  
      // Original Muxer by dmz / spacecat56 ?.
      public class BlueMux : Muxer
      {
          private TsFileType fileType;
          private List<StreamInfo> StreamsToKeep;
          private TsIo tsiow;
          private List<TsPacket> buffer;
          private Int64 currentPcr;
          private Int64 lastPcr;
          private Int64 lastDelta;
          private Int64 pcrOffset;
          private Int64 currentPts;
          private Int64 lastPts;
          private Int64 ptsDelta;
          private Int64 ptsOffset;
          private Int64 ptsCount;
          private PatPacket pat;
          private PmtPacket pmt;
          private SitPacket sit;
          private PcrPacket pcrPacket;
          private Dictionary<ushort, ushort> pidMappings;
          private byte[] header;
          private Dictionary<ushort, byte> Continuities;
          private byte[] supHeader;
          private FileStream fsw;
          private byte VideoType;
          private PesPacket lastPacket;
          private byte sitCount;
          private UInt32 packetCount;
          private List<EpElement> epData;
          private List<StreamInfo> pmtStreams;
          private Dictionary<ushort, Int64> lastPtsList;
          private Dictionary<ushort, List<byte>> soundFrames;
          private bool processAudio;
          private bool MlpToAc3;
  
          public BlueMux(string fileName, TsFileType fileType, List<StreamInfo> StreamsToKeep, bool fAsync, bool fProcessAudio, bool fMlpToAc3)
          {
              fsw = new FileStream(fileName, FileMode.Create, FileAccess.Write, FileShare.None, Constants.DISK_BUFFER, fAsync);
              this.tsiow = new TsIo(null, fsw, Constants.DISK_BUFFER);
              this.fileType = fileType;
              this.StreamsToKeep = StreamsToKeep;
              this.buffer = new List<TsPacket>();
              this.currentPcr = -1;
              this.lastPcr = -1;
              this.lastDelta = -1;
              this.pcrOffset = -1;
              this.lastPts = -1;
              this.currentPts = -1;
              this.ptsDelta = -1;
              this.ptsOffset = 0;
              this.ptsCount = -1;
              this.pcrPacket = null;
              this.header = new byte[4];
              this.supHeader = new byte[10];
              this.supHeader[0] = (byte)'P';
              this.supHeader[1] = (byte)'G';
              this.VideoType = 0;
              this.lastPacket = null;
              this.sitCount = 0;
              this.packetCount = 0;
              this.epData = new List<EpElement>();
              this.pmtStreams = null;
              this.lastPtsList = new Dictionary<ushort, long>();
              this.soundFrames = new Dictionary<ushort, List<byte>>();
              this.processAudio = fProcessAudio;
              this.MlpToAc3 = fMlpToAc3;
              CreatePsi();
          }
  
          public override UInt32 GetCurrentPacketNumber()
          {
              return packetCount + (UInt32)buffer.Count;
          }
  
          public override EpElement[] GetEpData()
          {
              return epData.ToArray();
          }
  
          public override StreamInfo[] GetPsi()
          {
              return pmtStreams.ToArray();
          }
          /*
          public UInt32 CurrentPacketNumber
          {
              get { return packetCount + (UInt32)buffer.Count; }
          }
  
          public EpElement[] EpData
          {
              get
              {
                  return epData.ToArray();
              }
          }
  
          public StreamInfo[] Psi
          {
              get { return pmtStreams.ToArray(); }
          }
          */
  
          public override void MuxPacket(PesPacket pp)
          {
              if (fileType == TsFileType.TS || fileType == TsFileType.M2TS || fileType == TsFileType.BLU_RAY)
                  MuxTsPacket(pp);
              else if (fileType == TsFileType.SUP_ELEMENTARY && pidMappings.ContainsKey(pp.PID))
              {
                  byte[] data = pp.GetData();
                  PesHeader ph = pp.GetHeader();
                  if (ph.HasPts)
                  {
                      Int64 Pts = ph.Pts;
                      Int64 Dts = 0;
                      if (ph.HasDts)
                          Dts = ph.Dts;
                      supHeader[2] = (byte)((Pts >> 24) & 0xff);
                      supHeader[3] = (byte)((Pts >> 16) & 0xff);
                      supHeader[4] = (byte)((Pts >> 8) & 0xff);
                      supHeader[5] = (byte)(Pts & 0xff);
                      supHeader[6] = (byte)((Dts >> 24) & 0xff);
                      supHeader[7] = (byte)((Dts >> 16) & 0xff);
                      supHeader[8] = (byte)((Dts >> 8) & 0xff);
                      supHeader[9] = (byte)(Dts & 0xff);
                      tsiow.Write(supHeader, 0, supHeader.Length);
                  }
                  tsiow.Write(data, ph.HeaderLength + 9, data.Length - (ph.HeaderLength + 9));
              }
              else if (fileType == TsFileType.PES_ELEMENTARY && pidMappings.ContainsKey(pp.PID))
              {
                  byte[] data = pp.GetData();
                  tsiow.Write(data, 0, data.Length);
              }
              else if (fileType == TsFileType.ELEMENTARY && pidMappings.ContainsKey(pp.PID))
              {
                  byte[] data = pp.GetData();
                  PesHeader ph = pp.GetHeader();
                  tsiow.Write(data, ph.HeaderLength + 9, data.Length - (ph.HeaderLength + 9));
              }
          }
  
          public override void PcrChanged(Int64 pcr)
          {
              lastPcr = currentPcr;
              currentPcr = pcr;
              Int64 newPcr = -1;
              if (lastPcr == -1)
              {
                  pcrOffset = currentPcr - Constants.MPEG2TS_CLOCK_RATE;
                  newPcr = currentPcr - pcrOffset;
                  if (newPcr < 0)
                      newPcr += Constants.MAX_MPEG2TS_CLOCK;
                  if (newPcr >= Constants.MAX_MPEG2TS_CLOCK)
                      newPcr -= Constants.MAX_MPEG2TS_CLOCK;
                  pcrPacket = new PcrPacket(newPcr,0,Constants.DEFAULT_PCR_PID);
                  PcrPacket pp = new PcrPacket(newPcr, pcrPacket.ContinuityCounter, Constants.DEFAULT_PCR_PID);
                  buffer.Add(pp);
                  return;
              }
              Int64 delta = currentPcr - lastPcr;
              if (delta < 0)
                  delta += Constants.MAX_MPEG2TS_CLOCK;
              if (delta > Constants.MPEG2TS_CLOCK_RATE)
              {
                  // discontinuity - adjust offset
                  Int64 predictedPcr = lastPcr + (buffer.Count * lastDelta);
                  if (predictedPcr >= Constants.MAX_MPEG2TS_CLOCK)
                      predictedPcr -= Constants.MAX_MPEG2TS_CLOCK;
                  pcrOffset += currentPcr - predictedPcr;
                  pcrPacket.IncrementContinuityCounter();
                  PcrPacket pp = new PcrPacket(newPcr, pcrPacket.ContinuityCounter, Constants.DEFAULT_PCR_PID);
                  buffer.Add(pp);
                  return;
              }
              newPcr = currentPcr - pcrOffset;
              if (newPcr < 0)
                  newPcr += Constants.MAX_MPEG2TS_CLOCK;
              if (newPcr >= Constants.MAX_MPEG2TS_CLOCK)
                  newPcr -= Constants.MAX_MPEG2TS_CLOCK;
              pat.IncrementContinuityCounter();
              buffer.Add(new PatPacket(pat.GetData()));
              pmt.IncrementContinuityCounter();
              buffer.Add(new PmtPacket(pmt.GetData()));
              sitCount += 1;
              sitCount %= 10;
              if (sitCount == 0)
              {   // one SIT for every 10 PMT/PAT
                  sit.IncrementContinuityCounter();
                  buffer.Add(new SitPacket(sit.GetData()));
              }
              pcrPacket.IncrementContinuityCounter();
              PcrPacket p = new PcrPacket(newPcr, pcrPacket.ContinuityCounter, Constants.DEFAULT_PCR_PID);
              buffer.Add(p);
              lastDelta = delta / buffer.Count;
              Int64 stamp = lastPcr % Constants.MAX_BLURAY_CLOCK;
              for (int i = 0; i < buffer.Count; i++)
              {
                  stamp += lastDelta;
                  stamp %= Constants.MAX_BLURAY_CLOCK;
                  header[0] = (byte)((stamp >> 24) & 0x3f);
                  header[1] = (byte)((stamp >> 16) & 0xff);
                  header[2] = (byte)((stamp >> 8) & 0xff);
                  header[3] = (byte)(stamp & 0xff);
                  if (fileType == TsFileType.M2TS || fileType == TsFileType.BLU_RAY)
                  {
                      tsiow.Write(header, 0, 4);
                      tsiow.Write(buffer[i].GetData(), 0, Constants.TS_SIZE);
                  }
                  else if (fileType == TsFileType.TS)
                  {
                      tsiow.Write(buffer[i].GetData(), 0, Constants.TS_SIZE);
                  }
              }
              packetCount += (UInt32)buffer.Count;
              buffer.Clear();
          }
  
          public override void Close()
          {
              if (fileType == TsFileType.BLU_RAY)
              {
                  if ((CurrentPacketNumber % 32) > 0)
                  {
                      for (UInt32 nullPacketNum = 32 - (CurrentPacketNumber % 32); nullPacketNum > 0; nullPacketNum--)
                      {
                          buffer.Add(new TsPacket());
                      }
                  }
              }
              if (buffer.Count > 0)
              {
                  Int64 stamp = currentPcr % Constants.MAX_BLURAY_CLOCK;
  
gt:                 for (int i = 0; i < buffer.Count; i++)