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) : NExpr {
				Walk(
					tree,
					fun(expr : NExpr) {
						| Assign(left, right) =>
							match(left) {
								| Var(var_) when var_ == var => right
								| _ => null
							}
						| _ => null
					}
				) :> NExpr
			}
			
			def FindAllAssigns(tree, var) : list [object] {
				WalkAll(
					tree,
					fun(expr : NExpr) {
						| Assign(left, _) =>
							match(left) {
								| Var(var_) when var_ == var => expr
								| _ => null
							}
						| _ => null
					}
				)
			}
			
			def Evaluate(expr : NExpr, stack) {
				match(expr) {
					| Number(num) => num
					| 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(initialState) => expr
						| _ => null
					}
				) :> NExpr;
				def stack = case :: stack;
				def stateExpr = FindAssign(case, stateObfsVar);
				def nextState = (Evaluate(stateExpr, stack) + stateOff) :> uint;
				DefyLogic(nextState, stack);
			}
			
			def tree = Dump.DumpTree;
			// NExpr.Print(tree)
			
			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
			);
			
			DefyLogic(initialState, []);
		}
		
		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
			}
		}
	}
}