#include <fstream>
#include <vector>
#include <string>
#include <unordered_map>
#include <cctype>
#include <cassert>

#include "parse.h"

void trim(string &str){
	while(str.size()>0&&isspace(str[0]))str.erase(0,1);
	while(str.size()>0&&isspace(str.back()))str.erase(str.size()-1,1);
}

void squashSpaces(string &str){
	bool sp=false;
	for(size_t i=0;i<str.size();i++){
		if(isspace(str[i])){
			if(sp){
				str.erase(i,1);
				i--;
			} else {
				str[i]=' ';
				sp=true;
			}
		} else {
			sp=false;
		}
	}
}

vector<string> splittrim(const string &str,char sep){
	if(str.size()==0)return {};
	vector<string> res;
	size_t cursor=0;
	while(true){
		size_t idx=str.find(sep,cursor);
		if(idx==string::npos){
			string s=str.substr(cursor);
			trim(s);
			res.push_back(s);
			return res;
		}
		string s=str.substr(cursor,idx-cursor);
		trim(s);
		res.push_back(s);
		cursor=idx+1;
	}
}

void lowercase(string &str){
	for(size_t i=0;i<str.size();i++){
		str[i]=tolower(str[i]);
	}
}

Argument parseArgument(const string &str,const unordered_map<string,Location> &labels,Location ip){
	if(str[0]=='#'||str[0]=='%'){
		bool remote=str[0]=='%';
		if(isdigit(str[1])){
			int num=atoi(str.data()+1);
			assert(num>=1&&num<=20);
			return {arg_t::var,remote,num,{}};
		} else {
			return {arg_t::name,remote,0,str.substr(1)};
		}
	} else if(str[0]=='@'){
		if(labels.find(str.substr(1))==labels.end()){
			return {arg_t::undeflabel,false,0,str.substr(1)};
		}
		Location ref=labels.find(str.substr(1))->second;
		assert(ref.bank==ip.bank);
		return {arg_t::number,false,ref.pos-ip.pos,{}};
	} else if(str[0]=='$'){
		return {arg_t::constant,false,0,str.substr(1)};
	} else if(isdigit(str[0])){
		return {arg_t::number,false,atoi(str.data()),{}};
	} else {
		assert(false);
	}
}

Team assemble(const string &source){
	Team team;
	unordered_map<string,Location> labels;

	for(string line : splittrim(source,'\n')){
		if(line.size()==0)continue;

		string word1=line.substr(0,line.find(' '));
		string word2r=line.find(' ')==string::npos?"":line.substr(line.find(' ')+1);
		if(word1=="published"){
			string word2=word2r.substr(0,word2r.find(' '));
			string word3r=word2r.find(' ')==string::npos?"":word2r.substr(word2r.find(' ')+1);
			if(word2=="name"){
				team.name=word3r;
			}
		} else if(word1=="bank"){
			team.banks.emplace_back();
		} else if(word1[0]=='@'){
			assert(line.find(' ')==string::npos);
			assert(labels.find(word1.substr(1))==labels.end());
			labels[word1.substr(1)]={(int)team.banks.size()-1,(int)team.banks.back().size()};
		} else {
			assert(team.banks.size()>0);
			vector<string> args=splittrim(word2r,',');
			vector<Argument> pargs;
			for(const string &s : args){
				pargs.push_back(parseArgument(s,labels,
					{(int)team.banks.size()-1,(int)team.banks.back().size()}));
			}
			if(word1=="debugger"||word1=="die"||word1=="move")assert(pargs.size()==0);
			else if(word1=="jump"||word1=="turn"||word1=="scan")assert(pargs.size()==1);
			else if(word1=="create")assert(pargs.size()==3);
			else assert(pargs.size()==2);

			Script &s=team.banks.back();
			if(word1=="set")s.push_back(Instruction::make(ins_t::set,{pargs[0],pargs[1]}));
			else if(word1=="add")s.push_back(Instruction::make(ins_t::add,{pargs[0],pargs[1]}));
			else if(word1=="sub")s.push_back(Instruction::make(ins_t::sub,{pargs[0],pargs[1]}));
			else if(word1=="comp")s.push_back(Instruction::make(ins_t::comp,{pargs[0],pargs[1]}));
			else if(word1=="trans")s.push_back(Instruction::make(ins_t::trans,{pargs[0],pargs[1]}));
			else if(word1=="jump")s.push_back(Instruction::make(ins_t::jump,{pargs[0]}));
			else if(word1=="bjump")s.push_back(Instruction::make(ins_t::bjump,{pargs[0],pargs[1]}));
			else if(word1=="die")s.push_back(Instruction::make(ins_t::die,{}));
			else if(word1=="move")s.push_back(Instruction::make(ins_t::move,{}));
			else if(word1=="turn")s.push_back(Instruction::make(ins_t::turn,{pargs[0]}));
			else if(word1=="scan")s.push_back(Instruction::make(ins_t::scan,{pargs[0]}));
			else if(word1=="create")s.push_back(Instruction::make(ins_t::create,{pargs[0],pargs[1],pargs[2]}));
			else if(word1=="debugger")s.push_back(Instruction::make(ins_t::debugger,{}));
			else assert(false);
		}
	}

	int bank=0;
	for(Script &script : team.banks){
		int pos=0;
		for(Instruction &ins : script){
			for(Argument &arg : ins.args){
				if(arg.type==arg_t::undeflabel){
					assert(labels.find(arg.name)!=labels.end());
					assert(labels[arg.name].bank==bank);
					arg.type=arg_t::number;
					arg.num=labels[arg.name].pos-pos;
				}
			}
			pos++;
		}
		bank++;
	}

	return team;
}

string preprocess(istream &file){
	unordered_map<string,string> defines;

	string result;

	string line;
	while(getline(file,line)){
		if(line.find(';')!=string::npos){
			line.erase(line.find(';'));
		}
		trim(line);
		if(line.size()==0)continue;
		squashSpaces(line);
		lowercase(line);

		string word1=line.substr(0,line.find(' '));
		if(word1=="define"){
			string word2r=line.find(' ')==string::npos?"":line.substr(line.find(' ')+1);
			string word2=word2r.substr(0,word2r.find(' '));
			string word3r=word2r.find(' ')==string::npos?"":word2r.substr(word2r.find(' ')+1);
			assert(defines.find(word2.substr(1))==defines.end());
			assert(word3r[0]=='{');
			string value=word3r.substr(1);
			assert(word3r.find('}')==string::npos||word3r.find('}')==word3r.size()-1);
			if(word3r.back()!='}'){
				while(true){
					char c=file.get();
					if(c=='}')break;
					value+=c;
				}
				assert(file.get()=='\n');
			} else {
				value.erase(value.size()-1);
			}
			defines[word2.substr(1)]=value;
		} else {
			while(line.find('&')!=string::npos){
				size_t start=line.find('&')+1;
				size_t end=line.find_first_not_of("abcdefghijklmnopqrstuvwxyz0123456789._",start);
				if(end==string::npos)end=line.size();
				assert(defines.find(line.substr(start,end-start))!=defines.end());
				line=line.substr(0,start-1)+defines.find(line.substr(start,end-start))->second+line.substr(end);
			}
			result+=line+'\n';
		}
	}

	return result;
}