;;; claude-code-mcp-connection.el --- MCP WebSocket connection management -*- lexical-binding: t; -*-

;; Copyright (C) 2025

;; Author: DESKTOP2 <yuya373@DESKTOP2>
;; Keywords: tools, convenience

;; This program 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.

;; This program 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 this program.  If not, see <https://www.gnu.org/licenses/>.

;;; Commentary:

;; This module handles MCP WebSocket connection management including:
;; - Per-server-instance connection tracking
;; - Port discovery and registration
;; - Connection retry logic
;; - WebSocket lifecycle management
;; - Health monitoring via ping/pong heartbeat
;;
;; Connections are keyed by a server instance ID (a UUID generated by
;; each MCP server process) instead of the project root.  Claude Code
;; starts one MCP server process per session, so a single project can
;; have several live connections at once — one per running agent.
;; Keying by instance ID keeps their lifecycles independent: one
;; agent's shutdown must never tear down another agent's connection.

;;; Code:

(require 'websocket nil t)
(require 'projectile)
(require 'json)
(require 'claude-code-core)

;; Declare websocket functions to avoid eager macro-expansion failures
(declare-function websocket-open "websocket" (url &rest args))
(declare-function websocket-send-text "websocket" (websocket text))
(declare-function websocket-close "websocket" (websocket))
(declare-function websocket-openp "websocket" (websocket))

;; Forward declarations
(declare-function claude-code-mcp-on-message "claude-code-mcp-protocol" (_websocket frame instance-id))
(declare-function claude-code-mcp-on-error "claude-code-mcp-protocol" (_websocket type error &optional _instance-id))
(declare-function claude-code-mcp-on-close "claude-code-mcp-protocol" (websocket instance-id))

;;; Customization

(defgroup claude-code-mcp nil
  "MCP server integration for Claude Code Emacs."
  :group 'claude-code
  :prefix "claude-code-mcp-")

(defcustom claude-code-mcp-host "localhost"
  "Host for MCP server."
  :type 'string
  :group 'claude-code-mcp)

(defcustom claude-code-mcp-max-connection-attempts 10
  "Maximum number of connection attempts."
  :type 'integer
  :group 'claude-code-mcp)

(defcustom claude-code-mcp-connection-retry-delay 5
  "Delay in seconds between connection attempts."
  :type 'number
  :group 'claude-code-mcp)

(defcustom claude-code-mcp-ping-interval 30
  "Interval in seconds between WebSocket ping messages."
  :type 'integer
  :group 'claude-code-mcp)

(defcustom claude-code-mcp-ping-timeout 10
  "Timeout in seconds to wait for pong response."
  :type 'integer
  :group 'claude-code-mcp)

;;; Variables

(defvar claude-code-mcp-connections (make-hash-table :test 'equal)
  "Hash table mapping server instance IDs to connection info.
Each MCP server process generates a unique instance ID, so multiple
Claude Code sessions in the same project each get their own entry.
Each value is an alist with keys:
  - project-root: Normalized project root the server runs in
  - port: WebSocket port of the server's Emacs bridge
  - websocket: The WebSocket connection
  - request-id: Counter for JSON-RPC request IDs
  - pending-requests: Hash table of pending requests
  - connection-attempts: Number of connection attempts
  - ping-timer: Timer for periodic ping messages
  - ping-timeout-timer: Timer for ping timeout detection
  - last-pong-time: Time of last received pong")

;;; Connection info management

(defun claude-code-mcp-initialize-connection-info (instance-id project-root port)
  "Initialize and return connection info for INSTANCE-ID.
PROJECT-ROOT is normalized before being stored.  PORT is the WebSocket
port of the server's Emacs bridge."
  (let* ((normalized-root (claude-code-normalize-project-root project-root))
         (info (list (cons 'project-root normalized-root)
                     (cons 'port port)
                     (cons 'websocket nil)
                     (cons 'request-id 0)
                     (cons 'pending-requests (make-hash-table :test 'equal))
                     (cons 'connection-attempts 0)
                     (cons 'ping-timer nil)
                     (cons 'ping-timeout-timer nil)
                     (cons 'last-pong-time nil))))
    (puthash instance-id info claude-code-mcp-connections)
    info))

(defun claude-code-mcp-get-connection-info (instance-id)
  "Get connection info for INSTANCE-ID.
Returns nil if no connection info exists for the instance."
  (gethash instance-id claude-code-mcp-connections))

(defun claude-code-mcp-get-project-root (instance-id)
  "Get the normalized project root for INSTANCE-ID."
  (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
    (cdr (assoc 'project-root info))))

(defun claude-code-mcp-get-websocket (instance-id)
  "Get WebSocket for INSTANCE-ID."
  (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
    (cdr (assoc 'websocket info))))

(defun claude-code-mcp-set-websocket (websocket instance-id)
  "Set WEBSOCKET for INSTANCE-ID."
  (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
    (setcdr (assoc 'websocket info) websocket)))

;;; Port Registration

(defun claude-code-mcp-register-port (project-root port instance-id)
  "Register PORT for the MCP server INSTANCE-ID running in PROJECT-ROOT.
Called by the MCP server process (via emacsclient) after it starts its
WebSocket bridge.  Other instances of the same project are untouched.
If INSTANCE-ID is already registered (server restart with the same ID),
its previous connection is torn down first."
  (let ((normalized-root (claude-code-normalize-project-root project-root)))
    (when (claude-code-mcp-get-connection-info instance-id)
      (claude-code-mcp-disconnect instance-id))
    (claude-code-mcp-initialize-connection-info instance-id normalized-root port)
    (message "MCP server %s registered on port %d for project %s"
             instance-id port normalized-root)
    (claude-code-mcp-try-connect-async instance-id)))

(defun claude-code-mcp-unregister-port (instance-id)
  "Unregister the MCP server INSTANCE-ID and disconnect.
Called when that server process shuts down.  Only the instance's own
connection is removed; other instances keep their connections."
  (claude-code-mcp-disconnect instance-id))

;;; Connection Management

(defun claude-code-mcp-try-connect-async (instance-id)
  "Try to connect to MCP server INSTANCE-ID asynchronously.
Aborts silently when the instance has been unregistered in the
meantime (e.g. the server exited between retries)."
  (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
    (let ((attempts (cdr (assoc 'connection-attempts info)))
          (project-root (cdr (assoc 'project-root info))))
      (if (>= attempts claude-code-mcp-max-connection-attempts)
          (message "Failed to connect to MCP server %s after %d attempts for project %s"
                   instance-id claude-code-mcp-max-connection-attempts project-root)
        (setcdr (assoc 'connection-attempts info) (1+ attempts))
        (message "Attempting to connect to MCP server %s (attempt %d/%d) for project %s..."
                 instance-id
                 (1+ attempts)
                 claude-code-mcp-max-connection-attempts
                 project-root)
        (claude-code-mcp-connect
         instance-id
         (lambda (connected)
           (unless connected
             ;; Connection failed, retry
             (run-at-time claude-code-mcp-connection-retry-delay nil
                          #'claude-code-mcp-try-connect-async instance-id))))))))

(defun claude-code-mcp-connect (instance-id &optional callback)
  "Connect to MCP server WebSocket for INSTANCE-ID.
Project root and port are taken from the instance's connection info.
If CALLBACK is provided, call it with connection result."
  (if-let ((info (claude-code-mcp-get-connection-info instance-id)))
      (let ((project-root (cdr (assoc 'project-root info)))
            (port (cdr (assoc 'port info))))
        (condition-case err
            (progn
              (websocket-open
               (format "ws://%s:%d/?session=%s&instance=%s"
                       claude-code-mcp-host
                       port
                       (url-hexify-string project-root)
                       (url-hexify-string instance-id))
               :on-open (lambda (websocket)
                          (message "MCP WebSocket opened for instance %s (project %s)"
                                   instance-id project-root)
                          (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
                            (setcdr (assoc 'connection-attempts info) 0))
                          (claude-code-mcp-set-websocket websocket instance-id)
                          ;; Start ping timer
                          (claude-code-mcp-start-ping-timer instance-id)
                          (when callback (funcall callback t)))
               :on-message (lambda (websocket frame)
                             (claude-code-mcp-on-message websocket frame instance-id))
               :on-error (lambda (websocket type error)
                           (claude-code-mcp-on-error websocket type error instance-id))
               :on-close (lambda (websocket)
                           (claude-code-mcp-on-close websocket instance-id)))
              (message "Initiating MCP WebSocket connection on port %d for instance %s"
                       port instance-id)
              t)
          (error
           (message "Failed to open WebSocket: %s" err)
           (when callback (funcall callback nil))
           nil)))
    (when callback (funcall callback nil))
    nil))

(defun claude-code-mcp-disconnect (instance-id)
  "Disconnect from MCP server INSTANCE-ID and remove its entry.
Only this instance is affected; other instances of the same project
keep their connections."
  ;; Stop ping timers
  (claude-code-mcp-stop-ping-timer instance-id)
  (claude-code-mcp-stop-ping-timeout instance-id)
  (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
    (let ((websocket (cdr (assoc 'websocket info)))
          (pending-requests (cdr (assoc 'pending-requests info)))
          (project-root (cdr (assoc 'project-root info))))
      (when (and websocket (websocket-openp websocket))
        (websocket-close websocket))
      (when pending-requests
        (clrhash pending-requests))
      (remhash instance-id claude-code-mcp-connections)
      (message "Disconnected from MCP server %s for project %s"
               instance-id project-root))))

;;; Ping/Pong functionality

(defun claude-code-mcp-send-ping (instance-id)
  "Send ping message to MCP server INSTANCE-ID."
  (let ((websocket (claude-code-mcp-get-websocket instance-id)))
    (when (and websocket (websocket-openp websocket))
      (condition-case err
          (progn
            (websocket-send-text websocket "{\"type\":\"ping\"}")
            ;; Set up timeout timer
            (claude-code-mcp-start-ping-timeout instance-id))
        (error
         (message "Error sending ping to MCP server %s: %s" instance-id err)
         (claude-code-mcp-handle-connection-lost instance-id))))))

(defun claude-code-mcp-start-ping-timer (instance-id)
  "Start periodic ping timer for INSTANCE-ID."
  (claude-code-mcp-stop-ping-timer instance-id)
  (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
    (let ((timer (run-with-timer claude-code-mcp-ping-interval
                                 claude-code-mcp-ping-interval
                                 #'claude-code-mcp-send-ping
                                 instance-id)))
      (setcdr (assoc 'ping-timer info) timer))))

(defun claude-code-mcp-stop-ping-timer (instance-id)
  "Stop ping timer for INSTANCE-ID."
  (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
    (let ((timer (cdr (assoc 'ping-timer info))))
      (when (timerp timer)
        (cancel-timer timer))
      (setcdr (assoc 'ping-timer info) nil))))

(defun claude-code-mcp-start-ping-timeout (instance-id)
  "Start ping timeout timer for INSTANCE-ID."
  (claude-code-mcp-stop-ping-timeout instance-id)
  (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
    (let ((timer (run-with-timer claude-code-mcp-ping-timeout
                                 nil
                                 #'claude-code-mcp-handle-ping-timeout
                                 instance-id)))
      (setcdr (assoc 'ping-timeout-timer info) timer))))

(defun claude-code-mcp-stop-ping-timeout (instance-id)
  "Stop ping timeout timer for INSTANCE-ID."
  (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
    (let ((timer (cdr (assoc 'ping-timeout-timer info))))
      (when (timerp timer)
        (cancel-timer timer))
      (setcdr (assoc 'ping-timeout-timer info) nil))))

(defun claude-code-mcp-handle-ping-timeout (instance-id)
  "Handle ping timeout for INSTANCE-ID."
  (message "MCP WebSocket ping timeout for instance %s" instance-id)
  (claude-code-mcp-handle-connection-lost instance-id))

(defun claude-code-mcp-handle-pong (instance-id)
  "Handle pong response for INSTANCE-ID."
  ;; Cancel timeout timer
  (claude-code-mcp-stop-ping-timeout instance-id)
  ;; Update last pong time
  (when-let ((info (claude-code-mcp-get-connection-info instance-id)))
    (setcdr (assoc 'last-pong-time info) (current-time))))

(defun claude-code-mcp-handle-connection-lost (instance-id)
  "Handle lost connection for INSTANCE-ID.
The instance's entry is removed; the server re-registers itself if it
is still alive, and other instances are unaffected."
  (message "MCP WebSocket connection lost for instance %s" instance-id)
  (claude-code-mcp-disconnect instance-id))

;;; Event notification functions

(defun claude-code-mcp-send-event-to-project (project-root event-name params)
  "Send an event notification to every MCP server of PROJECT-ROOT.
A project can have several connected server instances (one per running
Claude Code session), so the event is broadcast to all of them.
EVENT-NAME is the event type (e.g., \"bufferListUpdated\").
PARAMS is an alist of event parameters."
  (let ((normalized-root (claude-code-normalize-project-root project-root))
        (message-text (json-encode
                       `((jsonrpc . "2.0")
                         (method . ,(concat "emacs/" event-name))
                         (params . ,params)))))
    (maphash
     (lambda (instance-id info)
       (let ((websocket (cdr (assoc 'websocket info))))
         (when (and (equal (cdr (assoc 'project-root info)) normalized-root)
                    websocket
                    (websocket-openp websocket))
           (condition-case err
               (websocket-send-text websocket message-text)
             (error
              (message "Error sending event %s to MCP server %s: %s"
                       event-name instance-id err))))))
     claude-code-mcp-connections)))

(provide 'claude-code-mcp-connection)
;;; claude-code-mcp-connection.el ends here
