package org.bucksch.zeipis.recorder;
import org.bucksch.zeipis.all.Util;
import org.bucksch.zeipis.store.StreamCopy;
import java.io.*;
import java.util.*;
import java.util.Map.Entry;

/**
 * Takes an MPEG TS, and passes the output to several
 * InputStreams.
 * Each input stream can filter for certain PIDs (and they
 * can be different for each InputStream).
 * 
 * This allows to record several channels from the same
 * DVB Transport Stream.
 *
 * Implementation note:
 * I am not happy with the use of threads and streams here.
 */
public class MPEGTSSplit
{
  public void setSource(InputStream source)
  {
    assert(mySource == null);
    mySource = source;
  }

  /**
   * Set up a new target stream, where the data will be written to.
   * It will contain all parts of the source stream which match
   * the PIDs you pass in here (or in addPIDToTarget() later).
   *
   * The same PID may be used for several targets, i.e. you can
   * add two targets, one getting PIDs 0 and 256 and the other
   * getting PIDs 0 and 512.
   *
   * @param pids   which parts of the source stream to write here.
   * @return   where the data will be sent to.
   */
  public InputStream addTarget(Set<Integer> pids)
  {
    if (pids == null)
      pids = Util.newSet();

    InputOutputStream target;
    synchronized (myTargets)
    {
      target = new InputOutputStream(1024 * 500, new CloseCallback() {
        public void streamClosed(InputOutputStream s) {
          closeTarget(s);
        }});
      myTargets.put(target, pids);      
    }
    return target;
  }
  protected void closeTarget(InputOutputStream target)
  {
    synchronized (myTargets)
    {
      myTargets.remove(target);
      if (myTargets.isEmpty())
        myStopSignal = true; // TODO or close mySource directly?
    }
  }

  /**
   * Add a PID to a an already existing target.
   */
  public void addPIDToTarget(InputOutputStream target, int pid)
  {
    Set<Integer> existingPIDs = myTargets.get(target);
    assert(existingPIDs != null);
    existingPIDs.add(pid);
  }

  
  public void start()
  {
    myStopSignal = false;
    Thread t = new Thread("MPEG TS splitting")
    {
      @Override
      public void run()
      {
        try
        {
          copyInThread();
        } catch (Exception e) {
          Util.error(e); // XXX pass to consumer
          e.printStackTrace();
        }
      }
    };
    t.start();
  }
  protected void copyInThread()
    throws IOException
  {
    try {
      int count;
      byte[] buffer = new byte[1024];
      while ( ! myStopSignal)
      {
        count = mySource.read(buffer); // blocks until input data is available or EOF
        if (count == -1) // EOF
        {
          break;
        }

        // TODO read PID

        synchronized (myTargets)
        {
          for (Entry<InputOutputStream, Set<Integer>> entry : myTargets.entrySet())
          {
            InputOutputStream target = entry.getKey();
            target.write(buffer, 0, count);
          }          
        }
        Thread.yield();
      }
    } finally {
      synchronized (myTargets)
      {
        mySource.close();
        for (InputOutputStream target : myTargets.keySet())
          target.eof();
        myIsClosed = true;
      }
    }
  }
  public void stop()
  {
    myStopSignal = true;
    while ( ! myIsClosed) // TODO synchronized?
      Thread.yield();
  }
  public boolean isRunning()
  {
    // TODO here we could have a race: the time between myStopSignal and mySource.close()
    return ! myStopSignal;
  }

  
  // Test
  public static void main(String[] args)
  {
    try
    {
      MPEGTSSplit ts = new MPEGTSSplit();
      String sourceFilename = null;
      Map<Set<Integer>, OutputStream> outstreams = Util.newMap();
      for (int i = 0; i < args.length; i++)
      {
        if (args[i].equals("--source") &&
                ++i < args.length)
          sourceFilename = args[i];
        else if (args[i].equals("--out") &&
            ++i + 1 < args.length)
        {
          String targetFilename = args[i];
          String pidsString = args[++i];
          OutputStream out = new BufferedOutputStream(new FileOutputStream(targetFilename));
          String[] pidsStrings = pidsString.split(",");
          Set<Integer> pids = Util.newSet();
          for (String pidString : pidsStrings)
            pids.add(Integer.parseInt(pidString));
          outstreams.put(pids, out);
        }
      }
      if (sourceFilename == null || outstreams.isEmpty())
        throw new Exception("Need --source <filename> --out <filename> <pid>,<pid>,... --out ...");

      ts.setSource(new BufferedInputStream(new FileInputStream(sourceFilename)));
      //Map<Set<Integer>, InputStream> instreams = Util.newMap();
      //Set<StreamCopy> threads = Util.newSet();
      for (Entry<Set<Integer>, OutputStream> entry : outstreams.entrySet())
      {
        //instreams.put(entry.getKey(), ts.addTarget(entry.getKey()));
        InputStream in = ts.addTarget(entry.getKey());
        StreamCopy thread = new StreamCopy();
        thread.copy(in, entry.getValue(), false, "Copy thread");
      }
      
      ts.start();
    } catch (Exception e) { Util.error(e); }
  }


  // Data 
  protected InputStream mySource = null;
  protected Map<InputOutputStream, Set<Integer>> myTargets = Util.newMap(); // target stream / Set of PIDs

  protected boolean myStopSignal = false;
  protected boolean myIsClosed = false;

}
