import fs from 'node:fs'
import path from 'node:path'
import type { FimOrder } from '../../shared/prompt.ts'

export type ContextMode = 'symbols' | 'simple' | 'raw'

/** What a model was trained on; the toggle only offers these modes. */
export interface ModelProfile {
  path: string
  order: FimOrder
  modes: ContextMode[]
  defaultMode: ContextMode
  /** Symbol-context format the model was trained on (shared/symbols.ts): 1 = Stage C, 2 = + namespace imports. */
  contextFormat?: 1 | 2
  /** Trained with the TS-errors header section (Stage D+). */
  errors?: boolean
  /** Override symbols.members for this model (member lists after `expr.`). */
  members?: boolean
  /** Trained for next-edit prediction (shared/nextedit.ts), enables /v1/next-edit. */
  nextEdit?: boolean
  /** Another profile (trained with nextEdit) that serves /v1/next-edit for this model, in a second llama-server. */
  nextEditModel?: string
}

export interface DaemonConfig {
  host: string
  port: number
  tokenizer: { json: string; config: string }
  llamaServer: {
    binary: string
    port: number
    ctx: number
    /** Extra flags, e.g. ["-ngl", "99"] on Apple Silicon, ["-t", "16"] on the laptop CPU. */
    args: string[]
    /** Extra flags for the next-edit companion server only. Default: n-gram speculative decoding (the output is
     *  mostly a copy of the region, so drafts from the prompt are accepted; exact with greedy sampling). */
    nextEditArgs?: string[]
  }
  models: Record<string, ModelProfile>
  defaultModel: string
  prompt: {
    prefixTokens: number
    suffixTokens: number
    /** Prefix window start moves in multiples of this many lines (KV-cache stability). */
    prefixChunkLines: number
    maxLines: number
    nPredictSingleLine: number
    nPredictMultiLine: number
  }
  symbols: {
    /** Token budget for the __symbols__.d.ts context (Stage C was trained with up to 2048). */
    budget: number
    /** How long a request waits for fresh context when the anchor changes before using stale/none. */
    waitMs: number
    /** Member completions after `expr.`: in training, but every `.` changes the header (cold prefill). */
    members: boolean
  }
  thisRepair: { enabled: boolean; waitMs: number }
  /** Cut a half-typed word off the prompt and force the output to start with it (default on). */
  tokenHealing?: boolean
  nextEditOptions?: { regionBefore: number; regionAfter: number; maxTokens: number; beforeTokens: number; afterTokens: number }
  telemetry: { enabled: boolean; dir: string }
  /** Where opted-in edit recordings from editors are appended (the opt-in is the editor setting). */
  recording: { dir: string }
  stateFile: string
  /** Workspace roots for clients that can't send one (Zed via /v1/completions): files are looked up under these. */
  workspaceRoots?: string[]
  /** Exit after this many minutes without any request (0 = never). Editors poll every 15 s while open. */
  idleShutdownMinutes?: number
  /** Opaque tag reported by /health (the VS Code extension uses it to replace daemons left over from older versions). */
  instanceTag?: string
}

/** Toggles that survive daemon restarts and are shared by every editor talking to this daemon. */
export interface RuntimeState {
  enabled: boolean
  model: string
  mode: ContextMode
  /** config.defaultModel when this state was saved: a changed default (editor model setting, new release) wins. */
  defaultModel?: string
}

export function loadConfig(configPath: string): DaemonConfig {
  const dir = path.dirname(path.resolve(configPath))
  const raw = JSON.parse(fs.readFileSync(configPath, 'utf8')) as DaemonConfig
  const abs = (p: string) => (path.isAbsolute(p) || !p.includes('/') ? p : path.resolve(dir, p)) // bare names (e.g. "llama-server") stay on PATH
  raw.tokenizer = { json: abs(raw.tokenizer.json), config: abs(raw.tokenizer.config) }
  raw.llamaServer.binary = abs(raw.llamaServer.binary)
  for (const m of Object.values(raw.models)) m.path = abs(m.path)
  raw.telemetry.dir = abs(raw.telemetry.dir)
  raw.recording.dir = abs(raw.recording.dir)
  raw.stateFile = abs(raw.stateFile)
  if (!raw.models[raw.defaultModel]) throw new Error(`defaultModel ${raw.defaultModel} is not in models`)
  for (const [name, m] of Object.entries(raw.models)) {
    if (!m.modes.includes(m.defaultMode)) throw new Error(`model ${name}: defaultMode must be one of its modes`)
    if (m.nextEditModel && !raw.models[m.nextEditModel]?.nextEdit) throw new Error(`model ${name}: nextEditModel ${m.nextEditModel} must be a profile with nextEdit`)
  }
  return raw
}

export function loadState(config: DaemonConfig): RuntimeState {
  const fallback: RuntimeState = { enabled: true, model: config.defaultModel, mode: config.models[config.defaultModel].defaultMode, defaultModel: config.defaultModel }
  try {
    const s = JSON.parse(fs.readFileSync(config.stateFile, 'utf8')) as Partial<RuntimeState>
    const keepModel = s.defaultModel === config.defaultModel && s.model && config.models[s.model]
    const model = keepModel ? s.model! : fallback.model
    const mode = s.mode && config.models[model].modes.includes(s.mode) ? s.mode : config.models[model].defaultMode
    return { enabled: s.enabled ?? true, model, mode, defaultModel: config.defaultModel }
  } catch {
    return fallback
  }
}

export function saveState(config: DaemonConfig, state: RuntimeState): void {
  fs.mkdirSync(path.dirname(config.stateFile), { recursive: true })
  fs.writeFileSync(config.stateFile, JSON.stringify(state, null, 2) + '\n')
}