/*
 * Process Hacker - 
 *   thread provider
 * 
 * Copyright (C) 2008-2009 wj32
 * 
 * This file is part of Process Hacker.
 * 
 * Process Hacker is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Process Hacker 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 General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Process Hacker.  If not, see <http://www.gnu.org/licenses/>.
 */

using System;
using System.Collections.Generic;
using System.Threading;
using ProcessHacker.Common;
using ProcessHacker.Common.Messaging;
using ProcessHacker.Native;
using ProcessHacker.Native.Api;
using ProcessHacker.Native.Objects;
using ProcessHacker.Native.Security;
using ProcessHacker.Native.Symbols;

namespace ProcessHacker
{
    public class ThreadItem : ICloneable
    {
        public object Clone()
        {
            return this.MemberwiseClone();
        }

        public int RunId;
        public int Tid;

        public long ContextSwitches;
        public long ContextSwitchesDelta;
        public ulong Cycles;
        public ulong CyclesDelta;
        public int PriorityI;
        public string Priority;
        public IntPtr StartAddressI;
        public string StartAddress;
        public SymbolResolveLevel StartAddressLevel;
        public KWaitReason WaitReason;
        public bool IsGuiThread;
        public bool JustResolved;

        public ThreadHandle ThreadQueryLimitedHandle;
    }

    public class ThreadProvider : Provider<int, ThreadItem>
    {
        private class ResolveMessage : Message
        {
            public int Tid;
            public string Symbol;
            public SymbolResolveLevel ResolveLevel;
        }

        public delegate void LoadingStateChangedDelegate(bool loading);
        private delegate void ResolveThreadStartAddressDelegate(int tid, ulong startAddress);

        private static readonly WorkQueue _symbolsWorkQueue = new WorkQueue() { MaxWorkerThreads = 1 };

        public event LoadingStateChangedDelegate LoadingStateChanged;

        private ProcessHandle _processHandle;
        private ProcessAccess _processAccess;
        private SymbolProvider _symbols;
        private bool _kernelSymbolsLoaded = false;
        private int _pid;
        private int _loading = 0;
        private MessageQueue _messageQueue = new MessageQueue();
        private EventWaitHandle _moduleLoadCompletedEvent = new EventWaitHandle(false, EventResetMode.ManualReset);
        private bool _waitedForLoad = false;

        public ThreadProvider(int pid)
            : base()
        {
            this.Name = this.GetType().Name;
            _pid = pid;

            _messageQueue.AddListener(
                new MessageQueueListener<ResolveMessage>((message) =>
                    {
                        if (message.Symbol != null)
                        {
                            this.Dictionary[message.Tid].StartAddress = message.Symbol;
                            this.Dictionary[message.Tid].StartAddressLevel = message.ResolveLevel;
                            this.Dictionary[message.Tid].JustResolved = true;
                        }
                    }));

            this.ProviderUpdate += new ProviderUpdateOnce(UpdateOnce);
            this.Disposed += ThreadProvider_Disposed;

            try
            {
                // Try to get a good process handle we can use the same handle for stack walking.
                try
                {
                    _processAccess = ProcessAccess.QueryInformation | ProcessAccess.VmRead;
                    _processHandle = new ProcessHandle(_pid, _processAccess);
                }
                catch
                {
                    try
                    {
                        if (KProcessHacker.Instance != null)
                        {
                            _processAccess = Program.MinProcessReadMemoryRights;
                            _processHandle = new ProcessHandle(_pid, _processAccess);
                        }
                        else
                        {
                            _processAccess = Program.MinProcessQueryRights;
                            _processHandle = new ProcessHandle(_pid, _processAccess);
                        }
                    }
                    catch (Exception ex)
                    {
                        Logging.Log(ex);
                    }
                }

                try
                {
                    // Needed (maybe) to display the EULA
                    Win32.SymbolServerSetOptions(SymbolServerOption.Unattended, 0);
                }
                catch (Exception ex)
                {
                    Logging.Log(ex);
                }

                // start loading symbols; avoid the UI blocking on the dbghelp call lock
                _symbolsWorkQueue.QueueWorkItemTag(new Action(() =>
                {
                    try
                    {
                        // Use the process handle if we have one, otherwise use the default ID generator.
                        if (_processHandle != null)
                            _symbols = new SymbolProvider(_processHandle);
                        else
                            _symbols = new SymbolProvider();

                        SymbolProvider.Options = SymbolOptions.DeferredLoads |
                            (Properties.Settings.Default.DbgHelpUndecorate ? SymbolOptions.UndName : 0);

                        if (Properties.Settings.Default.DbgHelpSearchPath != "")
                            _symbols.SearchPath = Properties.Settings.Default.DbgHelpSearchPath;

                        try
                        {
                            if (_pid != 4)
                            {
                                using (var phandle =
                                    new ProcessHandle(_pid, Program.MinProcessQueryRights | Program.MinProcessReadMemoryRights))
                                {
                                    // Load the process' modules.
                                    try { _symbols.LoadProcessModules(phandle); }
                                    catch { }

                                    // If the process is CSRSS we should load kernel modules 
                                    // due to the presence of kernel-mode threads.
                                    if (phandle.GetKnownProcessType() == KnownProcess.WindowsSubsystem)
                                        this.LoadKernelSymbols();
                                }
                            }
                            else
                            {
                                this.LoadKernelSymbols();
                            }
                        }
                        catch (Exception ex)
                        {
                            Logging.Log(ex);
                        }
                    }
                    finally
                    {
                        lock (_moduleLoadCompletedEvent)
                        {
                            if (!_moduleLoadCompletedEvent.SafeWaitHandle.IsClosed)
                                _moduleLoadCompletedEvent.Set();
                        }
                    }
                }), "symbols-load");
            }
            catch (Exception ex)
            {
                Logging.Log(ex);
            }
        }

        public ProcessAccess ProcessAccess
        {
            get { return _processAccess; }
        }

        public ProcessHandle ProcessHandle
        {
            get { return _processHandle; }
        }

        public void LoadKernelSymbols()
        {
            lock (_symbols)
            {
                if (!_kernelSymbolsLoaded)
                {
                    if (KProcessHacker.Instance != null)
                        _symbols.LoadKernelModules();

                    _kernelSymbolsLoaded = true;
                }
            }
        }

        private void ThreadProvider_Disposed(IProvider provider)
        {
            if (_symbols != null)
                _symbols.Dispose();
            if (_processHandle != null)
                _processHandle.Dispose();
            _symbols = null;

            lock (_moduleLoadCompletedEvent)
                _moduleLoadCompletedEvent.Close();

            foreach (int tid in this.Dictionary.Keys)
            {
                ThreadItem item = this.Dictionary[tid];

                if (item.ThreadQueryLimitedHandle != null)
                    item.ThreadQueryLimitedHandle.Dispose();
            }
        }

        private void ResolveThreadStartAddress(int tid, ulong startAddress)
        {
            ResolveMessage result = new ResolveMessage();

            result.Tid = tid;

            if (!_moduleLoadCompletedEvent.SafeWaitHandle.IsClosed)
            {
                try
                {
                    _moduleLoadCompletedEvent.WaitOne();
                }
                catch
                { }
            }

            if (_symbols == null)
                return;

            try
            {
                Interlocked.Increment(ref _loading);

                if (this.LoadingStateChanged != null)
                    this.LoadingStateChanged(Thread.VolatileRead(ref _loading) > 0);

                try
                {
                    result.Symbol = _symbols.GetSymbolFromAddress(startAddress, out result.ResolveLevel);
                    _messageQueue.Enqueue(result);
                }
                catch
                { }
            }
            finally
            {
                Interlocked.Decrement(ref _loading);

                if (this.LoadingStateChanged != null)
                    this.LoadingStateChanged(Thread.VolatileRead(ref _loading) > 0);
            }
        }

        public void QueueThreadResolveStartAddress(int tid)
        {
            this.QueueThreadResolveStartAddress(tid, this.Dictionary[tid].StartAddressI.ToUInt64());
        }

        public void QueueThreadResolveStartAddress(int tid, ulong startAddress)
        {
            _symbolsWorkQueue.QueueWorkItemTag(
                new ResolveThreadStartAddressDelegate(this.ResolveThreadStartAddress),
                "thread-resolve",
                tid, startAddress
                );
        }

        private string GetThreadBasicStartAddress(ulong startAddress, out SymbolResolveLevel level)
        {
            ulong modBase;
            string fileName = _symbols.GetModuleFromAddress(startAddress, out modBase);

            if (fileName == null)
            {
                level = SymbolResolveLevel.Address;
                return "0x" + startAddress.ToString("x");
            }
            else
            {
                level = SymbolResolveLevel.Module;
                return (new System.IO.FileInfo(fileName)).Name + "+0x" +
                    (startAddress - modBase).ToString("x");
            }
        }

        private void UpdateOnce()
        {
            var threads = Windows.GetProcessThreads(_pid);
            Dictionary<int, ThreadItem> newdictionary = new Dictionary<int, ThreadItem>(this.Dictionary);

            if (threads == null)
                threads = new Dictionary<int, SystemThreadInformation>();

            // look for dead threads
            foreach (int tid in Dictionary.Keys)
            {
                if (!threads.ContainsKey(tid))
                {
                    ThreadItem item = this.Dictionary[tid];

                    if (item.ThreadQueryLimitedHandle != null)
                        item.ThreadQueryLimitedHandle.Dispose();

                    this.OnDictionaryRemoved(item);
                    newdictionary.Remove(tid);
                }
            }

            // Get resolve results.
            _messageQueue.Listen();

            // look for new threads
            foreach (int tid in threads.Keys)
            {
                var t = threads[tid];

                if (!Dictionary.ContainsKey(tid))
                {
                    ThreadItem item = new ThreadItem();

                    item.RunId = this.RunCount;
                    item.Tid = tid;
                    item.ContextSwitches = t.ContextSwitchCount;
                    item.WaitReason = t.WaitReason;

                    try
                    {
                        item.ThreadQueryLimitedHandle = new ThreadHandle(tid, Program.MinThreadQueryRights);

                        try
                        {
                            item.PriorityI = (int)item.ThreadQueryLimitedHandle.GetPriorityLevel();
                            item.Priority = item.ThreadQueryLimitedHandle.GetPriorityLevel().ToString();
                        }
                        catch
                        { }

                        if (KProcessHacker.Instance != null)
                        {
                            try
                            {
                                item.IsGuiThread = KProcessHacker.Instance.KphGetThreadWin32Thread(item.ThreadQueryLimitedHandle) != 0;
                            }
                            catch
                            { }
                        }

                        if (OSVersion.HasCycleTime)
                        {
                            try
                            {
                                item.Cycles = item.ThreadQueryLimitedHandle.GetCycleTime();
                            }
                            catch
                            { }
                        }
                    }
                    catch
                    { }

                    if (KProcessHacker.Instance != null && item.ThreadQueryLimitedHandle != null)
                    {
                        try
                        {
                            item.StartAddressI = 
                                KProcessHacker.Instance.GetThreadStartAddress(item.ThreadQueryLimitedHandle).ToIntPtr();
                        }
                        catch
                        { }
                    }
                    else
                    {
                        try
                        {
                            using (ThreadHandle thandle =
                                new ThreadHandle(tid, ThreadAccess.QueryInformation))
                            {
                                item.StartAddressI = thandle.GetWin32StartAddress();
                            }
                        }
                        catch
                        {
                            item.StartAddressI = t.StartAddress;
                        }
                    }

                    if (!_waitedForLoad)
                    {
                        _waitedForLoad = true;

                        try
                        {
                            if (_moduleLoadCompletedEvent.WaitOne(0, false))
                            {
                                item.StartAddress = this.GetThreadBasicStartAddress(
                                    item.StartAddressI.ToUInt64(), out item.StartAddressLevel);
                            }
                        }
                        catch
                        { }
                    }

                    if (string.IsNullOrEmpty(item.StartAddress))
                    {
                        item.StartAddress = "0x" + item.StartAddressI.ToString("x8");
                        item.StartAddressLevel = SymbolResolveLevel.Address;
                    }

                    this.QueueThreadResolveStartAddress(tid, item.StartAddressI.ToUInt64());

                    newdictionary.Add(tid, item);
                    this.OnDictionaryAdded(item);
                }
                // look for modified threads
                else
                {
                    ThreadItem item = Dictionary[tid];
                    ThreadItem newitem = item.Clone() as ThreadItem;

                    newitem.JustResolved = false;
                    newitem.ContextSwitchesDelta = t.ContextSwitchCount - newitem.ContextSwitches;
                    newitem.ContextSwitches = t.ContextSwitchCount;
                    newitem.WaitReason = t.WaitReason;

                    try
                    {
                        newitem.PriorityI = (int)newitem.ThreadQueryLimitedHandle.GetPriorityLevel();
                        newitem.Priority = newitem.ThreadQueryLimitedHandle.GetPriorityLevel().ToString();
                    }
                    catch
                    { }

                    if (KProcessHacker.Instance != null)
                    {
                        try
                        {
                            newitem.IsGuiThread = KProcessHacker.Instance.KphGetThreadWin32Thread(newitem.ThreadQueryLimitedHandle) != 0;
                        }
                        catch
                        { }
                    }

                    if (OSVersion.HasCycleTime)
                    {
                        try
                        {
                            ulong thisCycles = newitem.ThreadQueryLimitedHandle.GetCycleTime();

                            newitem.CyclesDelta = thisCycles - newitem.Cycles;
                            newitem.Cycles = thisCycles;
                        }
                        catch
                        { }
                    }

                    if (newitem.StartAddressLevel == SymbolResolveLevel.Address)
                    {
                        if (_moduleLoadCompletedEvent.WaitOne(0, false))
                        {
                            newitem.StartAddress = this.GetThreadBasicStartAddress(
                                newitem.StartAddressI.ToUInt64(), out newitem.StartAddressLevel);
                        }

                        // If we couldn't resolve it to a module+offset, 
                        // use the StartAddress (instead of the Win32StartAddress)
                        // and queue the resolve again.
                        if (
                            item.StartAddressLevel == SymbolResolveLevel.Address &&
                            item.JustResolved)
                        {
                            if (item.StartAddressI != t.StartAddress)
                            {
                                item.StartAddressI = t.StartAddress;
                                this.QueueThreadResolveStartAddress(tid, item.StartAddressI.ToUInt64());
                            }
                        }
                    }

                    if (
                        newitem.ContextSwitches != item.ContextSwitches || 
                        newitem.ContextSwitchesDelta != item.ContextSwitchesDelta ||
                        newitem.Cycles != item.Cycles || 
                        newitem.CyclesDelta != item.CyclesDelta || 
                        newitem.IsGuiThread != item.IsGuiThread || 
                        newitem.Priority != item.Priority || 
                        newitem.StartAddress != item.StartAddress ||
                        newitem.WaitReason != item.WaitReason || 
                        item.JustResolved
                        )
                    {
                        newdictionary[tid] = newitem;
                        this.OnDictionaryModified(item, newitem);
                    }
                }
            }

            Dictionary = newdictionary;
        }

        public SymbolProvider Symbols
        {
            get { return _symbols; }
        }

        public int Pid
        {
            get { return _pid; }
        }
    }
}