import { FileHelper } from '@start9labs/start-sdk'
import { z } from 'zod'
import { sdk } from '../sdk'

export const sparkConfigSchema = z.object({
  spark1_host: z.string().catch(''),
  spark1_user: z.string().catch(''),
  spark2_host: z.string().catch(''),
  spark2_user: z.string().catch(''),
  // Optional vLLM port override (Spark 1). Blank => 8888 (launch-cluster.sh default).
  vllm_port: z.string().catch(''),
  // Optional per-service overrides. Blank => use spark2_host / spark2_user.
  parakeet_host: z.string().catch(''),
  parakeet_user: z.string().catch(''),
  parakeet_container: z.string().catch(''),
  kokoro_host: z.string().catch(''),
  kokoro_user: z.string().catch(''),
  kokoro_container: z.string().catch(''),
  // Optional overrides for the embedding server (spark-embed) + Qdrant.
  embed_host: z.string().catch(''),
  embed_user: z.string().catch(''),
  embed_container: z.string().catch(''),
  qdrant_host: z.string().catch(''),
  qdrant_user: z.string().catch(''),
  qdrant_container: z.string().catch(''),
  qdrant_collection: z.string().catch(''),
  // Optional matrix-bridge bot. Blank => no tile. Host reuses Spark 2.
  matrix_bridge_user: z.string().catch(''),
  // Optional Open WebUI deep-link
  open_webui_url: z.string().catch(''),
  // Optional NGC API key for pulling NIM containers from nvcr.io/nim/...
  ngc_api_key: z.string().catch(''),
})

export type SparkConfig = z.infer<typeof sparkConfigSchema>

export const sparkConfigYaml = FileHelper.yaml(
  { base: sdk.volumes.main, subpath: 'config.yaml' },
  sparkConfigSchema,
)