#include "forensics-gecode.hh"
#include "error.hh"
#include "util.hh"
#include "filewrapper.hh"
#include "imagedatasource.hh"
#include "shiftdatasource.hh"
#include "interdatasource.hh"
#include "blockdatasource.hh"
#include "minsearch.hh"
#include "patternminsearcher.hh"
#include "execpropagator.hh"
#include "extract.hh"
#include "ntfsbitmap.hh"
#include "ntfsallocation.hh"
#include "ntfsrunlistpropagator.hh"


using namespace Gecode;

/* ************************************************************************************************** */

DataSource& getFileDS (Space& home);
DataSource& getInterDS (Space& home);
BlockDataSource& getBlocksDS (Space& home);

/* ************************************************************************************************** */

enum SearchArea { 
  SEARCH_ALL,
  SEARCH_ALLOCATED,
  SEARCH_UNALLOCATED,
};

/* ************************************************************************************************** */

const int max_file_length_to_search_for = 1 * 1024 * 1024;

class MySpace : public Space
{
public:
  const int num_clusters;
  filewrapper& fw;
  ImageDataSource fileDS;
  ShiftDataSource shiftDS;
  InterDataSource interDS;
  BlockDataSource blocksDS;
  IntVar fileLength;
  IntVar clusterStart;
  IntVar offsetStart;
  IntVar clusterEnd;
  IntVar offsetEnd;

  MySpace (filewrapper& fw, const SearchArea area) : 
    Space (),
    num_clusters ((int) ((fw.size () / BLOCK_SIZE) - 63 - 1)),
    fw (fw),
    fileDS (*this, fw),
    shiftDS (*this, fileDS, 63),
    interDS (*this, shiftDS, IntSet (0, Int::Limits::max)),
    blocksDS (*this, interDS, (max_file_length_to_search_for / BLOCK_SIZE) + 1, 0, num_clusters), // +1 needed when offsetStart need not be zero.
    fileLength (*this, 0, max_file_length_to_search_for),
    clusterStart (*this, 0, num_clusters),
    offsetStart (*this, 0, BLOCK_SIZE - 1),
    clusterEnd (*this, 0, num_clusters),
    offsetEnd (*this, 0, BLOCK_SIZE - 1)
  {
    unsigned char start_pattern_array[] = { 0xFF, 0xD8, 0xFF, 0xE0, 0x00, 0x10 };
    std::list<unsigned char> start_pattern (&start_pattern_array[0], &start_pattern_array[6]);
    IntVar clusterStartE (*this, 0, num_clusters);
    IntVar offsetStartE (*this, 0, BLOCK_SIZE - 1);
    pattern_min_search (*this, clusterStart, offsetStart, clusterStartE, offsetStartE, start_pattern, getInterDS);

    unsigned char end_pattern_array[] = { 0xFF, 0xD9 };
    std::list<unsigned char> end_pattern (&end_pattern_array[0], &end_pattern_array[2]);
    IntVar clusterEndE (*this, 0, num_clusters);
    IntVar offsetEndE (*this, 0, BLOCK_SIZE - 1);
    pattern_min_search (*this, clusterEnd, offsetEnd, clusterEndE, offsetEndE, end_pattern, getInterDS);

    IntVarArray& blocks = blocksDS.blocks ();
    const int num_blocks = blocks.size ();
    // Caution: if offsetStart need not be 0, then we may use one more block.
    for (int i = 1; i < num_blocks; i++) {
      rel (*this, (((offsetStart + fileLength)) > i * BLOCK_SIZE) >> (blocks[i-1] + 1 == blocks[i]));
      rel (*this, (((offsetStart + fileLength)) <= i * BLOCK_SIZE) >> (0 == blocks[i]));
    }

    if (area == SEARCH_ALL) {
      // No further constraint.
    } else if (area == SEARCH_ALLOCATED) {
      NTFSBitmapExtractor* bitmapExtractor = new NTFSBitmapExtractor (interDS);
      ntfs_allocated (*this, clusterStart, bitmapExtractor);
      ntfs_allocated (*this, clusterEnd, bitmapExtractor);
    } else if (area == SEARCH_UNALLOCATED) {
      NTFSBitmapExtractor* bitmapExtractor = new NTFSBitmapExtractor (interDS);
      ntfs_unallocated (*this, clusterStart, bitmapExtractor);
      ntfs_unallocated (*this, clusterEnd, bitmapExtractor);
    }
    
    rel (*this, clusterStart == blocks[0]);
    rel (*this, clusterStart <= clusterEnd);
    rel (*this, clusterEnd - clusterStart <= (max_file_length_to_search_for / BLOCK_SIZE) + 3);
    rel (*this, fileLength == (clusterEnd - clusterStart) * BLOCK_SIZE + (offsetEnd - offsetStart) + end_pattern.size ());
    // rel (*this, offsetStart == 0);


    std::string executable ("./is-valid-jpeg 2>/dev/null >/dev/null");
    exec_propagator (*this, offsetStart, fileLength, executable, getBlocksDS);

   
    IntVarArgs l (5);
    l[0] = clusterStart;
    l[1] = offsetStart;
    l[2] = clusterEnd;
    l[3] = offsetEnd;
    l[4] = fileLength;
    
    branch (*this, l, INT_VAR_NONE, INT_VAL_SPLIT_MIN);
  }


  MySpace (bool share, MySpace& s) :
    Space (share, s),
    num_clusters (s.num_clusters),
    fw (s.fw),
    fileDS (share, *this, s.fileDS),
    shiftDS (share, *this, fileDS, s.shiftDS),
    interDS (share, *this, shiftDS, s.interDS),
    blocksDS (share, *this, interDS, s.blocksDS)
  {
    fileLength.update (*this, share, s.fileLength);
    clusterStart.update (*this, share, s.clusterStart);
    offsetStart.update (*this, share, s.offsetStart);
    clusterEnd.update (*this, share, s.clusterEnd);
    offsetEnd.update (*this, share, s.offsetEnd);
  }


  virtual Space* copy (bool share) 
  {
    MySpace* result = new MySpace (share, *this);
    return result;
  }


  void print (std::ostream& os) const
  {
    os << "fileDS = ";
    fileDS.print (os) ;
    os << std::endl;
    os << "clusterStart = " << clusterStart 
       << ", offsetStart = " << offsetStart 
       << ", clusterEnd = " << clusterEnd 
       << ", offsetEnd = " << offsetEnd 
       << ", fileLength = " << fileLength;
    //    << ", blocksDS = ";
    // blocksDS.print (os);
    os << std::endl;
  }

  void print (void) const
  {
    print (std::cout);
  }

};

/* ************************************************************************************************** */

DataSource& getFileDS (Space& home)
{
  MySpace& myhome = static_cast<MySpace&> (home);
  DataSource& result = myhome.fileDS;
  return result;
}

/* ************************************************************************************************** */

DataSource& getInterDS (Space& home)
{
  MySpace& myhome = static_cast<MySpace&> (home);
  DataSource& result = myhome.interDS;
  return result;
}

/* ************************************************************************************************** */

BlockDataSource& getBlocksDS (Space& home)
{
  MySpace& myhome = static_cast<MySpace&> (home);
  BlockDataSource& result = myhome.blocksDS;
  return result;
}

/* ************************************************************************************************** */

int main (int argc, char** argv) 
{
  try {

    if (argc != 3) error ("wrong number of arguments");

    filewrapper fw (argv[1]);

    SearchArea area = SEARCH_ALL;
    if (!strcmp (argv[2], "all")) {
      area = SEARCH_ALL;
    } else if (!strcmp (argv[2], "allocated")) {
      area = SEARCH_ALLOCATED;
    } else if (!strcmp (argv[2], "unallocated")) {
      area = SEARCH_UNALLOCATED;
    } else {
      error ("%s should be \"all\", \"allocated\", or \"unallocated\"");
    }
    
    MySpace* m = new MySpace (fw, area);

    std::cout << std::endl << "Results of trivial propagation" << std::endl << std::endl;
    m->print ();
    std::cout << std::endl;

    std::cout << "Results of main propagation" << std::endl << std::endl;
    m->status ();
    m->print ();
    std::cout << std::endl;

    std::cout << "Create search engine" << std::endl;
    std::cout << std::endl << "===================================================================" << std::endl << std::endl;

    DFS<MySpace> e (m);
    int num_solutions = 0;
    while (MySpace* s = e.next ()) {

      int64_t start = s->clusterStart.val ();
      start *= BLOCK_SIZE;
      start += s->offsetStart.val ();
      int64_t end = s->clusterEnd.val ();
      end *= BLOCK_SIZE;
      end += s->offsetEnd.val ();
      int length = s->fileLength.val ();

      std::ostringstream oss (std::ostringstream::out);
      oss << "carving/output-";
      oss.fill ('0'); oss.width (12); oss << start;
      oss << "-";
      oss.fill ('0'); oss.width (12); oss << end;
      oss << ".jpg";
      std::string file_name_string = oss.str ();

      Extract extract (s->blocksDS, file_name_string);
      extract.write_file (s->offsetStart.val (), length);
      // Extract extract (s->fileDS, file_name_string);
      // extract.write_file (start, length);

      print_statistics (std::cout, e.statistics ());
      s->print ();
      delete s;
      num_solutions++;

      std::cout << std::endl << "===================================================================" << std::endl << std::endl;
    }

    std::cout << std::endl;

    print_statistics (std::cout, e.statistics ());
    std::cout << "number of solutions = " << num_solutions << std::endl;

    delete m;

    std::cout << "cache hits = " << cache_hits << "; cache_misses = " << cache_misses << std::endl;
    std::cout << "pointless_finds = " << pointless_finds << std::endl;
    std::cout << "num_bytes_written_to_exec = " << num_bytes_written_to_exec << std::endl;

  } catch (Exception e) {
    std::cerr << "Gecode exception: " << e.what () << std::endl;
    return 1;
  }

  return 0;
}

/* ************************************************************************************************** */
