#include <unordered_map>
#include <utility>
#include <cstdlib>
#include <cassert>
#include <cstring>

#include "world.h"

using namespace std;


Argument::~Argument(){}


Instruction Instruction::make(ins_t op,vector<Argument> args){
	return {op,args};
}


void Robot::load(int idx,const Script &scr){
	banks[idx]=scr;
}

WorldAction Robot::tick(World &world){
	if(active==0)return WorldAction::none;  // TODO: check whether `waited` needs to be reset
	if(!ipInRange()){
		ip.bank=ip.pos=0;
		if(!ipInRange())return WorldAction::die;
	}
	waited++;
	Instruction ins=resolve(world);
	int dur=calcDuration(ins);
	if(waited<dur)return WorldAction::none;
	waited=0;
	return execute(ins,world);
}

uint16_t* Robot::resolveVar(const Argument &arg,World &world){
	// cout<<"resolveVar("<<(int)arg.type<<","<<arg.remote<<","<<arg.num<<","<<arg.name<<")"<<endl;
	if(arg.remote){
		assert(arg.type==arg_t::name);
		assert(arg.name=="active");
		Robot *target=world.targetbot(this);
		if(target)return &target->active;
		else return nullptr;
	} else {
		if(arg.type==arg_t::var){
			return &vars.at(arg.num-1);
		} else if(arg.type==arg_t::name){
			assert(arg.name=="active");
			return &active;
		} else assert(false);
	}
}

Instruction Robot::resolve(World &world){
	Instruction ins=banks.at(ip.bank).at(ip.pos);
	// cout<<this<<" Resolving "<<ins<<endl;
	vector<bool> dstvar(ins.args.size(),false);
	if(ins.op==ins_t::set||ins.op==ins_t::add||ins.op==ins_t::sub||ins.op==ins_t::scan){
		assert(dstvar.size()>0);
		dstvar[0]=true;
	}
	for(size_t i=0;i<ins.args.size();i++){
		Argument &arg=ins.args[i];
		if(dstvar[i]){
			assert(arg.type==arg_t::var||arg.type==arg_t::name);
			continue;
		}
		switch(arg.type){
			case arg_t::number:
				break;
			case arg_t::var:{
				uint16_t *v=resolveVar(arg,world);
				if(v)arg.num=*v;
				else arg.num=0;
				arg.type=arg_t::number;
				break;
			}
			case arg_t::name:{
				assert(arg.name=="active");
				uint16_t *v=resolveVar(arg,world);
				if(v)arg.num=*v;
				else arg.num=0;
				arg.type=arg_t::number;
				break;
			}
			case arg_t::constant:
				assert(!arg.remote);
				if(arg.name=="fields"){
					arg.type=arg_t::number;
					arg.num=SIZE;
				} else assert(false);
				break;

			default:
				assert(false);
		}
	}
	return ins;
}

int Robot::calcDuration(Instruction &ins){
	int d=0;
	for(const Argument &arg : ins.args){
		if(arg.remote)d+=C::pen_remote;
	}
	int base=C::baseDuration[(int)ins.op];
	switch(ins.op){
		case ins_t::set: d+=base; break;
		case ins_t::add: d+=base; break;
		case ins_t::sub: d+=base; break;
		case ins_t::comp: d+=base; break;
		case ins_t::trans: d+=base+C::pen_transinstr*banks.at(ins.args[0].num-1).size(); break;
		case ins_t::jump: d+=base; break;
		case ins_t::bjump: d+=base; break;
		case ins_t::die: d+=base; break;
		case ins_t::move: d+=base; break;
		case ins_t::turn: d+=base; break;
		case ins_t::scan: d+=base; break;
		case ins_t::create: d+=(base+C::pen_createbank*ins.args[1].num)*C::pen_createmobilemult+C::pen_createmobile*ins.args[2].num+C::pen_createiset1*(ins.args[0].num==1)+C::pen_createiset2*(ins.args[0].num==2); break;
		case ins_t::debugger: d+=base; break;
		default: assert(false);
	}
	return d;
}

WorldAction Robot::execute(Instruction &ins,World &world){
	// cout<<this<<" Executing "<<ins<<endl;
	WorldAction action=WorldAction::none;
	switch(ins.op){
		case ins_t::set:{
			uint16_t *v=resolveVar(ins.args[0],world);
			if(v)*v=ins.args[1].num;
			break;
		}

		case ins_t::add:{
			uint16_t *v=resolveVar(ins.args[0],world);
			if(v)*v+=ins.args[1].num;
			break;
		}

		case ins_t::sub:{
			uint16_t *v=resolveVar(ins.args[0],world);
			if(v)*v-=ins.args[1].num;
			break;
		}

		case ins_t::comp:
			if(ins.args[0].num==ins.args[1].num){
				advanceIP();
			}
			break;

		case ins_t::trans:{
			assert(iset>=1);
			Robot *r=world.targetbot(this);
			if(r){
				assert(ins.args[0].num>=1&&ins.args[0].num<=(int)banks.size()&&
				       ins.args[1].num>=1&&ins.args[1].num<=(int)r->banks.size());
				r->banks[ins.args[1].num-1]=banks[ins.args[0].num-1];
			}
			break;
		}

		case ins_t::jump:
			ip.pos+=ins.args[0].num-1;
			break;

		case ins_t::bjump:
			ip.bank=ins.args[0].num-1;
			ip.pos=ins.args[1].num-1;
			break;

		case ins_t::die:
			return WorldAction::die;

		case ins_t::move:
			if(!mobile)return WorldAction::die;
			if(world.targetbot(this)==nullptr){
				action=WorldAction::move;
			}
			break;

		case ins_t::turn:
			if(ins.args[0].num==0)heading=(heading+3)%4;
			else heading=(heading+1)%4;
			break;

		case ins_t::scan:{
			assert(iset>=1);
			Robot *r=world.targetbot(this);
			uint16_t *v=resolveVar(ins.args[0],world);
			if(!v)break;
			if(r){
				if(r->team==team)*v=2;
				else *v=1;
			} else {
				*v=0;
			}
			break;
		}

		case ins_t::create:
			assert(iset>=2);
			if(world.targetbot(this)!=nullptr)break;
			assert(ins.args[0].num>=0&&ins.args[0].num<=2);
			assert(ins.args[1].num>=0);
			assert((bool)ins.args[2].num==ins.args[2].num);
			world.createInFront(this,this->team,ins.args[0].num,ins.args[1].num,ins.args[2].num);
			break;

		case ins_t::debugger:
			__asm("int3\n\t");
			break;

		default:
			assert(false);
	}

	advanceIP();
	return action;
}

void Robot::advanceIP(){
	ip.pos++;
	if(!ipInRange()){
		ip.bank=ip.pos=0;
		assert(ipInRange());
	}
}

bool Robot::ipInRange() const {
	if(ip.bank<0||ip.bank>=(int)banks.size())return false;
	if(ip.pos<0||ip.pos>=(int)banks[ip.bank].size())return false;
	return true;
}


World::World(){
	memset(&board[0][0],0,SIZE*SIZE*sizeof(Robot*));
}

World::~World(){
	for(int y=0;y<SIZE;y++){
		for(int x=0;x<SIZE;x++){
			if(board[y][x])delete board[y][x];
		}
	}
}

Robot& World::create(const Team *team,int iset,int nbanks,bool mobile){
	int x,y,heading;
	while(true){
		x=rand()%SIZE;
		y=rand()%SIZE;
		if(!board[y][x])break;
	}
	heading = rand()%4;

	return create(team, iset, nbanks, mobile, x, y, heading);
}

Robot& World::create(const Team *team,int iset,int nbanks,bool mobile,int x,int y,int heading){
	x = (x%SIZE + SIZE)%SIZE; // Negative values are wrapped
	y = (y%SIZE + SIZE)%SIZE;
	if(board[y][x])return *board[y][x];
	board[y][x]=new Robot();
	board[y][x]->team=team;
	board[y][x]->banks.resize(nbanks);
	board[y][x]->iset=iset;
	board[y][x]->mobile=mobile;
	board[y][x]->heading=heading%4;
	return *board[y][x];
}

Robot& World::createInFront(const Robot *caller,const Team *team,int iset,int nbanks,bool mobile){
	Robot **ptr=targetbotptr(caller);
	Robot *r=new Robot();
	r->team=team;
	r->banks.resize(nbanks);
	r->iset=iset;
	r->mobile=mobile;
	r->heading=caller->heading;
	*ptr=r;
	return *r;
}

void World::removeRobot(const Robot *r){
	*botptr(r)=nullptr;
}

void World::removeTeam(const Team *team){
	for(int y=0;y<SIZE;y++){
		for(int x=0;x<SIZE;x++){
			if(board[y][x]&&board[y][x]->team==team){
				board[y][x]=nullptr;
			}
		}
	}
}

void World::tick(){
	vector<pair<int,int>> candidates;
	candidates.reserve(SIZE*SIZE);

	for(int y=0;y<SIZE;y++){
		for(int x=0;x<SIZE;x++){
			if(board[y][x]&&board[y][x]->active){
				candidates.emplace_back(x,y);
			}
		}
	}

	for(const pair<int,int> &p : candidates){
		int x=p.first,y=p.second;
		if(board[y][x]->active==0)continue;
		switch(board[y][x]->tick(*this)){
			case WorldAction::none: break;

			case WorldAction::die:
				delete board[y][x];
				board[y][x]=nullptr;
				break;

			case WorldAction::move:{
				Robot** ptr=nullptr;
				switch(board[y][x]->heading%4){
					case 0: ptr=&board[(y+SIZE-1)%SIZE][x]; break;
					case 1: ptr=&board[y][(x+1)%SIZE]; break;
					case 2: ptr=&board[(y+1)%SIZE][x]; break;
					case 3: ptr=&board[y][(x+SIZE-1)%SIZE]; break;
				}
				*ptr=board[y][x];
				board[y][x]=nullptr;
				break;
			}
		}
	}
}

Robot** World::botptr(const Robot *r){
	for(int y=0;y<SIZE;y++){
		for(int x=0;x<SIZE;x++){
			if(board[y][x]==r)return &board[y][x];
		}
	}
	assert(false);
}

Robot** World::targetbotptr(const Robot *r){
	for(int y=0;y<SIZE;y++){
		for(int x=0;x<SIZE;x++){
			if(board[y][x]!=r)continue;
			switch(r->heading%4){
				case 0: return &board[(y+SIZE-1)%SIZE][x];
				case 1: return &board[y][(x+1)%SIZE];
				case 2: return &board[(y+1)%SIZE][x];
				case 3: return &board[y][(x+SIZE-1)%SIZE];
			}
		}
	}
	assert(false);
}

Robot* World::targetbot(const Robot *r){
	return *targetbotptr(r);
}

static int colourCode(const Team *team){
	static unordered_map<const Team*,int> mappings;
	static int nextCode=1;
	auto it=mappings.find(team);
	if(it!=mappings.end())return it->second;
	int code=nextCode;
	mappings[team]=code;
	nextCode=nextCode%6+1;
	return code;
}

void World::print(ScreenBuffer &sb) const {
	for(int y=0;y<SIZE;y++){
		sb.moveto(0,y);
		for(int x=0;x<SIZE;x++){
			if(board[y][x]==nullptr)sb.printf("..");
			else {
				char c2="^>v<"[board[y][x]->heading%4];
				char c1=board[y][x]->active?c2:'X';
				sb.setfg(colourCode(board[y][x]->team));
				sb.printf("%c%c",c1,c2);
				sb.setfg(9);
			}
			sb.printf(" ");
		}
	}
}


ostream& operator<<(ostream &os,const Argument &arg){
	if(arg.remote){
		switch(arg.type){
			case arg_t::number: return os<<arg.num;
			case arg_t::var: assert(false);
			case arg_t::name: return os<<'%'<<arg.name;
			case arg_t::constant: assert(false);
			case arg_t::undeflabel: return os<<"\x1B[31;43mUNDEFLABEL\x1B[0m";
		}
	} else {
		switch(arg.type){
			case arg_t::number: return os<<arg.num;
			case arg_t::var: return os<<'#'<<arg.num;
			case arg_t::name: return os<<'#'<<arg.name;
			case arg_t::constant: return os<<'$'<<arg.name;
			case arg_t::undeflabel: return os<<"\x1B[31;43mUNDEFLABEL\x1B[0m";
		}
	}
}

ostream& operator<<(ostream &os,const ins_t &type){
	switch(type){
		case ins_t::set: return os<<"set";
		case ins_t::add: return os<<"add";
		case ins_t::sub: return os<<"sub";
		case ins_t::comp: return os<<"comp";
		case ins_t::trans: return os<<"trans";
		case ins_t::jump: return os<<"jump";
		case ins_t::bjump: return os<<"bjump";
		case ins_t::die: return os<<"die";
		case ins_t::move: return os<<"move";
		case ins_t::turn: return os<<"turn";
		case ins_t::scan: return os<<"scan";
		case ins_t::create: return os<<"create";
		default: assert(false);
	}
}

ostream& operator<<(ostream &os,const Instruction &ins){
	os<<ins.op;
	bool first=true;
	for(const Argument &arg : ins.args){
		if(first)os<<' ';
		else os<<", ";
		first=false;
		os<<arg;
	}
	return os;
}

ostream& operator<<(ostream &os,const Team &team){
	os<<"TEAM: "<<team.name<<endl;
	int bank=0;
	for(const Script &scr : team.banks){
		os<<"Bank "<<bank<<endl;
		for(const Instruction &ins : scr){
			os<<'\t'<<ins<<endl;
		}
		bank++;
	}
	return os;
}