# TortoiseHgOverlayServer.py - icon overlay server for Windows shellext
#
# Copyright 2009 TK Soh <teekaysoh@gmail.com>
#
# This software may be used and distributed according to the terms of the
# GNU General Public License version 2 or any later version.

# Creates a task-bar icon.  Run from Python.exe to see the
# messages printed. Takes an optional logfile as first command
# line parameter

from __future__ import annotations

import os
import sys

if hasattr(sys, "frozen"):
    if sys.frozen == 'windows_exe':
        # sys.stdin is invalid, should be None.  Fixes svn, git subrepos
        sys.stdin = None

        # py2exe blackholes stdout with a custom class to prevent random
        # exceptions when writing to it, but doesn't provide the ``.buffer``
        # attribute that mercurial.utils.procutil wants.  As of 0.11.1.0, it
        # also doesn't return the number of bytes written, and therefore
        # violates the interface contract.  For details, see:
        # https://github.com/py2exe/py2exe/blob/af0e841bffcf9c64abce6204718fecad17a59506/py2exe/boot_common.py#L3
        #
        # TODO: upstream write() return value and PROGRAMDATA usage in py2exe,
        #  and then most or all of this can go away.
        sys.stdout = open(os.devnull, "w")

        # py2exe writes its redirected stderr log file next to the executable
        # (so typically in %PROGRAMFILES%), but UAC blocks that when it is
        # enabled.  Replacing stderr with our own log file to an always
        # writeable area prevents that, fixes the missing write() return value,
        # and provides the `.buffer` attribute that procutil assumes is present.
        program_data = os.environ.get("PROGRAMDATA")

        if program_data is not None:
            import atexit, ctypes

            dir_name = os.path.join(program_data, "TortoiseHg")
            log_name = os.path.join(
                dir_name,
                os.path.splitext(
                    os.path.basename(sys.executable)
                )[0] + '.log'
            )

            try:
                os.makedirs(dir_name, exist_ok=True)
                sys.stderr = open(log_name, 'a', buffering=1)
            except Exception as details:
                sys.stderr = sys.stdout  # blackhole messages
                ctypes.windll.user32.MessageBoxW(
                    0,
                    "The logfile '%s' could not be opened:\n %s" % (
                        log_name, details,
                    ),
                    "Errors in %r" % os.path.basename(sys.executable),
                    0
                )
            else:
                orig_size = sys.stderr.seek(0, os.SEEK_END)

                def exit_check():
                    log_size = sys.stderr.tell()

                    if orig_size != log_size:
                        ctypes.windll.user32.MessageBoxW(
                            0,
                            "See the logfile '%s' for details" % log_name,
                            "Errors in %r" % os.path.basename(sys.executable),
                            0,
                        )

                atexit.register(exit_check)
        else:
            # py2exe messages come through sys.stderr; mercurial.ui.error()
            # messages come through sys.stderr.buffer.  Drop everything for the
            # *w.exe executable if the log file cannot be created, in order to
            # prevent various issues mentioned above.
            sys.stderr = sys.stdout

import errno
import time
import threading
import traceback
import gc

try:
    from win32api import *
    from win32gui import *
    import win32pipe
    import win32con
    import win32event
    import win32file
    import winerror
    import pywintypes
    import win32security
except ImportError as e:
    print('Fatal error at startup', e)
    sys.exit(1)

from mercurial import demandimport
demandimport.IGNORES.add('win32com.shell')
demandimport.enable()
from mercurial import (
    error,
    pycompat,
    util,
)

from mercurial.windows import posixfile, unlink, rename
from tortoisehg.util.i18n import _
from tortoisehg.util import hglib, thread2, paths, shlib, version

if hglib.TYPE_CHECKING:
    from typing import (
        Iterable,
        List,
        Set,
        Text,
        Tuple,
    )


APP_TITLE = _('TortoiseHg Overlay Icon Server')

EXIT_CMD = 1025

class Logger:
    def __init__(self):
        self.file = None

    def setfile(self, name: Text) -> None:
        oname = name + '.old'
        try:
            rename(hglib.fromunicode(name), hglib.fromunicode(oname))
        except:
            pass
        self.file = posixfile(hglib.fromunicode(name), b'wb')
        self.msg('%s, Version %s' % (APP_TITLE, version.version()))
        self.msg('Logging to file started')

    def msg(self, msg: Text) -> None:
        ts = '[%s] ' % time.strftime('%c')
        f = self.file
        if f:
            f.write(hglib.fromunicode(ts + msg) + b'\n')
            f.flush()
            os.fsync(f.fileno())
            print('L' + ts + msg)
        else:
            print(ts + msg)

logger = Logger()

iconcache = {}

def load_icon(name: Text):
    from tortoisehg.util.paths import get_tortoise_icon
    iconPathName = get_tortoise_icon(name)
    if iconPathName and os.path.isfile(iconPathName):
        icon_flags = win32con.LR_LOADFROMFILE | win32con.LR_DEFAULTSIZE
        hicon = LoadImage(win32con.NULL, iconPathName, win32con.IMAGE_ICON, 16, 16, icon_flags)
    else:
        print("Can't find a Python icon file - using default")
        hicon = LoadIcon(0, win32con.IDI_APPLICATION)
    return hicon

def get_icon(name: Text):
    try:
        return iconcache[name]
    except KeyError:
        iconcache[name] = load_icon(name)
        return iconcache[name]

def SetIcon(hwnd, name: Text, add: bool = False) -> None:
    # Try and find a custom icon
    if '--noicon' in sys.argv:
        return
    print("SetIcon(%s)" % name)

    hicon = get_icon(name)

    flags = NIF_ICON | NIF_MESSAGE | NIF_TIP
    nid = (hwnd, 0, flags, win32con.WM_USER+20, hicon, APP_TITLE)
    action = NIM_MODIFY
    if add:
        action = NIM_ADD
    try:
        Shell_NotifyIcon(action, nid)
    except pywintypes.error:
        # This is common when windows is starting, and this code is hit
        # before the taskbar has been created.
        print("Failed to add the taskbar icon - is explorer running?")
        # but keep running anyway - when explorer starts, we get the
        # TaskbarCreated message.

class MainWindow:
    def __init__(self):
        self.pipethread = None
        msg_TaskbarRestart = RegisterWindowMessage("TaskbarCreated")
        message_map = {
                msg_TaskbarRestart: self.OnRestart,
                win32con.WM_DESTROY: self.OnDestroy,
                win32con.WM_COMMAND: self.OnCommand,
                win32con.WM_USER+20 : self.OnTaskbarNotify,
        }
        # Register the Window class.
        wc = WNDCLASS()
        hinst = wc.hInstance = GetModuleHandle(None)
        wc.lpszClassName = "THgRpcServer"
        wc.style = win32con.CS_VREDRAW | win32con.CS_HREDRAW
        wc.hCursor = LoadCursor( 0, win32con.IDC_ARROW )
        wc.hbrBackground = win32con.COLOR_WINDOW
        wc.lpfnWndProc = message_map # could also specify a wndproc.
        classAtom = RegisterClass(wc)
        # Create the Window.
        style = win32con.WS_OVERLAPPED | win32con.WS_SYSMENU
        self.hwnd = CreateWindow( classAtom, APP_TITLE, style,
                0, 0, win32con.CW_USEDEFAULT, win32con.CW_USEDEFAULT,
                0, 0, hinst, None)
        UpdateWindow(self.hwnd)
        self._DoCreateIcons()

    def _DoCreateIcons(self):
        show, highlight = get_config()
        if show:
            SetIcon(self.hwnd, "hg.ico", add=True)
        # start namepipe server for hg status
        self.start_pipe_server()

    def OnRestart(self, hwnd, msg, wparam, lparam):
        logger.msg("MainWindow.OnRestart")
        self._DoCreateIcons()

    def OnDestroy(self, hwnd, msg, wparam, lparam):
        logger.msg("MainWindow.OnDestroy")
        nid = (self.hwnd, 0)
        try:
            Shell_NotifyIcon(NIM_DELETE, nid)
        except pywintypes.error:
            pass # happens when we run without icon
        PostQuitMessage(0) # Terminate the app.

    def OnTaskbarNotify(self, hwnd, msg, wparam, lparam):
        if lparam==win32con.WM_RBUTTONUP or lparam==win32con.WM_LBUTTONUP:
            menu = CreatePopupMenu()
            # AppendMenu(menu, win32con.MF_SEPARATOR, 0, '')
            AppendMenu(menu, win32con.MF_STRING, EXIT_CMD, _('Exit'))
            pos = GetCursorPos()
            # See http://msdn.microsoft.com/library/default.asp?url=/library/en-us/winui/menus_0hdi.asp
            try:
                SetForegroundWindow(self.hwnd)
                TrackPopupMenu(menu, win32con.TPM_LEFTALIGN, pos[0], pos[1], 0, self.hwnd, None)
                PostMessage(self.hwnd, win32con.WM_NULL, 0, 0)
            except pywintypes.error:
                pass
        return 1

    def OnCommand(self, hwnd, msg, wparam, lparam):
        id = LOWORD(wparam)
        if id == EXIT_CMD:
            self.exit_application()
        else:
            print("Unknown command -", id)

    def exit_application(self):
        logger.msg("MainWindow.exit_application")
        if self.stop_pipe_server():
            DestroyWindow(self.hwnd)
        logger.msg("Goodbye")

    def stop_pipe_server(self):
        logger.msg("MainWindow.stop_pipe_server")
        if not self.pipethread.is_alive():
            logger.msg("pipethread is not alive")
            return True

        # Try the nice way first
        self.svc.SvcStop()

        max_try = 10
        cnt = 1
        while cnt <= max_try and self.pipethread.is_alive():
            print("testing pipe [try %d] ..." % cnt)
            try:
                try:
                    self.pipethread.terminate()
                except ValueError:
                    pass
                win32pipe.CallNamedPipe(PIPENAME, b'', PIPEBUFSIZE, 0)
            except:
                logger.msg(traceback.format_exc())
                pass
            cnt += 1

        if self.pipethread.is_alive():
            msg = "WARNING: unable to stop server after %d trys." % max_try
            logger.msg(msg)
            return False
        else:
            return True

    def start_pipe_server(self):
        if self.pipethread is not None:
            return

        def servepipe():
            self.svc = PipeServer(self.hwnd)
            self.svc.SvcDoRun()

        self.pipethread = thread2.Thread(target=servepipe)
        self.pipethread.start()


PIPENAME = r"\\.\pipe\TortoiseHgRpcServer-bc0c27107423-"
PIPENAME += GetUserName()

PIPEBUFSIZE = 4096

def getrepos(batch: Iterable[Text]) -> Tuple[Set[Text], Set[Text]]:
    roots = set()
    notifypaths = set()
    for path in batch:
        r = paths.find_root(path)
        if r is None:
            try:
                for n in os.listdir(path):
                    r = paths.find_root(os.path.join(path, n))
                    if r is not None:
                        roots.add(r)
                        notifypaths.add(r)
            except Exception as e:
                # This exception raises in case of fixutf8 extension enabled
                # and folder name contains '0x5c'(backslash).
                logger.msg('Failed listdir %s (%s)' % (path, hglib.exception_str(e)))
        else:
            roots.add(r)
            notifypaths.add(path)
    return roots, notifypaths

def update_batch(batch: Iterable[Text]) -> None:
    '''updates thgstatus for all paths in batch'''
    roots, notifypaths = getrepos(batch)
    if roots:
        _ui = hglib.loadui()
        failedroots = set()
        _stderr = sys.stderr

        # TODO: Figure out if this hack for py2exe is still needed, and possibly
        #       wrap this to handle unicode on py3
        if not pycompat.ispy3:
            errorstream = util.bytesio()
            sys.stderr = errorstream
        try:
            # Ensure that all unset dirstate entries can be updated.
            time.sleep(2)
            updated_any = False
            for r in sorted(roots):
                try:
                    if shlib.update_thgstatus(
                        _ui, hglib.fromunicode(r), wait=False
                    ):
                        updated_any = True
                    shlib.shell_notify([hglib.fromunicode(r)], noassoc=True)
                    logger.msg('Updated ' + r)
                except (IOError, OSError):
                    print("IOError or OSError on updating %s (check permissions)" % r)
                    logger.msg('Failed updating %s (check permissions)' % r)
                    failedroots.add(r)
                except (error.Abort, error.ConfigError, error.RepoError,
                        error.RevlogError, ImportError) as e:
                    logger.msg('Failed updating %s (%s)' % (r, hglib.exception_str(e)))
                    failedroots.add(r)
            notifypaths -= failedroots
            if notifypaths:
                shlib.shell_notify(
                    [hglib.fromunicode(p) for p in notifypaths],
                    noassoc=not updated_any
                )
                logger.msg('Shell notified')
            if not pycompat.ispy3:
                errmsg = errorstream.getvalue()
                if errmsg:
                    logger.msg('stderr: %s' % errmsg)
        finally:
            sys.stderr = _stderr

requests = pycompat.queue.Queue(0)

def get_config():
    show_taskbaricon = True
    hgighlight_taskbaricon = True
    version2cmenu = False
    try:
        from mercurial.windows import winreg
        HKEY_CURRENT_USER = winreg.HKEY_CURRENT_USER
        OpenKey = winreg.OpenKey
        QueryValueEx = winreg.QueryValueEx

        with OpenKey(HKEY_CURRENT_USER, r'Software\TortoiseHg') as hkey:
            t = ('1', 'True')
            try: show_taskbaricon = QueryValueEx(hkey, 'ShowTaskbarIcon')[0] in t
            except OSError: pass
            try: hgighlight_taskbaricon = QueryValueEx(hkey, 'HighlightTaskbarIcon')[0] in t
            except OSError: pass

            # Upgrade user's context menu, once per major release
            try: version2cmenu = QueryValueEx(hkey, 'ContextMenuVersion')[0] == '2'
            except OSError: pass

            try:
                if not version2cmenu:
                    CreateKey = winreg.CreateKey
                    SetValueEx = winreg.SetValueEx
                    REG_SZ = winreg.REG_SZ

                    try: promoted = QueryValueEx(hkey, 'PromotedItems')[0]
                    except OSError: promoted = ''
                    plist = [i.strip() for i in promoted.split(',')]

                    with CreateKey(HKEY_CURRENT_USER, r'Software\TortoiseHg') as hkey2:
                        if u'log' in plist:
                            idx = plist.index(u'log')
                            plist[idx] = u'workbench'
                            SetValueEx(hkey2, 'PromotedItems', 0, REG_SZ, ','.join(plist))
                        SetValueEx(hkey2, 'ContextMenuVersion', 0, REG_SZ, '2')
            except OSError:
                pass
    except (ImportError, OSError):
        pass
    return show_taskbaricon, hgighlight_taskbaricon

def update(args: List[Text], hwnd) -> None:
    batch = []
    r = args[0]
    print("got update request %s (first in batch)" % r)
    batch.append(r)
    print("wait a bit for additional requests...")
    show, highlight = get_config()
    if show and highlight:
        SetIcon(hwnd, "hgB.ico")
    time.sleep(0.2)
    deferred_requests = []
    try:
        while True:
            req = requests.get_nowait()
            s = req.split('|')
            cmd, args = s[0], s[1:]
            if cmd == 'update':
                print("got update request %s" % req)
                batch.append(args[0])
            else:
                deferred_requests.append(req)
    except pycompat.queue.Empty:
        pass
    for req in deferred_requests:
        requests.put(req)
    msg = "processing batch with %i update requests"
    print(msg % len(batch))
    update_batch(batch)
    if show and highlight:
        SetIcon(hwnd, "hg.ico")

def remove(args: List[Text]) -> None:
    path = args[0]
    logger.msg('Removing ' + path)
    roots, notifypaths = getrepos([path])
    if roots:
        for r in sorted(roots):
            tfn = os.path.join(r, '.hg', 'thgstatus')
            fn = hglib.fromunicode(tfn)
            try:
                f = posixfile(fn, b'rb')
                e = f.readline()
                f.close()
                if not e.startswith(b'@@noicons'):
                    unlink(fn)
            except (IOError, OSError) as e:
                if e.errno != errno.ENOENT:
                    logger.msg("Error while trying to remove %s (%s)" % (tfn, e))
        if notifypaths:
            shlib.shell_notify([hglib.fromunicode(p) for p in notifypaths])

def dispatch(req, cmd: Text, args: List[Text], hwnd) -> None:
    print("dispatch(%s)" % req)
    if cmd == 'update':
        update(args, hwnd)
    elif cmd == 'remove':
        remove(args)
    elif cmd == 'error':
        logger.msg("**** Error: %s" % args[0])
    else:
        logger.msg("**** Error: unknown request '%s'" % req)

class Updater(threading.Thread):
    def __init__(self, hwnd):
        threading.Thread.__init__(self)
        self.hwnd = hwnd

    def run(self):
        while True:
            req = requests.get()
            s = req.split('|')
            cmd, args = s[0], s[1:]
            if cmd == 'terminate':
                logger.msg('Updater thread terminating')
                return
            dispatch(req, cmd, args, self.hwnd)
            gc.collect()

class PipeServer:
    def __init__(self, hwnd):
        self.hwnd = hwnd
        self.updater = Updater(hwnd)
        self.updater.start()

        # Create an event which we will use to wait on.
        # The "service stop" request will set this event.
        self.hWaitStop = win32event.CreateEvent(None, 0, 0, None)

        # We need to use overlapped IO for this, so we dont block when
        # waiting for a client to connect.  This is the only effective way
        # to handle either a client connection, or a service stop request.
        self.overlapped = pywintypes.OVERLAPPED()

        # And create an event to be used in the OVERLAPPED object.
        self.overlapped.hEvent = win32event.CreateEvent(None,0,0,None)

    def SvcStop(self):
        logger.msg("PipeServer.SvcStop")
        win32event.SetEvent(self.hWaitStop)
        requests.put('terminate|')

    def SvcDoRun(self):
        logger.msg("PipeServer.SvcDoRun")
        # We create our named pipe.
        pipeName = PIPENAME
        openMode = win32pipe.PIPE_ACCESS_DUPLEX | win32file.FILE_FLAG_OVERLAPPED
        pipeMode = win32pipe.PIPE_TYPE_MESSAGE

        # When running as a service, we must use special security for the pipe
        sa = pywintypes.SECURITY_ATTRIBUTES()
        # Say we do have a DACL, and it is empty
        # (ie, allow full access!)
        sa.SetSecurityDescriptorDacl ( 1, None, 0 )

        pipeHandle = win32pipe.CreateNamedPipe(pipeName,
            openMode,
            pipeMode,
            win32pipe.PIPE_UNLIMITED_INSTANCES,
            0, 0, 6000, # default buffers, and 6 second timeout.
            sa)

        # Loop accepting and processing connections
        while True:
            try:
                hr = win32pipe.ConnectNamedPipe(pipeHandle, self.overlapped)
            except pywintypes.error as inst:
                logger.msg("Error connecting pipe: %s" % inst)
                pipeHandle.Close()
                break

            if hr==winerror.ERROR_PIPE_CONNECTED:
                # Client is fast, and already connected - signal event
                win32event.SetEvent(self.overlapped.hEvent)
            # Wait for either a connection, or a service stop request.
            timeout = win32event.INFINITE
            waitHandles = self.hWaitStop, self.overlapped.hEvent
            rc = win32event.WaitForMultipleObjects(waitHandles, 0, timeout)
            if rc==win32event.WAIT_OBJECT_0:
                # Stop event
                return
            else:
                # read pipe and process request
                try:
                    hr, data = win32file.ReadFile(pipeHandle, PIPEBUFSIZE)
                    if not data:
                        raise SystemExit  # signal by dispatch terminate
                    win32pipe.DisconnectNamedPipe(pipeHandle)
                except win32file.error:
                    # Client disconnected without sending data
                    # or before reading the response.
                    # Thats OK - just get the next connection
                    continue

                try:
                    requests.put(hglib.tounicode(data))
                    if data == 'terminate|':
                        logger.msg('PipeServer received terminate from pipe')
                        PostMessage(self.hwnd, win32con.WM_COMMAND, EXIT_CMD, 0)
                        break
                except SystemExit:
                    raise SystemExit # interrupted by thread2.terminate()
                except:
                    logger.msg("WARNING: something went wrong in requests.put")
                    logger.msg(traceback.format_exc())
                    status = "ERROR"
        # Clean up when we exit
        self.SvcStop()

RUNMUTEXNAME = 'thgtaskbar-' + GetUserName()

def ehook(etype, values, tracebackobj):
    elist = traceback.format_exception(etype, values, tracebackobj)
    logger.msg(''.join(elist))

def main():
    args = sys.argv[1:]
    sa = win32security.SECURITY_ATTRIBUTES()
    sa.SetSecurityDescriptorDacl(1, None, 0) # allow full access
    runmutex = win32event.CreateMutex(sa, 1, RUNMUTEXNAME)
    if GetLastError() == winerror.ERROR_ALREADY_EXISTS:
        print("another instance is already running")
        return

    logfilename = None
    for arg in args:
        if arg[0] == '-':
            pass
        else:
            logfilename = arg
    if logfilename:
        logger.setfile(logfilename)
    else:
        try:
            from win32com.shell import shell, shellcon
            appdir = shell.SHGetSpecialFolderPath(0, shellcon.CSIDL_APPDATA)
        except pywintypes.com_error:
            appdir = os.environ['APPDATA']
        logfilename = os.path.join(appdir, 'TortoiseHg', 'OverlayServerLog.txt')
        try:
            os.makedirs(os.path.dirname(logfilename))
        except EnvironmentError:
            pass
        logger.setfile(logfilename)

    sys.excepthook = ehook

    w=MainWindow()
    PumpMessages()

if __name__=='__main__':
    main()
