using NObfs.Tree;
using NObfs.Dump;
using System;

namespace NObfs {
	public module NObfs {
		public Main() : void {
			def Expr(expr : NExpr) {
				| Var(var) => var : object
				| Number(num) => num : object
				| x => x
			}
			
			def FindAssign(tree, var, ass : NExpr = null) : NExpr {
				Walk(
					tree,
					fun(expr : NExpr) {
						| Assign(left, right) when (right : object) != ass =>
							match(left) {
								| Var(var_) when var_ == var => 
									right
								| _ => null
							}
						| _ => null
					}
				) :> NExpr
			}
			
			def tree = Dump.DumpTree;
			
			def stateDeobfsVar = Walk(
				tree,
				fun(expr : NExpr) {
					| Switch(expr, _) => Expr(expr)
					| _ => null
				}
			) :> string;
			Console.WriteLine(
				"Deobfs state == {0}",
				stateDeobfsVar
			);
			
			def (stateObfsVar, stateOff) = match(FindAssign(tree, stateDeobfsVar)) {
				| Arith(type_, left, right) =>
					(
						Expr(left) :> string,
						match(type_) {
							| Add =>   Expr(right) :> long
							| Sub => -(Expr(right) :> long)
							| _ => throw Exception("WTF?")
						}
					)
				| _ => throw Exception("Uh?")
			}
			Console.WriteLine(
				"Obfs state == {0}",
				stateObfsVar
			);
			Console.WriteLine(
				"State offset == {0}",
				stateOff
			);
			
			def initialState = (Expr(FindAssign(tree, stateObfsVar)) :> long + stateOff) :> uint;
			Console.WriteLine(
				"Initial state == {0}",
				initialState
			);
			
			def EvaluateVar(var, ass, stack) {
				match(stack) {
					| [] =>
						def expr = NExpr.Number(-1);
						(expr, expr, stack)
					| case :: rest =>
						def expr = FindAssign(case, var, ass);
						if(expr == null)
							EvaluateVar(var, ass, rest)
						else
							(expr, expr, stack)
				}
			}
			
			def Evaluate(expr : NExpr, ass, stack) {
				match(expr) {
					| Comp(type_, left, right) =>
						def left  = Evaluate(left, ass, stack);
						def right = Evaluate(right, ass, stack);
						if(
							match(type_) {
								| EQ  => left == right
								| NE  => left != right
								| SGE => left >= right
								| UGE => left >= right
								| SLE => left <= right
								| ULE => left <= right
								| SGT => left >  right
								| UGT => left >  right
								| SLT => left <  right
								| ULT => left <  right
							}
						)
							1
						else
							0
					
					| Arith(type_, left, right) =>
						def left  = Evaluate(left, ass, stack);
						def right = Evaluate(right, ass, stack);
						match(type_) {
							| Or   => (left |  right)
							| Xor  =>  left ^  right
							| And  =>  left &  right
							| SShr =>  left >> right
							| UShr =>  left >> right
							| Shl  =>  left << right
							| Add  =>  left +  right
							| Sub  =>  left -  right
							| Mul  =>  left *  right
							| SDiv =>  left /  right
							| UDiv =>  left /  right
							| SMod =>  left %  right
							| UMod =>  left %  right
						}
					
					| Cast(type_, expr) =>
						def value = Evaluate(expr, ass, stack);
						def sign = {
							if(value >= 0) 1
							else -1
						};
						match(type_) {
							| "char" => (value % 128) * sign
							| x => throw Exception(String.Format("Unknown cast type '{0}'", x))
						}
					
					| Var(var) =>
						Evaluate(EvaluateVar(var, ass, stack))
					
					| Number(num) => num :> int
					| x => throw Exception(String.Format("Unknown type {0}", x))
				}
			}
			
			def DefyLogic(state, stack) : void {
				def case = Walk(
					tree,
					fun(expr : NExpr) {
						| Case(values, _) when values.Contains(state) => expr
						| _ => null
					}
				) :> NExpr;
				def stack = case :: stack;
				def stateExpr = FindAssign(case, stateObfsVar);
				def nextState = (Evaluate(stateExpr, stateExpr, stack) + stateOff) :> uint;
				Console.WriteLine("Next state == {0}", nextState);
				unless(nextState == 0U)
					DefyLogic(nextState, stack);
			}
			
			DefyLogic(initialState, [tree]);
		}
		
		Walk(tree : NExpr, func : NExpr -> object) : object {
			def WalkList(exprs, func) {
				match(exprs) {
					| [] => null
					| head :: tail =>
						def ret = Walk(head, func);
						if(ret == null)
							WalkList(tail, func)
						else
							ret
				}
			}
			
			def ret = func(tree);
			if(ret == null) {
				match(tree) {
					| Loc(_, body)
					| Case(_, body)
					| While(_, body) => Walk(body, func)
					
					| Switch(_, body)
					| Block(body) => WalkList(body, func)
					
					| If(_, then, else_) =>
						def ret = Walk(then, func);
						if(ret == null)
							Walk(else_, func)
						else
							ret
					| _ => null
				}
			} else
				ret
		}
		
		WalkAll(tree : NExpr, func : NExpr -> object, mutable accum : list [object] = null) : list [object] {
			def WalkList(exprs, func, accum) {
				match(exprs) {
					| [] => accum
					| head :: tail =>
						WalkList(
							tail,
							func,
							WalkAll(head, func, accum)
						)
				}
			}
			
			when(accum == null)
				accum = [];
			
			def ret = func(tree);
			unless(ret == null)
				accum = ret :: accum;
				
			match(tree) {
				| Loc(_, body)
				| Case(_, body)
				| While(_, body) => WalkAll(body, func, accum)
				
				| Switch(_, body)
				| Block(body) => WalkList(body, func, accum)
				
				| If(_, then, else_) =>
					WalkAll(
						else_, 
						func,
						WalkAll(
							then,
							func,
							accum
						)
					)
				| _ => accum
			}
		}
	}
}