#lang racket/base

(require ffi/unsafe
         ffi/unsafe/define
         ffi/unsafe/os-async-channel
         ffi/winapi
         racket/class
         racket/draw
         racket/gui/base
         racket/match)

(provide mk-tray
         tray-close
         tray-set-icon!
         tray-set-menu!)

;; Shell_NotifyIcon constants
(define NIM_ADD        #x00000000)
(define NIM_MODIFY     #x00000001)
(define NIM_DELETE     #x00000002)
(define NIM_SETVERSION #x00000004)

(define NIF_MESSAGE #x00000001)
(define NIF_ICON    #x00000002)
(define NIF_TIP     #x00000004)
(define NIF_SHOWTIP #x00000080)

(define NOTIFYICON_VERSION_4 4)

;; Window messages used by NOTIFYICON_VERSION_4.
(define WM_CONTEXTMENU #x007B)
(define WM_NCDESTROY   #x0082)
(define WM_USER        #x0400)
(define WM_APP         #x8000)
(define NIN_SELECT     (+ WM_USER 0))
(define NIN_KEYSELECT  (+ WM_USER 1))

(define IMAGE_ICON       1)
(define LR_LOADFROMFILE  #x00000010)
(define SM_CXSMICON      49)
(define SM_CYSMICON      50)

(define BI_RGB         0)
(define DIB_RGB_COLORS 0)

;; Windows scalar aliases. Pointer-sized values matter on 64-bit Windows.
(define _HWND      _pointer)
(define _HICON     _pointer)
(define _UINT_PTR  _uintptr)
(define _DWORD_PTR _uintptr)
(define _WPARAM    _uintptr)
(define _LPARAM    _intptr)
(define _LRESULT   _intptr)
(define _HBITMAP   _pointer)

(define-cstruct _BITMAPINFOHEADER
  ([biSize          _uint32]
   [biWidth         _int32]
   [biHeight        _int32]
   [biPlanes        _uint16]
   [biBitCount      _uint16]
   [biCompression   _uint32]
   [biSizeImage     _uint32]
   [biXPelsPerMeter _int32]
   [biYPelsPerMeter _int32]
   [biClrUsed       _uint32]
   [biClrImportant  _uint32]))

(define-cstruct _ICONINFO
  ([fIcon    _int32]
   [xHotspot _uint32]
   [yHotspot _uint32]
   [hbmMask  _HBITMAP]
   [hbmColor _HBITMAP]))

;; NOTIFYICONDATAW as defined by shellapi.h. The uTimeout/uVersion union is
;; represented by one UINT because this package only uses uVersion.
(define-cstruct _NOTIFYICONDATAW
  ([cbSize           _uint32]
   [hWnd             _HWND]
   [uID              _uint32]
   [uFlags           _uint32]
   [uCallbackMessage _uint32]
   [hIcon            _HICON]
   [szTip            (_array _uint16 128)]
   [dwState          _uint32]
   [dwStateMask      _uint32]
   [szInfo           (_array _uint16 256)]
   [uVersion         _uint32]
   [szInfoTitle      (_array _uint16 64)]
   [dwInfoFlags      _uint32]
   [guidItem         (_array _uint8 16)]
   [hBalloonIcon     _HICON]))

(define shell32  (ffi-lib "shell32.dll"))
(define user32   (ffi-lib "user32.dll"))
(define gdi32    (ffi-lib "gdi32.dll"))
(define comctl32 (ffi-lib "comctl32.dll"))

(define-ffi-definer define-shell32 shell32)
(define-ffi-definer define-user32 user32)
(define-ffi-definer define-gdi32 gdi32)
(define-ffi-definer define-comctl32 comctl32)

(define-shell32 Shell_NotifyIconW
  (_fun #:abi winapi
        _uint32 _NOTIFYICONDATAW-pointer
        -> _int32))

(define-user32 LoadImageW
  (_fun #:abi winapi
        _pointer _string/utf-16 _uint32 _int32 _int32 _uint32
        -> _pointer))

(define-user32 DestroyIcon
  (_fun #:abi winapi _HICON -> _int32))

(define-user32 CreateIconIndirect
  (_fun #:abi winapi _ICONINFO-pointer -> _HICON))

(define-user32 GetSystemMetrics
  (_fun #:abi winapi _int32 -> _int32))

(define-gdi32 CreateDIBSection
  (_fun #:abi winapi
        _pointer _BITMAPINFOHEADER-pointer _uint32 _pointer _pointer _uint32
        -> _HBITMAP))

(define-gdi32 CreateBitmap
  (_fun #:abi winapi
        _int32 _int32 _uint32 _uint32 _pointer
        -> _HBITMAP))

(define-gdi32 DeleteObject
  (_fun #:abi winapi _pointer -> _int32))

(define _SUBCLASSPROC
  (_fun #:abi winapi
        _HWND _uint32 _WPARAM _LPARAM _UINT_PTR _DWORD_PTR
        -> _LRESULT))

(define-comctl32 SetWindowSubclass
  (_fun #:abi winapi
        _HWND _pointer _UINT_PTR _DWORD_PTR
        -> _int32))

(define-comctl32 RemoveWindowSubclass
  (_fun #:abi winapi
        _HWND _pointer _UINT_PTR
        -> _int32))

(define-comctl32 DefSubclassProc
  (_fun #:abi winapi
        _HWND _uint32 _WPARAM _LPARAM
        -> _LRESULT))

(struct tray (frame
              hwnd
              id
              callback-message
              on-click
              eventspace
              event-channel
              event-thread
              subclass-proc
              icon
              menu
              closed?)
  #:mutable)

(define next-tray-id 1)

(define (allocate-tray-id)
  (define id next-tray-id)
  (when (> id #xffff)
    (error 'mk-tray "too many tray icons have been created in this process"))
  (set! next-tray-id (add1 next-tray-id))
  id)

(define (bool-result? v)
  (not (zero? v)))

(define (call-in-eventspace eventspace thunk)
  (define handler-thread (eventspace-handler-thread eventspace))
  (unless handler-thread
    (error 'racket-tray "the frame's eventspace has been shut down"))
  (if (eq? (current-thread) handler-thread)
      (parameterize ([current-eventspace eventspace])
        (thunk))
      (let ([result-channel (make-channel)])
        (parameterize ([current-eventspace eventspace])
          (queue-callback
           (λ ()
             (with-handlers ([exn?
                              (λ (exn)
                                (channel-put result-channel (cons 'error exn)))])
               (channel-put result-channel (cons 'ok (thunk)))))))
        (match (channel-get result-channel)
          [(cons 'ok value) value]
          [(cons 'error exn) (raise exn)]))))

(define (make-notify-data hwnd id callback-message icon)
  (define data
    (cast (malloc _NOTIFYICONDATAW 'atomic)
          _pointer
          _NOTIFYICONDATAW-pointer))
  (memset data 0 0 (ctype-sizeof _NOTIFYICONDATAW))
  (set-NOTIFYICONDATAW-cbSize! data (ctype-sizeof _NOTIFYICONDATAW))
  (set-NOTIFYICONDATAW-hWnd! data hwnd)
  (set-NOTIFYICONDATAW-uID! data id)
  (set-NOTIFYICONDATAW-uCallbackMessage! data callback-message)
  (set-NOTIFYICONDATAW-hIcon! data icon)
  data)

(define (set-wide-array! array text capacity)
  (define source (cast text _string/utf-16 _pointer))
  (for ([i (in-range capacity)])
    (array-set! array i 0))
  (let loop ([i 0])
    (when (< i (sub1 capacity))
      (let ([code-unit (ptr-ref source _uint16 i)])
        (unless (zero? code-unit)
          (array-set! array i code-unit)
          (loop (add1 i)))))))

(define (load-ico-icon icon-file)
  (define width (GetSystemMetrics SM_CXSMICON))
  (define height (GetSystemMetrics SM_CYSMICON))
  (define icon
    (LoadImageW #f
                (path->string (path->complete-path icon-file))
                IMAGE_ICON
                width
                height
                LR_LOADFROMFILE))
  (unless icon
    (error 'mk-tray "could not load ICO file: ~a" icon-file))
  icon)

(define (png->icon icon-file)
  ;; Racket's drawing library already decodes PNG and preserves its alpha
  ;; channel. Scale it to the same system-small-icon size that is used for
  ;; ICO files, then turn the premultiplied pixels into a native HICON.
  (define source-bitmap
    (read-bitmap icon-file 'png/alpha #f #t))
  (unless (send source-bitmap ok?)
    (error 'mk-tray "could not load PNG file: ~a" icon-file))

  (define width (max 1 (GetSystemMetrics SM_CXSMICON)))
  (define height (max 1 (GetSystemMetrics SM_CYSMICON)))
  (define source-width (send source-bitmap get-width))
  (define source-height (send source-bitmap get-height))

  ;; Keep the aspect ratio and center non-square images in the tray-icon area.
  (define scale
    (min (/ width source-width)
         (/ height source-height)))
  (define draw-width (max 1 (inexact->exact (round (* source-width scale)))))
  (define draw-height (max 1 (inexact->exact (round (* source-height scale)))))
  (define draw-x (quotient (- width draw-width) 2))
  (define draw-y (quotient (- height draw-height) 2))

  (define bitmap (make-bitmap width height #t))
  (define dc (new bitmap-dc% [bitmap bitmap]))
  (send dc draw-bitmap-section-smooth
        source-bitmap
        draw-x
        draw-y
        draw-width
        draw-height
        0
        0
        source-width
        source-height)
  (send dc set-bitmap #f)

  (define pixel-count (* width height))
  (define argb (make-bytes (* pixel-count 4)))
  (send bitmap get-argb-pixels 0 0 width height argb #f #t)

  (define header
    (make-BITMAPINFOHEADER
     (ctype-sizeof _BITMAPINFOHEADER)
     width
     (- height) ; negative means top-down, matching Racket's row order
     1
     32
     BI_RGB
     (* pixel-count 4)
     0
     0
     0
     0))
  (define bits-out (malloc _pointer 'atomic))
  (ptr-set! bits-out _pointer #f)
  (define color-bitmap
    (CreateDIBSection #f header DIB_RGB_COLORS bits-out #f 0))
  (unless color-bitmap
    (error 'mk-tray "CreateDIBSection failed while loading PNG icon: ~a" icon-file))

  (define mask-bitmap #f)
  (define (cleanup-bitmaps!)
    (when mask-bitmap
      (DeleteObject mask-bitmap)
      (set! mask-bitmap #f))
    (when color-bitmap
      (DeleteObject color-bitmap)
      (set! color-bitmap #f)))

  (with-handlers ([exn?
                   (λ (exn)
                     (cleanup-bitmaps!)
                     (raise exn))])
    (define dib-bits (ptr-ref bits-out _pointer))
    (unless dib-bits
      (error 'mk-tray "CreateDIBSection did not return pixel storage for: ~a" icon-file))

    ;; Racket returns A,R,G,B. A Windows 32-bit DIB uses B,G,R,A byte order.
    ;; get-argb-pixels was requested premultiplied because that is the form
    ;; expected by alpha-blended Windows icons.
    (for ([pixel (in-range pixel-count)])
      (define source (* pixel 4))
      (define alpha (bytes-ref argb source))
      (define red (bytes-ref argb (+ source 1)))
      (define green (bytes-ref argb (+ source 2)))
      (define blue (bytes-ref argb (+ source 3)))
      (ptr-set! dib-bits _uint8 source blue)
      (ptr-set! dib-bits _uint8 (+ source 1) green)
      (ptr-set! dib-bits _uint8 (+ source 2) red)
      (ptr-set! dib-bits _uint8 (+ source 3) alpha))

    ;; For modern 32-bit alpha icons the color bitmap carries transparency.
    ;; ICONINFO still requires a same-sized monochrome mask for a color icon.
    (define mask-stride (* 2 (quotient (+ width 15) 16)))
    (define mask-size (* mask-stride height))
    (define mask-bits (malloc mask-size 'atomic))
    (memset mask-bits 0 0 mask-size)
    (set! mask-bitmap
          (CreateBitmap width height 1 1 mask-bits))
    (unless mask-bitmap
      (error 'mk-tray "CreateBitmap failed while loading PNG icon: ~a" icon-file))

    (define icon-info
      (make-ICONINFO 1 0 0 mask-bitmap color-bitmap))
    (define icon (CreateIconIndirect icon-info))

    ;; CreateIconIndirect copies both bitmaps, so the source GDI objects can
    ;; be released immediately. The returned HICON remains owned by us.
    (cleanup-bitmaps!)
    (unless icon
      (error 'mk-tray "CreateIconIndirect failed while loading PNG icon: ~a" icon-file))
    icon))

(define (load-icon icon-file)
  (define name
    (string-downcase
     (path->string (path->complete-path icon-file))))
  (cond
    [(regexp-match? #rx"[.]ico$" name)
     (load-ico-icon icon-file)]
    [(regexp-match? #rx"[.]png$" name)
     (png->icon icon-file)]
    [else
     (error 'mk-tray
            "expected an .ico or .png icon file; got: ~a"
            icon-file)]))

(define (low-word value)
  (bitwise-and value #xffff))

(define (signed-word value)
  (define n (low-word value))
  (if (>= n #x8000)
      (- n #x10000)
      n))

(define (x-from-wparam wparam)
  (signed-word wparam))

(define (y-from-wparam wparam)
  (signed-word (arithmetic-shift wparam -16)))

(define (menu-spec->popup-menu menu-spec)
  (cond
    [(not menu-spec) #f]
    [(is-a? menu-spec popup-menu%) menu-spec]
    [(list? menu-spec)
     (let ([popup (new popup-menu%)])
       (for ([entry (in-list menu-spec)])
         (match entry
           [(or #f 'separator)
            (new separator-menu-item% [parent popup])]
           [(list (? string? label) (? procedure? callback))
            (unless (procedure-arity-includes? callback 0)
              (error 'tray-set-menu!
                     "menu callback for ~e does not accept zero arguments"
                     label))
            (new menu-item%
                 [parent popup]
                 [label label]
                 [callback (λ (_item _event) (callback))])]
           [_
            (error 'tray-set-menu!
                   "expected a popup-menu% or a list containing (list label callback), #f, or 'separator; got: ~e"
                   entry)]))
       popup)]
    [else
     (error 'tray-set-menu!
            "expected a popup-menu%, menu specification list, or #f; got: ~e"
            menu-spec)]))

(define (queue-eventspace-callback t thunk)
  (define eventspace (tray-eventspace t))
  (unless (eventspace-shutdown? eventspace)
    (with-handlers ([exn:fail? (λ (_exn) (void))])
      (parameterize ([current-eventspace eventspace])
        (queue-callback
         (λ ()
           (unless (tray-closed? t)
             (thunk))))))))

(define (start-event-thread! t)
  (define event-thread
    (thread
     (λ ()
       (let loop ()
         (match (sync (tray-event-channel t))
           ['close
            (void)]
           ['activate
            (let ([callback (tray-on-click t)])
              (when callback
                (queue-eventspace-callback t callback)))
            (loop)]
           [(vector 'context-menu screen-x screen-y)
            (queue-eventspace-callback
             t
             (λ ()
               (define menu (tray-menu t))
               (when menu
                 (let-values ([(x y)
                               (send (tray-frame t)
                                     screen->client
                                     screen-x
                                     screen-y)])
                   (send (tray-frame t) popup-menu menu x y)))))
            (loop)]
           [_
            (loop)])))))
  (set-tray-event-thread! t event-thread))

(define (handle-native-tray-event t wparam lparam)
  ;; Racket CS evaluates foreign callbacks in atomic mode. Do not run GUI or
  ;; user code here. An OS async channel is explicitly safe to use from an OS
  ;; callback/thread, so forward the event to an ordinary Racket thread.
  (define notification (low-word lparam))
  (cond
    [(or (= notification NIN_SELECT)
         (= notification NIN_KEYSELECT))
     (os-async-channel-put (tray-event-channel t) 'activate)]
    [(= notification WM_CONTEXTMENU)
     (os-async-channel-put
      (tray-event-channel t)
      (vector 'context-menu
              (x-from-wparam wparam)
              (y-from-wparam wparam)))]
    [else
     (void)]))

(define (make-subclass-proc t)
  (λ (hwnd msg wparam lparam subclass-id ref-data)
    (cond
      [(= msg (tray-callback-message t))
       (handle-native-tray-event t wparam lparam)
       0]
      [(= msg WM_NCDESTROY)
       ;; The HWND is going away. Remove the notification-area icon while the
       ;; handle is still valid. Windows discards the subclass automatically
       ;; as part of window destruction.
       (unless (tray-closed? t)
         (let ([data
                (make-notify-data hwnd
                                  (tray-id t)
                                  (tray-callback-message t)
                                  (tray-icon t))])
           (Shell_NotifyIconW NIM_DELETE data)
           (when (tray-icon t)
             (DestroyIcon (tray-icon t))
             (set-tray-icon! t #f))
           (set-tray-closed?! t #t)
           (os-async-channel-put (tray-event-channel t) 'close)))
       (DefSubclassProc hwnd msg wparam lparam)]
      [else
       (DefSubclassProc hwnd msg wparam lparam)])))

(define (add-notify-icon! t icon tooltip)
  (define data
    (make-notify-data (tray-hwnd t)
                      (tray-id t)
                      (tray-callback-message t)
                      icon))
  (set-NOTIFYICONDATAW-uFlags!
   data
   (bitwise-ior NIF_MESSAGE NIF_ICON NIF_TIP NIF_SHOWTIP))
  (set-wide-array! (NOTIFYICONDATAW-szTip data) tooltip 128)

  (unless (bool-result? (Shell_NotifyIconW NIM_ADD data))
    (error 'mk-tray "Shell_NotifyIconW failed to add the tray icon"))

  (set-NOTIFYICONDATAW-uVersion! data NOTIFYICON_VERSION_4)
  (unless (bool-result? (Shell_NotifyIconW NIM_SETVERSION data))
    (Shell_NotifyIconW NIM_DELETE data)
    (error 'mk-tray "Shell_NotifyIconW could not enable NOTIFYICON_VERSION_4")))

(define (mk-tray frame icon-file on-click-cb)
  (unless (is-a? frame top-level-window<%>)
    (raise-argument-error 'mk-tray "(is-a?/c top-level-window<%>)" frame))
  (unless (or (not on-click-cb)
              (and (procedure? on-click-cb)
                   (procedure-arity-includes? on-click-cb 0)))
    (raise-argument-error 'mk-tray "(or/c #f (-> any))" on-click-cb))

  (define eventspace (send frame get-eventspace))
  (call-in-eventspace
   eventspace
   (λ ()
     (define hwnd (send frame get-handle))
     (unless hwnd
       (error 'mk-tray "the frame does not have a native HWND"))

     (let* ([id (allocate-tray-id)]
            [callback-message (+ WM_APP id)]
            [icon (load-icon icon-file)]
            [t (tray frame
                     hwnd
                     id
                     callback-message
                     on-click-cb
                     eventspace
                     (make-os-async-channel)
                     #f
                     #f
                     icon
                     #f
                     #f)]
            [subclass-proc
             (function-ptr (make-subclass-proc t) _SUBCLASSPROC)])
       ;; WM_APP through 0xBFFF is reserved for application-private messages.
       (when (> callback-message #xbfff)
         (DestroyIcon icon)
         (error 'mk-tray
                "too many simultaneously addressable tray callback messages"))

       (set-tray-subclass-proc! t subclass-proc)
       (start-event-thread! t)

       (unless (bool-result?
                (SetWindowSubclass hwnd subclass-proc id 0))
         (os-async-channel-put (tray-event-channel t) 'close)
         (DestroyIcon icon)
         (error 'mk-tray "SetWindowSubclass failed for the Racket frame"))

       (with-handlers ([exn?
                        (λ (exn)
                          (RemoveWindowSubclass hwnd subclass-proc id)
                          (os-async-channel-put (tray-event-channel t) 'close)
                          (DestroyIcon icon)
                          (raise exn))])
         (define frame-label (send frame get-label))
         (add-notify-icon! t icon
                           (if (string? frame-label) frame-label "Racket"))
         t)))))

(define (check-open-tray who t)
  (unless (tray? t)
    (raise-argument-error who "tray?" t))
  (when (tray-closed? t)
    (error who "tray icon is already closed")))

(define (tray-close t)
  (unless (tray? t)
    (raise-argument-error 'tray-close "tray?" t))
  (unless (tray-closed? t)
    (call-in-eventspace
     (tray-eventspace t)
     (λ ()
       (unless (tray-closed? t)
         (let ([data
                (make-notify-data (tray-hwnd t)
                                  (tray-id t)
                                  (tray-callback-message t)
                                  (tray-icon t))])
           (Shell_NotifyIconW NIM_DELETE data)
           (RemoveWindowSubclass (tray-hwnd t)
                                 (tray-subclass-proc t)
                                 (tray-id t))
           (when (tray-icon t)
             (DestroyIcon (tray-icon t))
             (set-tray-icon! t #f))
           (set-tray-closed?! t #t)
           (os-async-channel-put (tray-event-channel t) 'close))))))
  (void))

(define (tray-set-icon! t icon-file)
  (check-open-tray 'tray-set-icon! t)
  (call-in-eventspace
   (tray-eventspace t)
   (λ ()
     (define new-icon (load-icon icon-file))
     (define data
       (make-notify-data (tray-hwnd t)
                         (tray-id t)
                         (tray-callback-message t)
                         new-icon))
     (set-NOTIFYICONDATAW-uFlags! data NIF_ICON)
     (cond
       [(bool-result? (Shell_NotifyIconW NIM_MODIFY data))
        (let ([old-icon (tray-icon t)])
          (set-tray-icon! t new-icon)
          (when old-icon
            (DestroyIcon old-icon)))]
       [else
        (DestroyIcon new-icon)
        (error 'tray-set-icon! "Shell_NotifyIconW failed to change the icon")]))))

(define (tray-set-menu! t menu-spec)
  (check-open-tray 'tray-set-menu! t)
  (call-in-eventspace
   (tray-eventspace t)
   (λ ()
     (set-tray-menu! t (menu-spec->popup-menu menu-spec))))
  (void))