// Core imports below point at TypeScript source so tsc can type-check them.
// scripts/copy-bundled-plugins.cjs rewrites the prefix to the compiled
// output path after tsc -p tsconfig.plugins.json runs -- see that script.
import { getErrorMessage } from "../../../src/backend/utils/error-message.js";
import { StringDecoder } from "string_decoder";
import { Client as SSHClient } from "ssh2";
import { SSH_ALGORITHMS } from "../../../src/backend/utils/ssh-algorithms.js";
import { WebSocketServer, WebSocket } from "ws";
import { AuthManager } from "../../../src/backend/utils/auth-manager.js";
import { createCurrentHostResolutionRepository } from "../../../src/backend/database/repositories/factory.js";
import { systemLogger } from "../../../src/backend/utils/logger.js";
import type { SSHHost } from "../../../src/types/index.js";
import { applyAgentAuth } from "../../../src/backend/hosts/terminal-auth-helpers.js";
import {
  containerCommand,
  getContainerRuntimeConfig,
  type ContainerRuntime,
} from "./container-runtime.js";
import { resolveSshConnectConfigHost } from "../../../src/backend/hosts/ssh-dns.js";
import {
  hostAddressMismatch,
  HOST_ADDRESS_MISMATCH_MESSAGE,
  HOST_NOT_ON_THIS_SERVER_MESSAGE,
} from "../../../src/backend/hosts/host-identity.js";
import { extractWebSocketToken } from "../../../src/backend/utils/ws-auth.js";
import {
  asObject,
  asString,
  MAX_WS_MESSAGE_BYTES,
  parseWsMessage,
  toTerminalDimension,
} from "../../../src/backend/utils/ws-message.js";

const sshLogger = systemLogger;

interface SSHSession {
  client: SSHClient;
  stream: import("ssh2").ClientChannel | null;
  isConnected: boolean;
  containerId?: string;
  shell?: string;
  hostId?: number;
  containerRuntime?: ContainerRuntime;
}

const activeSessions = new Map<string, SSHSession>();

const wss = new WebSocketServer({
  host: "127.0.0.1",
  port: 30009,
  maxPayload: MAX_WS_MESSAGE_BYTES,
});

wss.on("error", (error) => {
  sshLogger.error("Docker console WebSocket server error", error, {
    operation: "wss_error",
  });
});

async function detectShell(
  session: SSHSession,
  containerId: string,
): Promise<string> {
  const shells = ["bash", "sh", "ash"];

  for (const shell of shells) {
    try {
      await new Promise<void>((resolve, reject) => {
        session.client.exec(
          containerCommand(
            session.containerRuntime,
            `exec ${containerId} which ${shell}`,
          ),
          (err, stream) => {
            if (err) return reject(err);

            let output = "";
            stream.on("data", (data: Buffer) => {
              output += data.toString();
            });

            stream.on("close", (code: number) => {
              if (code === 0 && output.trim()) {
                resolve();
              } else {
                reject(new Error(`Shell ${shell} not found`));
              }
            });

            stream.stderr.on("data", () => {});
            stream.stderr.on("error", () => {});
            stream.on("error", (streamErr) => {
              reject(streamErr);
            });
          },
        );
      });

      return shell;
    } catch {
      continue;
    }
  }

  return "sh";
}

async function createJumpHostChain(
  jumpHosts: Array<{ hostId: number }>,
  userId: string,
): Promise<SSHClient | null> {
  if (!jumpHosts || jumpHosts.length === 0) {
    return null;
  }

  let currentClient: SSHClient | null = null;
  const repository = createCurrentHostResolutionRepository();

  for (let i = 0; i < jumpHosts.length; i++) {
    const jumpHostId = jumpHosts[i].hostId;

    const resolvedJumpHost = await repository.findHostById(jumpHostId, userId);
    if (!resolvedJumpHost || resolvedJumpHost.userId !== userId) {
      throw new Error(`Jump host ${jumpHostId} not found`);
    }

    const jumpHost = resolvedJumpHost as unknown as SSHHost;
    if (typeof jumpHost.jumpHosts === "string" && jumpHost.jumpHosts) {
      try {
        jumpHost.jumpHosts = JSON.parse(jumpHost.jumpHosts);
      } catch (e) {
        sshLogger.error("Failed to parse jump hosts", e, {
          hostId: jumpHost.id,
        });
        jumpHost.jumpHosts = [];
      }
    }

    let resolvedCredentials: {
      password?: string;
      sshKey?: string;
      keyPassword?: string;
      authType?: string;
    } = {
      password: jumpHost.password,
      sshKey: jumpHost.key,
      keyPassword: jumpHost.keyPassword,
      authType: jumpHost.authType,
    };

    if (jumpHost.credentialId) {
      const credential = await repository.findCredentialByIdForUser(
        jumpHost.credentialId as number,
        userId,
      );

      if (credential) {
        resolvedCredentials = {
          password: credential.password as string | undefined,
          sshKey: (credential.key || credential.privateKey) as
            string | undefined,
          keyPassword: credential.keyPassword as string | undefined,
          authType: credential.authType as string | undefined,
        };
      }
    }

    const client = new SSHClient();

    const config: Record<string, unknown> = {
      host: jumpHost.ip?.replace(/^\[|\]$/g, "") || jumpHost.ip,
      port: jumpHost.port || 22,
      username: jumpHost.username,
      tryKeyboard: resolvedCredentials.authType !== "none",
      readyTimeout: 60000,
      keepaliveInterval: 30000,
      keepaliveCountMax: 120,
      tcpKeepAlive: true,
      tcpKeepAliveInitialDelay: 30000,
      algorithms: {
        kex: [
          "curve25519-sha256",
          "curve25519-sha256@libssh.org",
          "ecdh-sha2-nistp521",
          "ecdh-sha2-nistp384",
          "ecdh-sha2-nistp256",
          "diffie-hellman-group-exchange-sha256",
          "diffie-hellman-group18-sha512",
          "diffie-hellman-group17-sha512",
          "diffie-hellman-group16-sha512",
          "diffie-hellman-group15-sha512",
          "diffie-hellman-group14-sha256",
          "diffie-hellman-group14-sha1",
          "diffie-hellman-group-exchange-sha1",
          "diffie-hellman-group1-sha1",
        ],
        serverHostKey: [
          "ssh-ed25519",
          "ecdsa-sha2-nistp521",
          "ecdsa-sha2-nistp384",
          "ecdsa-sha2-nistp256",
          "rsa-sha2-512",
          "rsa-sha2-256",
          "ssh-rsa",
          "ssh-dss",
        ],
        cipher: SSH_ALGORITHMS.cipher,
        hmac: [
          "hmac-sha2-512-etm@openssh.com",
          "hmac-sha2-256-etm@openssh.com",
          "hmac-sha2-512",
          "hmac-sha2-256",
          "hmac-sha1",
          "hmac-md5",
        ],
        compress: ["none", "zlib@openssh.com", "zlib"],
      },
    };

    if (
      resolvedCredentials.authType === "password" &&
      resolvedCredentials.password
    ) {
      config.password = resolvedCredentials.password;
    } else if (
      resolvedCredentials.authType === "key" &&
      resolvedCredentials.sshKey
    ) {
      const cleanKey = resolvedCredentials.sshKey
        .trim()
        .replace(/\r\n/g, "\n")
        .replace(/\r/g, "\n");
      config.privateKey = Buffer.from(cleanKey, "utf8");
      if (resolvedCredentials.keyPassword) {
        config.passphrase = resolvedCredentials.keyPassword;
      }
    } else if (resolvedCredentials.authType === "agent") {
      const result = await applyAgentAuth(
        config,
        jumpHost.terminalConfig as unknown as
          Record<string, unknown> | undefined,
      );
      if ("error" in result) {
        throw new Error(result.error);
      }
    }

    client.on(
      "keyboard-interactive",
      (
        _name: string,
        _instructions: string,
        _lang: string,
        prompts: Array<{ prompt: string; echo: boolean }>,
        finish: (responses: string[]) => void,
      ) => {
        const responses = prompts.map((p) => {
          if (/password/i.test(p.prompt) && resolvedCredentials.password) {
            return resolvedCredentials.password as string;
          }
          return "";
        });
        finish(responses);
      },
    );

    if (currentClient) {
      await new Promise<void>((resolve, reject) => {
        currentClient!.forwardOut(
          "127.0.0.1",
          0,
          jumpHost.ip,
          jumpHost.port || 22,
          (err, stream) => {
            if (err) return reject(err);
            config.sock = stream;
            resolve();
          },
        );
      });
    }

    if (!config.sock) {
      await resolveSshConnectConfigHost(config);
    }

    await new Promise<void>((resolve, reject) => {
      client.on("ready", () => resolve());
      client.on("error", reject);
      client.connect(config);
    });

    currentClient = client;
  }

  return currentClient;
}

wss.on("connection", async (ws: WebSocket, req) => {
  const token = extractWebSocketToken(req);

  if (!token) {
    ws.close(1008, "Authentication required");
    return;
  }

  const authManagerInstance = AuthManager.getInstance();
  let payload;
  try {
    payload = await authManagerInstance.verifyJWTToken(token);
  } catch (error) {
    sshLogger.warn("Docker console JWT verification failed", {
      operation: "docker_console_auth_error",
      error: getErrorMessage(error),
    });
    ws.close(1008, "Authentication required");
    return;
  }
  if (!payload?.userId || payload.pendingTOTP) {
    ws.close(1008, "Authentication required");
    return;
  }

  const userId = payload.userId;
  const sessionId = `docker-console-${Date.now()}-${Math.random()}`;
  sshLogger.info("Docker console WebSocket connected", {
    operation: "docker_console_connect",
    sessionId,
    userId,
  });

  let sshSession: SSHSession | null = null;

  const wsPingInterval = setInterval(() => {
    if (ws.readyState === WebSocket.OPEN) {
      ws.ping();
    }
  }, 30000);

  const cleanup = () => {
    clearInterval(wsPingInterval);
    if (!sshSession) return;
    sshSession.stream?.end();
    sshSession.client.end();
    activeSessions.delete(sessionId);
    sshSession = null;
  };

  ws.on("message", async (data) => {
    try {
      const message = parseWsMessage(data);

      switch (message.type) {
        case "connect": {
          const connectData = asObject(message.data);
          const hostConfig = asObject(
            connectData.hostConfig,
          ) as unknown as SSHHost;
          const containerId = asString(connectData.containerId);
          const shell = asString(connectData.shell) || undefined;
          const cols = toTerminalDimension(connectData.cols) || 80;
          const rows = toTerminalDimension(connectData.rows) || 24;

          const hostId = hostConfig?.id;

          if (!hostId || !containerId) {
            ws.send(
              JSON.stringify({
                type: "error",
                message: "Host configuration and container ID are required",
              }),
            );
            return;
          }

          if (!/^[a-zA-Z0-9][a-zA-Z0-9_.-]*$/.test(containerId)) {
            ws.send(
              JSON.stringify({
                type: "error",
                message: "Invalid container ID",
              }),
            );
            return;
          }

          const allowedShells = ["bash", "sh", "ash", "zsh"];
          if (shell && !allowedShells.includes(shell)) {
            ws.send(
              JSON.stringify({
                type: "error",
                message: "Invalid shell",
              }),
            );
            return;
          }

          if (!hostConfig.enableDocker) {
            ws.send(
              JSON.stringify({
                type: "error",
                message: "Docker is not enabled on this host",
              }),
            );
            return;
          }

          try {
            // Resolve host with credentials server-side
            const { resolveHostById, resolveHostBySyncId } =
              await import("../../../src/backend/hosts/host-resolver.js");
            // syncId names the host on both sides of a sync pair; the numeric
            // id only names it in the database the client is displaying.
            const hostSyncId = hostConfig?.syncId;
            const resolvedHost = hostSyncId
              ? await resolveHostBySyncId(hostSyncId, userId)
              : await resolveHostById(hostId, userId);

            if (!resolvedHost) {
              ws.send(
                JSON.stringify({
                  type: "error",
                  message: hostSyncId
                    ? HOST_NOT_ON_THIS_SERVER_MESSAGE
                    : "Host not found",
                }),
              );
              return;
            }

            // The connection below dials resolvedHost.ip outright, so if this
            // server has a different machine under the id the client sent, the
            // console opens on that machine's Docker daemon instead.
            if (
              !hostSyncId &&
              hostAddressMismatch(hostConfig?.ip, resolvedHost.ip)
            ) {
              sshLogger.error(
                "Refusing Docker console: host id resolves to a different address here",
                undefined,
                {
                  operation: "docker_console_host_id_mismatch",
                  hostId,
                  userId,
                  clientIp: hostConfig?.ip,
                  resolvedIp: resolvedHost.ip,
                },
              );
              ws.send(
                JSON.stringify({
                  type: "error",
                  message: HOST_ADDRESS_MISMATCH_MESSAGE,
                }),
              );
              return;
            }

            if (!resolvedHost.enableDocker) {
              ws.send(
                JSON.stringify({
                  type: "error",
                  message:
                    "Docker is not enabled for this host. Enable it in Host Settings.",
                }),
              );
              return;
            }

            const client = new SSHClient();

            const config: Record<string, unknown> = {
              host: resolvedHost.ip?.replace(/^\[|\]$/g, "") || resolvedHost.ip,
              port: resolvedHost.port || 22,
              username: resolvedHost.username,
              tryKeyboard: true,
              readyTimeout: 60000,
              keepaliveInterval: 30000,
              keepaliveCountMax: 120,
              tcpKeepAlive: true,
              tcpKeepAliveInitialDelay: 30000,
            };

            if (resolvedHost.authType === "password" && resolvedHost.password) {
              config.password = resolvedHost.password;
            } else if (resolvedHost.authType === "key" && resolvedHost.key) {
              const cleanKey = resolvedHost.key
                .trim()
                .replace(/\r\n/g, "\n")
                .replace(/\r/g, "\n");
              config.privateKey = Buffer.from(cleanKey, "utf8");
              if (resolvedHost.keyPassword) {
                config.passphrase = resolvedHost.keyPassword;
              }
            } else if (resolvedHost.authType === "agent") {
              const result = await applyAgentAuth(
                config,
                resolvedHost.terminalConfig as unknown as
                  Record<string, unknown> | undefined,
              );
              if ("error" in result) {
                ws.send(
                  JSON.stringify({
                    type: "error",
                    message: result.error,
                  }),
                );
                return;
              }
            }

            if (resolvedHost.jumpHosts && resolvedHost.jumpHosts.length > 0) {
              const jumpClient = await createJumpHostChain(
                resolvedHost.jumpHosts,
                userId,
              );
              if (jumpClient) {
                const stream = await new Promise<import("ssh2").ClientChannel>(
                  (resolve, reject) => {
                    jumpClient.forwardOut(
                      "127.0.0.1",
                      0,
                      resolvedHost.ip,
                      resolvedHost.port || 22,
                      (err, stream) => {
                        if (err) return reject(err);
                        resolve(stream);
                      },
                    );
                  },
                );
                config.sock = stream;
              }
            }

            if (!config.sock) {
              await resolveSshConnectConfigHost(config);
            }

            await new Promise<void>((resolve, reject) => {
              client.on("ready", () => resolve());
              client.on("error", reject);
              client.connect(config);
            });

            const { runtime: containerRuntime } = getContainerRuntimeConfig(
              resolvedHost.dockerConfig,
            );

            sshSession = {
              client,
              stream: null,
              isConnected: true,
              containerId,
              hostId: resolvedHost.id,
              containerRuntime,
            };

            activeSessions.set(sessionId, sshSession);

            let shellToUse = shell || "bash";

            if (shell) {
              try {
                await new Promise<void>((resolve, reject) => {
                  client.exec(
                    containerCommand(
                      containerRuntime,
                      `exec ${containerId} which ${shell}`,
                    ),
                    (err, stream) => {
                      if (err) return reject(err);

                      let output = "";
                      stream.on("data", (data: Buffer) => {
                        output += data.toString();
                      });

                      stream.on("close", (code: number) => {
                        if (code === 0 && output.trim()) {
                          resolve();
                        } else {
                          reject(new Error(`Shell ${shell} not available`));
                        }
                      });

                      stream.stderr.on("data", () => {});
                      stream.stderr.on("error", () => {});
                      stream.on("error", (streamErr) => {
                        reject(streamErr);
                      });
                    },
                  );
                });
              } catch {
                sshLogger.warn(
                  `Requested shell ${shell} not found, detecting available shell`,
                  {
                    operation: "shell_validation",
                    sessionId,
                    containerId,
                    requestedShell: shell,
                  },
                );
                shellToUse = await detectShell(sshSession, containerId);
              }
            } else {
              shellToUse = await detectShell(sshSession, containerId);
            }

            sshSession.shell = shellToUse;

            const execCommand = containerCommand(
              containerRuntime,
              `exec -it ${containerId} /bin/${shellToUse}`,
            );
            sshLogger.info("Attaching to Docker container", {
              operation: "docker_attach",
              sessionId,
              userId,
              hostId: resolvedHost.id,
              containerId,
            });

            client.exec(
              execCommand,
              {
                pty: {
                  term: "xterm-256color",
                  cols,
                  rows,
                },
              },
              (err, stream) => {
                if (err) {
                  sshLogger.error("Failed to create docker exec", err, {
                    operation: "docker_exec",
                    sessionId,
                    containerId,
                  });

                  ws.send(
                    JSON.stringify({
                      type: "error",
                      message: `Failed to start console: ${err.message}`,
                    }),
                  );
                  return;
                }

                sshSession!.stream = stream;
                sshLogger.success("Docker container attached", {
                  operation: "docker_attach_success",
                  sessionId,
                  userId,
                  hostId: resolvedHost.id,
                  containerId,
                });

                // Buffers incomplete multi-byte UTF-8 sequences across chunk
                // boundaries so box-drawing/special characters don't get
                // corrupted when a character is split across TCP packets.
                const decoder = new StringDecoder("utf-8");

                stream.on("data", (data: Buffer) => {
                  if (ws.readyState === WebSocket.OPEN) {
                    const text = decoder.write(data);
                    if (!text) return;
                    ws.send(
                      JSON.stringify({
                        type: "output",
                        data: text,
                      }),
                    );
                  }
                });

                stream.stderr.on("data", () => {});
                stream.stderr.on("error", () => {});

                stream.on("error", (streamErr) => {
                  sshLogger.error("Docker console stream error", streamErr, {
                    operation: "docker_console_stream_error",
                    sessionId,
                    containerId,
                  });
                  if (ws.readyState === WebSocket.OPEN) {
                    ws.send(
                      JSON.stringify({
                        type: "error",
                        message: `Console error: ${streamErr.message}`,
                      }),
                    );
                  }
                });

                stream.on("close", () => {
                  if (ws.readyState === WebSocket.OPEN) {
                    ws.send(
                      JSON.stringify({
                        type: "disconnected",
                        message: "Console session ended",
                      }),
                    );
                  }

                  if (sshSession) {
                    sshSession.client.end();
                    activeSessions.delete(sessionId);
                  }
                });

                ws.send(
                  JSON.stringify({
                    type: "connected",
                    data: {
                      shell: shellToUse,
                      requestedShell: shell,
                      shellChanged: shell && shell !== shellToUse,
                    },
                  }),
                );
              },
            );
          } catch (error) {
            sshLogger.error("Failed to connect to container", error, {
              operation: "console_connect",
              sessionId,
              containerId,
            });

            ws.send(
              JSON.stringify({
                type: "error",
                message: getErrorMessage(
                  error,
                  "Failed to connect to container",
                ),
              }),
            );
          }
          break;
        }

        case "input": {
          if (sshSession && sshSession.stream) {
            const input = asString(message.data);
            if (input) sshSession.stream.write(input);
          }
          break;
        }

        case "resize": {
          if (sshSession && sshSession.stream) {
            const dimensions = asObject(message.data);
            const cols = toTerminalDimension(dimensions.cols);
            const rows = toTerminalDimension(dimensions.rows);
            if (cols && rows)
              sshSession.stream.setWindow(rows, cols, rows, cols);
          }
          break;
        }

        case "disconnect": {
          if (sshSession) {
            if (sshSession.stream) {
              sshSession.stream.end();
            }
            sshSession.client.end();
            activeSessions.delete(sessionId);

            ws.send(
              JSON.stringify({
                type: "disconnected",
                message: "Disconnected from container",
              }),
            );
          }
          break;
        }

        case "ping": {
          if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ type: "pong" }));
          }
          break;
        }

        default:
          sshLogger.warn("Unknown message type", {
            operation: "ws_message",
            type: message.type,
          });
      }
    } catch (error) {
      sshLogger.error("WebSocket message error", error, {
        operation: "ws_message",
        sessionId,
      });

      ws.send(
        JSON.stringify({
          type: "error",
          message: getErrorMessage(error, "An error occurred"),
        }),
      );
    }
  });

  ws.on("close", () => {
    sshLogger.info("Docker console disconnected", {
      operation: "docker_console_disconnect",
      sessionId,
      userId,
      hostId: sshSession?.hostId,
      containerId: sshSession?.containerId,
    });
    cleanup();
  });

  ws.on("error", (error) => {
    sshLogger.error("WebSocket error", error, {
      operation: "ws_error",
      sessionId,
    });

    cleanup();
  });
});

/**
 * Closes every live console session and the WebSocket server. Called by the
 * plugin's deactivate() as well as SIGTERM, since disabling the docker
 * plugin has to free port 30009 the same way process shutdown does.
 */
export function closeConsoleServer(): Promise<void> {
  activeSessions.forEach((session) => {
    if (session.stream) {
      session.stream.end();
    }
    session.client.end();
  });

  activeSessions.clear();

  return new Promise((resolve) => wss.close(() => resolve()));
}

process.on("SIGTERM", () => {
  void closeConsoleServer().then(() => process.exit(0));
});