using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Data;
using System.Windows.Documents;
using System.Windows.Input;
using System.Windows.Media;
using System.Windows.Media.Imaging;
using System.Windows.Navigation;
using System.Windows.Shapes;
using System.ComponentModel;
using System.Collections.ObjectModel;
using System.IO;

namespace Analyze
	{
	/// <summary>
	/// Interaction logic for MainWindow.xaml
	/// </summary>
	public partial class MainWindow : Window
		{
		public MainWindow()
			{
			InitializeComponent();
			Load();
			statBox.ItemsSource = observableStats;
			suffixBox.ItemsSource = observableSuffixes;
			fileBox.ItemsSource = fileStats;
			}

		ObservableCollection<FileStat> fileStats = new ObservableCollection<FileStat>();

		ObservableCollection<InstructionStat> observableStats = new ObservableCollection<InstructionStat>();
		List<InstructionStat> stats = new List<InstructionStat>();

		// perform the action on the UI thread
		void SendToUIThread(Action action)
			{
			Dispatcher.BeginInvoke(System.Windows.Threading.DispatcherPriority.Normal,action);
			}

		List<string> filenames = new List<string>();
		public class  Suffix
			{
			public string SuffixText {get; set; }
			public long Count {get; set; }
			}

		ObservableCollection<Suffix> observableSuffixes = new ObservableCollection<Suffix>();

		void TallyMessage(Loader.XMLArgs message)
			{
			switch (message.type)
				{
				case Loader.XMLArgs.Type.Error :
					logBox.Items.Add(message.Message);
					break;
				case Loader.XMLArgs.Type.File :
					fileStats.Add(message.fileStat);
					filenames.Add(message.Message);
					break;
				case  Loader.XMLArgs.Type.Suffix :
					int count = filenames.Count(f=>System.IO.Path.GetExtension(f).ToLower() == message.Message.ToLower());
					observableSuffixes.Add(new Suffix() { SuffixText = message.Message, Count = count });
					break;
				case Loader.XMLArgs.Type.TotalStat :
					observableStats.Add(message.instructionStat);
					stats.Add(message.instructionStat);
					break;
				}
			
			}

		void Finish(object sender, EventArgs param)
			{

			logBox.Items.Add("File count : " + filenames.Count);
			logBox.Items.Add("Suffix count : " + observableSuffixes.Count);
			logBox.Items.Add("Stat count : " + stats.Count);
			long insCount = stats.Sum(s=>s.Count);
			logBox.Items.Add("Instruction count : " + insCount);
			SaveSuffixes();
			SaveFiles();
			}

		void SaveSuffixes()
			{
			var filename = "suffixes.csv";
			StringBuilder sb = new StringBuilder("Suffix, Count,"+Environment.NewLine);
			foreach (var s in observableSuffixes)
				sb.Append(String.Format("{0},{1},"+Environment.NewLine,s.SuffixText,s.Count));
			File.WriteAllText(filename,sb.ToString());
			logBox.Items.Add("Suffix file " + filename + " saved");
			}

		void SaveFiles()
			{
			var filename = "files.csv";
			StringBuilder sb = new StringBuilder("Code bytes, Instruction bytes, Instructions, Filename,"+Environment.NewLine);
			foreach (var f in fileStats)
				{
				sb.Append(String.Format("{0},{1},{2},{3},"+Environment.NewLine,
					f.CodeBytes, f.InstructionBytes, f.Instructions, f.Filename));
				}
			File.WriteAllText(filename,sb.ToString());
			logBox.Items.Add("Files file " + filename + " saved");
			}

		void SaveStats(string filename, List<InstructionStat> stats)
			{
			StringBuilder sb = new StringBuilder(
				"Displacement Length, Immediate Length, Instruction Index, Mnemonic, "+
				"Mod byte, SIB byte, Opcode Length, Prefix, Prefix Length, Count, "+Environment.NewLine
			);
			foreach (var s in stats)
				{
				sb.Append(String.Format("{0},{1},{2},{3},{4},{5},{6},{7},{8},{9},"+Environment.NewLine,
					s.DisplacementLength, s.ImmediateLength, s.InstructionIndex, 
					s.Mnemonic, s.ModifierLength, s.SIBLength, s.OpcodeLength,
					s.Prefix, s.PrefixLength, s.Count
					));
				}
			File.WriteAllText(filename,sb.ToString());
			logBox.Items.Add("Stat file " + filename + " saved");
			}


		BackgroundWorker loader = new BackgroundWorker();
		void Load()
			{
			var filename  = @"Windows7x32Disasm.xml";
			if (!File.Exists(filename))
				{
				MessageBox.Show("Cannot find file " + filename + " in directory " + Directory.GetCurrentDirectory());
				return;
				}
			Loader cc = new Loader();
			cc.MessageEvent += (o, e) => SendToUIThread(() => TallyMessage(e));
			loader.DoWork += (o, e) => {
				cc.Parse(filename);
				};
			loader.RunWorkerCompleted += Finish;
			loader.RunWorkerAsync();
			}

		void ViewStats(List<InstructionStat> stats, string filename)
			{
			observableStats.Clear();
			foreach (var s in stats)
				observableStats.Add(s);
			if (filename != null)
				SaveStats(filename, stats);
			}

		private void unpackedView_Click(object sender, RoutedEventArgs e)
			{
			ViewStats(stats,"UnpackedStats.csv");
			}

		/// <summary>
		/// represent a frequency tuple
		/// </summary>
		class Tuple
			{
			public long Prefix { get; set; }
			public long Opcode { get; set; }
			public long Modifier { get; set; }
			public long SIB { get; set; }
			public long Displacement { get; set; }
			public long Immediate { get; set; }
			public long Count { get; set; }
			}

		private void lengthView_Click(object sender, RoutedEventArgs e)
			{ // for computing the complexity of instruction remapping:
			  // count unique tuples of (p,c,m,s,d,L)
			// which are lengths: prefix 0-4, opcode 0-4?, modrm+SIB 0-2, displacement 0-4, immediate 0-4.
			List<Tuple> tuples = new List<Tuple>();
			foreach (var s in stats)
				{
				Tuple t = new Tuple();
				t.Prefix = s.PrefixLength;
				t.Opcode = s.OpcodeLength;
				t.Modifier = s.ModifierLength;
				t.SIB = s.SIBLength;
				t.Displacement = s.DisplacementLength;
				t.Immediate = s.ImmediateLength;
				t.Count = s.Count;

				// see if exists
				Tuple found = null;
				foreach (var t2 in tuples)
					{
					if ((t.Immediate == t2.Immediate) &&
						(t.Modifier == t2.Modifier) &&
						(t.Opcode == t2.Opcode) &&
						(t.Prefix == t2.Prefix) &&
						(t.Displacement == t2.Displacement))
						found = t2;
					}

				if (found == null)
					tuples.Add(t);
				else
					found.Count += t.Count;
				}
			List<InstructionStat> newStats = new List<InstructionStat>();
			var filename = "LengthStats.csv";
			using (var file = File.CreateText(filename))
				{
				file.WriteLine("Prefix Length, Opcode Length, ModRM, SIB, Displacement Length, Immediate Length,Count,");
				foreach (var t in tuples)
					{
					InstructionStat stat = new InstructionStat();
					stat.Count = t.Count;
					stat.DisplacementLength = t.Displacement;
					stat.ImmediateLength = t.Immediate;
					stat.ModifierLength = t.Modifier;
					stat.OpcodeLength = t.Opcode;
					stat.PrefixLength = t.Prefix;
					// create length tuple, store as mnemonic
					long mod = t.Modifier, sib = 0;
					if (mod > 1)
						{
						sib = mod - 1;
						mod = 1;
						}
					var line = String.Format("{0},{1},{2},{3},{4},{5}",
						t.Prefix, t.Opcode, t.Modifier, t.SIB, t.Displacement, t.Immediate
						);
					file.WriteLine(line+","+stat.Count+",");
					stat.Mnemonic = "(" + line + ")";

					newStats.Add(stat);
					}
				}

			logBox.Items.Add("Stat file " + filename + " saved");
			ViewStats(newStats, null);
			}

		private void uniqueIDView_Click(object sender, RoutedEventArgs e)
			{
			HashSet<long> ids = new HashSet<long>();
			foreach (var s in stats)
				ids.Add(s.InstructionIndex);
			List<InstructionStat> newStats = new List<InstructionStat>();
			foreach (var id in ids)
				{
				var match = stats.Where(s => s.InstructionIndex == id);
				var instr = new InstructionStat();
				foreach (var m in match)
					{
					instr.InstructionIndex = m.InstructionIndex;
					instr.Mnemonic = m.Mnemonic;
					instr.Count += m.Count;
					}
				newStats.Add(instr);
				}
			ViewStats(newStats,"IDStats.csv");
			}

		private void mnenomicView_Click(object sender, RoutedEventArgs e)
			{
			HashSet<string> mnemonics = new HashSet<string>();
			foreach (var s in stats)
				mnemonics.Add(s.Mnemonic);
			List<InstructionStat> newStats = new List<InstructionStat>();
			foreach (var mn in mnemonics)
				{
				var match = stats.Where(s => s.Mnemonic == mn);
				var instr = new InstructionStat();
				foreach (var m in match)
					{
					instr.Mnemonic = m.Mnemonic;
					instr.Count += m.Count;
					}
				newStats.Add(instr);
				}
			ViewStats(newStats,"MnemonicStats.csv");
			}
		}
	}
