#pragma once
#include "control_point.hpp"
#include "opcode_utility.hpp"
#include "instruction_map.hpp"
#include "address_list.hpp"

#include <vector>
#include <iostream>
#include <algorithm>

#include "chocolat/string/string_aligner.hpp"

#include <boost/shared_ptr.hpp>
#include <boost/optional.hpp>
/** This is range of memory address, and node of control flow graph.*/
class BasicBlock
{
	BasicBlock()
			:m_is_valid(false)
	{
	}
	
public:
	typedef boost::shared_ptr<BasicBlock> shared_ptr;
private:
	unsigned m_start_address;
	unsigned m_end_address;
	AddressList::shared_ptr m_parent_routine;
	AddressList::shared_ptr m_predessors;
	AddressList::shared_ptr m_successors;
	bool m_is_valid;
	bool m_is_entry;
public:
	/** get invalid object.*/
	static shared_ptr getInvalid()
	{
		static shared_ptr invalid(new BasicBlock);
		return invalid;
	}
	static shared_ptr create(unsigned start_address)
	{
		BasicBlock::shared_ptr block(new BasicBlock(start_address));
		return block;
	}
	bool isValid()const
	{
		return m_is_valid;
	}
	/** if not set the end address, fake block size is undefined.*/
	BasicBlock(unsigned start_address)
			:m_start_address(start_address),
			 m_end_address(start_address),
			 m_is_valid(true),
			 m_is_entry(false)
	{
		m_parent_routine = AddressList::create();
		m_predessors = AddressList::create();
		m_successors = AddressList::create();
	}
	AddressList::shared_ptr predessorList()
	{
		return m_predessors;
	}
	AddressList::shared_ptr successorList()
	{
		return m_successors;
	}
	/* Self start address is in successors. */
	bool isSelfLoop()const
	{
		for (auto i=m_successors->begin(); i!=m_successors->end(); ++i) {
			if (*i == m_start_address) {
				return true;
			}
		}
		return false;
	}
	unsigned startAddress()const
	{
		return m_start_address;
	}
	unsigned endAddress()const
	{
		return m_end_address;
	}
	unsigned size()const
	{
		return m_end_address - m_start_address;
	}
	void setEnd(unsigned end_address) 
	{
		m_end_address = end_address;
	}
	void setEntryPoint()
	{
		m_is_entry = true;
	}
	bool isEntryPoint()const
	{
		return m_is_entry; 
	}
	void linkTo(BasicBlock::shared_ptr successor)
	{
		m_successors->push_back(successor->m_start_address);
		successor->m_predessors->push_back(m_start_address);
	}
	void display(std::ostream& os)const
	{
		AddressListManipulater succ(m_successors);
		AddressListManipulater pred(m_predessors);
		
		StringAligner left(20, StringAligner::LEFT);
		StringAligner right(20, StringAligner::RIGHT);
		os << right(pred.to_s())
		   << "-> " << rangeString()
		   << (m_is_entry? "*" : " ")
		   << "-> "
		   << left(succ.to_s());
	}
	void straighten()
	{
		AddressList::straighten(m_successors);
		AddressList::straighten(m_predessors);
	}
	std::string rangeString()const
	{
		return HexadecimalFormat(m_start_address).to_s()
				+ " (" + HexadecimalFormat(size()).width(4).to_s()+")";
	}
};

/** map of basic block.*/
class BasicBlockList
{
public:
	typedef std::map<unsigned, BasicBlock::shared_ptr> map;
	typedef boost::shared_ptr<map> shared_ptr;
	typedef map::iterator iterator;
	
	BasicBlockList::shared_ptr m_map;
	static shared_ptr create()
	{
		return shared_ptr(new map);
	}
	
	BasicBlockList()
			:m_map(new BasicBlockList::map)
	{
	}
	BasicBlockList(BasicBlockList::shared_ptr given_map)
			:m_map(given_map)
	{
	}
	/** if not exists, then create.*/
	BasicBlock::shared_ptr open(unsigned start_address)
	{
		if (m_map->find(start_address) == m_map->end()) {
			BasicBlock::shared_ptr block = BasicBlock::create(start_address);
			(*m_map)[start_address] = block;
		}
		return (*m_map)[start_address];
	}
	/* this doesn't create*/
	BasicBlock::shared_ptr get(unsigned start_address)
	{
		if (m_map->find(start_address) == m_map->end()) {
			return BasicBlock::getInvalid();
		}
		return (*m_map)[start_address];
		
	}
	iterator begin()const
	{
		return m_map->begin();
	}
	iterator end()const
	{
		return m_map->end();
	}
	void displayBriefResult(std::ostream& os)const
	{
		os << m_map->size() << " blocks found." << std::endl;
	}
	BasicBlockList::shared_ptr getList()
	{
		return m_map;
	}
};

class BasicBlockAnalyzer
{
	InstructionMap::shared_ptr m_map;
	ControlPointList::shared_ptr m_control_list;
	BasicBlockList m_block_list;
	AddressMarker m_visits;
public:
	BasicBlockAnalyzer(InstructionMap::shared_ptr map,
					   ControlPointList::shared_ptr control_list)
			:m_map(map),
			 m_control_list(control_list)
	{
	}
	/**/
	void analyze()
	{
		AddressList::shared_ptr entries = m_control_list->getEntryList();
		for (auto i=entries->begin(); i!=entries->end();++i) {
			createBasicBlocks(*i);
		}
		for (auto i=m_block_list.begin(); i!=m_block_list.end(); ++i) {
			i->second->straighten();
		}
	}
	BasicBlockList::shared_ptr getBlockList()
	{
		return m_block_list.getList();
	}
	void displayBriefResult(std::ostream& os)const
	{
		m_block_list.displayBriefResult(os);
	}
	void displayBlock(std::ostream& os)const
	{
		for (auto i=m_block_list.begin(); i!=m_block_list.end(); ++i) {
			i->second->display(os);
			os << std::endl;
		}
	}
	Instruction::shared_ptr getInstructionAt(unsigned address)const
	{
		if (m_map->find(address) != m_map->end()){
			return (*m_map)[address];
		} else {
			return Instruction::getInvalid();
		}
	}
	struct isNotCall
	{
		BasicBlockAnalyzer* that;
		isNotCall(BasicBlockAnalyzer* analyzer)
				:that(analyzer)
		{
		}
		bool operator()(unsigned address)const
		{
			
			opcode::OperationType optype(that->getInstructionAt(address)->op());			
			return !optype.is_call;
		}
	};
	Instruction::shared_ptr getNextLabel(unsigned current)
	{
		AddressListManipulater label_list(m_control_list->getLabelList());
		unsigned label_address = label_list.getNextFrom(current);
		return getInstructionAt(label_address);
	}
	Instruction::shared_ptr getNextBranch(unsigned current)
	{
		AddressListManipulater branch_list(m_control_list->getBranchList());
		isNotCall is_not_call(this);
		/* if current is branch, it must take in account.*/
		unsigned branch_address = branch_list.getNextIf(current-1,
														is_not_call);
		return getInstructionAt(branch_address);
	}
	void createBasicBlocks(unsigned entry_address) 
	{
		AddressMarker visits;
		AddressWorkList work_list;
		AddressListManipulater label_list(m_control_list->getLabelList());
		AddressListManipulater branch_list(m_control_list->getBranchList());
		
		work_list.push(entry_address);
		m_block_list.open(entry_address)->setEntryPoint();
		
		boost::optional<unsigned> current = boost::none;
		while (true) {
			if (!current || visits.found(*current) ) {
				if (!work_list.empty()) {
					current = work_list.pop();
				} else {
					break;
				}
			}
			visits.mark(*current);
			
			auto block = m_block_list.open(*current);
			if (m_visits.found(*current)) {
				continue;
			}
			m_visits.mark(*current);
			
			Instruction::shared_ptr label = getNextLabel(*current); 
			Instruction::shared_ptr branch = getNextBranch(*current);
			/*
			  std::cout << HexadecimalFormat(*current) << " "
			  << HexadecimalFormat(label->address) << " "
			  << HexadecimalFormat(branch->next) << " " << std::endl;
			//*/
			if (!label->isValid() && !branch->isValid()) {
				unsigned next_address = *current+1;
				for (auto i = m_map->find(*current); i!=m_map->end();++i){
					next_address = i->second->next();
				}
				block->setEnd(next_address);
				current = boost::none;
			} else if (!branch->isValid()
					   || (label->address() < branch->next()))
			{
				block->setEnd(label->address());
				block->linkTo(m_block_list.open(label->address()));
				current = label->address();
			} else {
				opcode::OperationType branch_type;
				branch_type.set(branch->op());
				
				if (branch_type.is_return) {
					block->setEnd(branch->next());
					current = boost::none;
				} else if (branch_type.is_conditional) {
					block->setEnd(branch->next());
					block->linkTo(m_block_list.open(branch->next()));
					work_list.push(branch->parameter());
					block->linkTo(m_block_list.open(branch->parameter()));
					current = branch->next();
				} else {
					block->setEnd(branch->next());
					block->linkTo(m_block_list.open(branch->parameter()));
					current = branch->parameter();
				}
			}
		}
	}
};
/** container adaptor for map, check in and check out the basic block.*/
class BasicBlockBox
{
	std::map<unsigned, BasicBlock::shared_ptr> m_map;
public:
	BasicBlockBox()
	{
	}
	void push(BasicBlock::shared_ptr block)
	{
		m_map[block->startAddress()] = block;
	}
	BasicBlock::shared_ptr pop()
	{
		auto pos = m_map.begin();
		m_map.erase(pos);
		return pos->second;
	}
	BasicBlock::shared_ptr checkOut(unsigned address)
	{
		auto pos = m_map.find(address);
		if (pos != m_map.end()){
			m_map.erase(pos);
			return pos->second;
		} else {
			return BasicBlock::getInvalid();
		}
	}
	bool empty()const
	{
		return m_map.empty();
	}
};
