#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);

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

enum SearchArea { 
  SEARCH_ALL,
  SEARCH_ALLOCATED,
  SEARCH_UNALLOCATED,
};

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

class MySpace : public Space
{
public:
  const int num_clusters;
  filewrapper& fw;
  ImageDataSource fileDS;
  ShiftDataSource shiftDS;
  InterDataSource interDS;
  IntVar clusterFILE;
  IntVar offsetFILE;

  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)),
    clusterFILE (*this, 0, num_clusters),
    offsetFILE (*this, 0, BLOCK_SIZE - 1)
  {
    // Search for "FILE" on cluster boundaries.
    unsigned char file_pattern_array[] = "FILE";
    std::list<unsigned char> file_pattern (&file_pattern_array[0], &file_pattern_array[4]);
    rel (*this, offsetFILE == 0);

    IntVar clusterFILEE (*this, 0, num_clusters);
    IntVar offsetFILEE (*this, 0, BLOCK_SIZE - 1);
    pattern_min_search (*this, clusterFILE, offsetFILE, clusterFILEE, offsetFILEE, file_pattern, getInterDS);

    if (area == SEARCH_ALL) {
      // No further constraint.
    } else if (area == SEARCH_ALLOCATED) {
      NTFSBitmapExtractor* bitmapExtractor = new NTFSBitmapExtractor (interDS);
      ntfs_allocated (*this, clusterFILE, bitmapExtractor);
    } else if (area == SEARCH_UNALLOCATED) {
      NTFSBitmapExtractor* bitmapExtractor = new NTFSBitmapExtractor (interDS);
      ntfs_unallocated (*this, clusterFILE, bitmapExtractor);
    }
    
    IntVarArgs l (2);
    l[0] = clusterFILE;
    l[1] = offsetFILE;
    
    branch (*this, l, INT_VAR_NONE, INT_VAL_SPLIT_MIN);
    // branch (*this, l, INT_VAR_NONE, INT_VAL_RANGE_MIN);
    // branch (*this, l, INT_VAR_NONE, INT_VAL_MIN);
    // branch (*this, offset1, INT_VAL_SPLIT_MIN);
    // branch (*this, offset1, INT_VAL_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)
  {
    clusterFILE.update (*this, share, s.clusterFILE);
    offsetFILE.update (*this, share, s.offsetFILE);
  }


  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;
    double percent = 0;
    if (clusterFILE.assigned()) {
      percent = 100*clusterFILE.val ();
    }
    percent /= num_clusters;
    os << "clusterFILE = " << clusterFILE << " out of " << num_clusters << " [ " << percent << "% ]"
       << ", offsetFILE = " << offsetFILE;
    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;
}

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

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 ()) {

      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;
}

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

