import { sdk } from '../sdk'
import { configFile } from '../file-models/config'
import { sshKeyFile, hfTokenFile } from '../file-models/secrets'

const { InputSpec, Value } = sdk

const inputSpec = InputSpec.of({
  primarySparkHost: Value.text({
    name: 'Primary Spark Host',
    description: 'Hostname or IP of the head DGX Spark (reachable over SSH). Serves models, hosts the model proxy, and runs the grading panel.',
    required: true,
    default: null,
    placeholder: 'spark-01.local',
  }),
  primarySparkUser: Value.text({
    name: 'SSH User',
    description: 'The login user on the Spark (DGX OS default is "nvidia").',
    required: true,
    default: 'nvidia',
  }),
  sshPort: Value.number({
    name: 'SSH Port',
    description: 'SSH port on the Spark.',
    required: true,
    default: 22,
    integer: true,
    min: 1,
    max: 65535,
  }),
  sshPrivateKey: Value.textarea({
    name: 'SSH Private Key',
    description:
      'A private key (PEM/OpenSSH) whose public half is in the Spark user\'s ' +
      '~/.ssh/authorized_keys. Stored in this service\'s private volume and ' +
      'used only to reach your Sparks. Paste the FULL key including header/footer.',
    warning:
      'This is a credential. It is written to the service volume and never ' +
      'shown again. Use a dedicated key for this service.',
    required: true,
    default: null,
    minRows: 6,
    maxRows: 14,
    placeholder: '-----BEGIN OPENSSH PRIVATE KEY-----\n...\n-----END OPENSSH PRIVATE KEY-----',
  }),
  useBothSparks: Value.toggle({
    name: 'Use Both Sparks',
    description:
      'Allow models to be served on a second Spark (over ConnectX/200GbE) for ' +
      'extra capacity. NOTE: in air-gapped grading mode all models must run on ' +
      'the head Spark; the second Spark is used only in local-services mode.',
    default: false,
  }),
  secondarySparkHost: Value.text({
    name: 'Secondary Spark Host',
    description: 'Hostname/IP of the second Spark. Required only if "Use Both Sparks" is on.',
    required: false,
    default: null,
    placeholder: 'spark-02.local',
  }),
  headInternalHost: Value.text({
    name: 'Head Internal Host',
    description:
      'Address the head Spark uses for its own preflight checks against the local ' +
      'model proxy. Single Spark: 127.0.0.1 is fine.',
    required: true,
    default: '127.0.0.1',
  }),
  remoteWorkDir: Value.text({
    name: 'Remote Work Directory',
    description: 'Absolute path on the head Spark for staged deck text, the HF cache, and logs.',
    required: true,
    default: '/home/nvidia/boardroom-map',
  }),
  servingImage: Value.text({
    name: 'vLLM Image Tag',
    description: 'The vLLM serving image built on the Sparks (e.g. via spark-vllm-docker).',
    required: true,
    default: 'boardroom-vllm:latest',
  }),
  graderImage: Value.text({
    name: 'Grader Image Tag',
    description: 'The sandboxed grader image built on the head Spark from sandbox/build.sh.',
    required: true,
    default: 'boardroom-grader:latest',
  }),
  hfToken: Value.text({
    name: 'Hugging Face Token (optional)',
    description:
      'Only needed for gated/private models or to warm a model the first time. ' +
      'Leave empty to keep the existing token unchanged. Passed to the vLLM ' +
      'container at serve time.',
    required: false,
    default: null,
    masked: true,
  }),
})

export const configureSparks = sdk.Action.withInput(
  'configure-sparks',

  async ({ effects }) => ({
    name: 'Configure Sparks',
    description: 'Set the DGX Spark connection details, SSH credentials, and image tags.',
    warning: null,
    allowedStatuses: 'any',
    group: 'Setup',
    visibility: 'enabled',
  }),

  inputSpec,

  // Prefill non-secret fields from existing config. Never prefill the key/token.
  async ({ effects }) => {
    const cfg = await configFile.read().const(effects)
    if (!cfg) return {}
    return {
      primarySparkHost: cfg.primarySparkHost || undefined,
      primarySparkUser: cfg.primarySparkUser,
      sshPort: cfg.sshPort,
      useBothSparks: cfg.useBothSparks,
      secondarySparkHost: cfg.secondarySparkHost ?? undefined,
      headInternalHost: cfg.headInternalHost,
      remoteWorkDir: cfg.remoteWorkDir,
      servingImage: cfg.servingImage,
      graderImage: cfg.graderImage,
    }
  },

  async ({ effects, input }) => {
    // Persist the private key to its own file (600 enforced in-container).
    await sshKeyFile.write(effects, input.sshPrivateKey.trim() + '\n')

    let hfTokenSet = (await configFile.read().const(effects))?.hfTokenSet ?? false
    if (input.hfToken && input.hfToken.trim()) {
      await hfTokenFile.write(effects, input.hfToken.trim())
      hfTokenSet = true
    }

    await configFile.merge(effects, {
      primarySparkHost: input.primarySparkHost,
      primarySparkUser: input.primarySparkUser,
      sshPort: input.sshPort,
      useBothSparks: input.useBothSparks,
      secondarySparkHost: input.secondarySparkHost,
      headInternalHost: input.headInternalHost,
      remoteWorkDir: input.remoteWorkDir,
      servingImage: input.servingImage,
      graderImage: input.graderImage,
      hfTokenSet,
    })

    return {
      version: '1',
      title: 'Sparks Configured',
      message:
        'Saved. Use "Test Spark Connection" to verify SSH + GPU access, then set ' +
        '"Configure Models" and "Configure Graders".',
      result: { type: 'single', value: input.primarySparkHost, copyable: false, qr: false, masked: false },
    }
  },
)