Diagnostic Explorer

#region Copyright

// Diagnostic Explorer, a .Net diagnostic toolset
// Copyright (C) 2010 Cameron Elliot
// 
// This file is part of Diagnostic Explorer.
// 
// Diagnostic Explorer is free software: you can redistribute it and/or modify
// it under the terms of the GNU Lesser General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
// 
// Diagnostic Explorer is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Lesser General Public License for more details.
// 
// You should have received a copy of the GNU Lesser General Public License
// along with Diagnostic Explorer.  If not, see <http://www.gnu.org/licenses/>.
// 
// http://diagexplorer.sourceforge.net/

#endregion

using System;
using System.CodeDom.Compiler;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Reflection;
using System.Text;
using System.Text.RegularExpressions;
using log4net;
using System.Linq;

namespace DiagnosticExplorer
{
	/// <summary>Enabled trace to a single source through method calls</summary>
	public class TraceScope : IDisposable
	{
		[ThreadStatic] private static List<TraceScope> _scopeStack;
		private DateTime _created = DateTime.UtcNow;
		private DateTime? _disposed;
		private List<TraceItem> _traceItems = new List<TraceItem>();
		private string _name;
		private bool _forceTrace;
		private bool _isRoot;

		/// <summary>
		/// Sorted dictionary of time in milliseconds vs the trace method
		/// which should be used if the operation exceeds that time
		/// </summary>
		private SortedDictionary<int, Action<string>> _traceMethods;

		#region Constructors

		public TraceScope()
		{
			Setup(GetCurrentMethodName(), null, null, false);
		}

		public TraceScope(TraceMode mode)
		{
			Setup(GetCurrentMethodName(), null, mode, false);
		}

		public TraceScope(Action<string> traceMethod)
		{
			Setup(GetCurrentMethodName(), traceMethod, null, false);
		}

		public TraceScope(Action<string> traceMethod, TraceMode mode)
		{
			Setup(GetCurrentMethodName(), traceMethod, mode, false);
		}

		public TraceScope(string name)
		{
			Setup(name, null, null, false);
		}

		public TraceScope(string name, TraceMode mode)
		{
			Setup(name, null, mode, false);
		}

		public TraceScope(string name, Action<string> traceMethod)
		{
			Setup(name, traceMethod, null, false);
		}

		public TraceScope(string name, Action<string> traceMethod, TraceMode mode)
		{
			Setup(name, traceMethod, mode, false);
		}
	
		public TraceScope(string name, Action<string> traceMethod, bool forceTrace)
		{
			Setup(name, traceMethod, null, forceTrace);
		}

		public TraceScope(Action<string> traceMethod, bool forceTrace)
		{
			Setup(GetCurrentMethodName(), traceMethod, null, forceTrace);
		}

		public TraceScope(string name, Action<string> traceMethod,  TraceMode mode, bool forceTrace)
		{
			Setup(name, traceMethod, mode, forceTrace);
		}

		private void Setup(string name, Action<string> traceMethod, TraceMode? mode, bool forceTrace)
		{
			_traceMethods = new SortedDictionary<int, Action<string>>();
			SetTraceAction(0, traceMethod);

			_name = name;
			ScopeTraceMode = mode;
			_forceTrace = forceTrace;

			_isRoot = ScopeStack.Count == 0;

			if (ScopeStack.Count != 0)
				CurrentScope._traceItems.Add(new TraceItem(this));

			ScopeStack.Add(this);
		}

		#endregion

		public TraceMode? ScopeTraceMode { get; set; }

		public TimeSpan? SuppressDetailThreshold { get; set; }

		public void SetTraceAction(Action<string> traceMethod)
		{
			SetTraceAction(0, traceMethod);
		}

		public void SetTraceAction(int timeMillis, Action<string> traceMethod)
		{
			_traceMethods[timeMillis] = traceMethod;
		}

		public static TraceMode TraceMode
		{
			get
			{
				if (_scopeStack == null) return TraceMode.Normal;

				for (int i = _scopeStack.Count - 1; i >= 0; i--)
				{
					TraceScope scope = _scopeStack[i];
					if (scope.ScopeTraceMode.HasValue)
						return scope.ScopeTraceMode.Value;
				}
				return TraceMode.Normal;
			}
		}

		public static string GetCurrentMethodName()
		{
			StackTrace t = new StackTrace(false);

			MethodBase method = t.GetFrame(2).GetMethod();
			StringBuilder sb = new StringBuilder();
			AppendType(sb, method.DeclaringType);
			sb.Append(".");
			sb.Append(method.Name);
			if (method.IsGenericMethod)
				WriteGenericArguments(sb, method.GetGenericArguments());
			sb.Append("(");
			foreach (ParameterInfo info in method.GetParameters())
			{
				AppendType(sb, info.ParameterType);
				sb.Append(", ");
			}

			if (sb[sb.Length - 1] != '(')
				sb.Length -= 2;
			sb.Append(")");
			return sb.ToString();
		}

		private static void WriteGenericArguments(StringBuilder sb, Type[] types)
		{
			sb.Append("<");
			foreach (Type t in types)
			{
				AppendType(sb, t);
				sb.Append(", ");
			}
			if (sb[sb.Length - 1] != '<')
				sb.Length -= 2;
			sb.Append(">");
		}

		private static void AppendType(StringBuilder sb, Type type)
		{
			if (!type.IsGenericType)
			{
				sb.Append(type.Name);
			}
			else
			{
				sb.Append(Regex.Replace(type.UnderlyingSystemType.Name, "`.*", ""));
				WriteGenericArguments(sb, type.GetGenericArguments());
			}
		}

		public TimeSpan Age
		{
			get { return DateTime.UtcNow.Subtract(_created); }
		}

		private static List<TraceScope> ScopeStack
		{
			get
			{
				if (_scopeStack == null)
					_scopeStack = new List<TraceScope>();
				return _scopeStack;
			}
		}

		public override string ToString()
		{
			StringBuilder sb = new StringBuilder();
			StringWriter writer = new StringWriter(sb);
			IndentedTextWriter indentedWriter = new IndentedTextWriter(writer);

			DateTime lastMessage = _created;
			ToString(indentedWriter, _created, ref lastMessage, 0);
			indentedWriter.Flush();
			writer.Flush();
			return sb.ToString();
		}

		private void ToString(IndentedTextWriter writer, DateTime scopeStart, ref DateTime lastMessage, int indent)
		{
			if (_traceItems.Count == 0)
			{
				WriteBeginEndScope(this, writer, scopeStart, ref lastMessage, false);
			}
			else
			{
				WriteBeginScope(writer, scopeStart, ref lastMessage);

				writer.Indent += indent;
				foreach (TraceItem item in _traceItems)
				{
					if (item.TraceScope == null)
					{
						WriteString(writer, item.Message, item.Created, scopeStart, ref lastMessage);
					}
					else if (FullTraceRequired(item))
					{
						item.TraceScope.ToString(writer, scopeStart, ref lastMessage, 1);
					}
					else
					{
						WriteBeginEndScope(item.TraceScope, writer, scopeStart, ref lastMessage, true);
					}
				}
				writer.Indent -= indent;

				WriteEndScope(writer, scopeStart, ref lastMessage);
			}
		}

		private bool FullTraceRequired(TraceItem item)
		{
			if (item.TraceScope == null) return false;
			if (item.TraceScope._forceTrace) return false;
			if (item.TraceScope.SuppressDetailThreshold == null) return true;

			TimeSpan thresh = item.TraceScope.SuppressDetailThreshold.Value;
			TimeSpan duration = item.TraceScope._disposed.Value - item.TraceScope._created;
			return duration >= thresh;
		}

		private void WriteBeginScope(IndentedTextWriter writer, DateTime scopeStart, ref DateTime lastMessage)
		{
			if (_name == null) return;

			string message = string.Format("BEGIN {0}", _name);
			if (_disposed != null)
				message = string.Format("{0} ({1:N3} seconds)", message, _disposed.Value.Subtract(_created).TotalSeconds);

			WriteString(writer, message, _created, scopeStart, ref lastMessage);
		}

		private void WriteEndScope(IndentedTextWriter writer, DateTime scopeStart, ref DateTime lastMessage)
		{
			if (_name == null) return;
			if (_disposed == null) return;

			string message = string.Format("END {0} ({1:N3} seconds)", _name, _disposed.Value.Subtract(_created).TotalSeconds);

			WriteString(writer, message, _disposed.Value, scopeStart, ref lastMessage);
		}

		private static void WriteBeginEndScope(TraceScope scope, IndentedTextWriter writer,
			DateTime scopeStart, ref DateTime lastMessage, bool suppressed)
		{
			if (scope._name == null) return;
			if (scope._disposed == null) return;

			TimeSpan duration = scope._disposed.Value.Subtract(scope._created);

			string message = string.Format("BEGIN/END{0} {1} ({2:N3} seconds)",
			                               suppressed ? "*" : "",
			                               scope._name,
			                               duration.TotalSeconds);

			WriteString(writer, message, scope._disposed.Value, scopeStart, ref lastMessage);
		}

		private static void WriteString(IndentedTextWriter writer, string message, DateTime itemDate, DateTime scopeStart, ref DateTime lastMessage)
		{
			double age = itemDate.Subtract(scopeStart).TotalSeconds;
			double split = itemDate.Subtract(lastMessage).TotalSeconds;

			string[] parts = Regex.Split(message, @"\r?\n");
			writer.Write("[{0:00.000}] [{1:00.000}] ", age, split);
			writer.WriteLine(parts[0].Trim());

			for (int i = 1; i < parts.Length; i++)
			{
				writer.Write("                  ");
				writer.WriteLine(parts[i]);
			}

			lastMessage = itemDate;
		}

		public void Dispose()
		{
			Dispose(true);
			GC.SuppressFinalize(this);
		}

		private void Dispose(bool disposing)
		{
			if (disposing)
			{
				_disposed = DateTime.UtcNow;

				if (ScopeStack.Count == 0) return;

				ScopeStack.RemoveAt(ScopeStack.Count - 1);
				if (ScopeStack.Count == 0)
					_scopeStack = null;
				TraceMessage();
			}
		}

		private void TraceMessage()
		{
			try
			{
				KeyValuePair<int, Action<string>> tracePair = GetTraceMethod(Age.TotalMilliseconds);
				if (tracePair.Value == null) return;
				if (!_isRoot && !_forceTrace) return;

				tracePair.Value(ToString());
			}
			catch (Exception ex)
			{
				LogManager.GetLogger(GetType()).Error(ex.Message, ex);
			}
		}

		private KeyValuePair<int, Action<string>> GetTraceMethod(double millis)
		{
			foreach (var pair in _traceMethods.Reverse())
				if (millis >= pair.Key)
					return pair;

			return _traceMethods.FirstOrDefault();
		}

		#region Write Methods

		public static void Trace(StackTrace trace)
		{
			string s = trace.ToString();
			if (s.Length > 200)
				s = s.Substring(0, 200);
			Trace(s);
		}

		public static void Trace(object o)
		{
			Trace(o == null ? "" : o.ToString());
		}

		public static void TraceIf(bool condition, object o)
		{
			if (condition)
				Trace(o == null ? "" : o.ToString());
		}

		public static void Trace(string format, params object[] args)
		{
			if (ScopeStack == null) return;
			if (ScopeStack.Count == 0) return;

			format = TryFormat(format, args);
			CurrentScope._traceItems.Add(new TraceItem(format));
		}

		public static void TraceIf(bool condition, string format, params object[] args)
		{
			if (condition)
				Trace(format, args);
		}
		
		private static TraceScope CurrentScope
		{
			get
			{
				if (_scopeStack == null) return null;
				if (_scopeStack.Count == 0) return null;
				return _scopeStack[_scopeStack.Count - 1];
			}
		}

		private static string TryFormat(string format, object[] args)
		{
			try
			{
				if (format == null) return "";
				if (args != null && args.Length != 0)
					format = string.Format(format, args);
				
				return format;
			}
			catch
			{
				return format;
			}
		}

		#endregion
	}
}