import { stat } from 'node:fs/promises'
import { join } from 'node:path'
import type { ModelLoadInfo, ModelProgress } from '../decisions/adapters/crivo.js'
import type { RuntimeMetrics } from '../decisions/adapters/runtime.js'

/** Human-readable diagnostics only; stdout remains the JSON response. */
export class TerminalProgress {
  private readonly started = performance.now()
  private info?: ModelLoadInfo
  private missing = new Set<string>()
  private steps = new Map<string, number>()

  constructor(
    private readonly write: (message: string) => void = (message) => {
      process.stderr.write(`${message}\n`)
    }
  ) {}

  async loading(info: ModelLoadInfo): Promise<void> {
    this.info = info
    const modelFile = info.dtype === 'q8' ? 'onnx/model_quantized.onnx' : 'onnx/model.onnx'
    const weights = info.weightFiles ?? [modelFile]
    const files = ['config.json', 'tokenizer.json', 'tokenizer_config.json', ...weights]
    const cached = await Promise.all(
      files.map(async (file) => {
        if (!info.cacheDir) return false
        try {
          const entry = await stat(join(info.cacheDir, info.modelId, file))
          return entry.isFile() && entry.size > 0
        } catch {
          return false
        }
      })
    )
    this.missing = new Set(files.filter((_, index) => !cached[index]))
    this.write(
      `Model: ${info.modelId}${info.variant ? ` / ${info.variant}` : ''} | precision: ${info.dtype} | providers: ${info.providers.join(', ')}`
    )
    if (this.missing.size) {
      const weightMiB =
        info.weightMiB ??
        (info.modelId === 'dizyann/crivo-150m-web' ? (info.dtype === 'q8' ? 144 : 571) : undefined)
      const size =
        weights.some((file) => this.missing.has(file)) && weightMiB !== undefined
          ? ` Weights: approximately ${weightMiB} MiB.`
          : ''
      this.write(`Files missing from cache; download required.${size}`)
      this.write(`Download: ${[...this.missing].join(', ')}`)
    } else {
      this.write('Files found in cache; loading the model session.')
    }
    this.write('The first run may take a while to download files and initialize the backend.')
  }

  progress(event: ModelProgress): void {
    const file = event.file
    if (!file || !this.missing.has(file)) return
    if (
      event.status === 'progress' &&
      event.progress !== undefined &&
      Number.isFinite(event.progress)
    ) {
      const step = Math.min(100, Math.floor(event.progress / 10) * 10)
      if (step <= (this.steps.get(file) ?? -1)) return
      this.steps.set(file, step)
      this.write(`Downloading ${file}: ${step}%`)
    } else if (event.status === 'done') {
      this.missing.delete(file)
      this.write(`File ready: ${file}`)
      if (!this.missing.size) this.write('Files available; initializing the model session…')
    }
  }

  summary(metrics: RuntimeMetrics, count: number): void {
    const seconds = (ms: number) => `${(ms / 1000).toFixed(2)} s`
    this.write(
      `Completed: ${count} decisions | total: ${seconds(performance.now() - this.started)}`
    )
    this.write(
      `Loading/download: ${seconds(metrics.loadMs)} | inference: ${seconds(metrics.inferenceMs)} | cleanup: ${seconds(metrics.disposeMs)}`
    )
    if (this.info)
      this.write(
        `Model: ${this.info.modelId}${this.info.variant ? ` / ${this.info.variant}` : ''} | precision: ${this.info.dtype} | providers: ${this.info.providers.join(', ')}`
      )
  }
}