// Parser for pulseq .seq files, formats v1.4.x and v1.5.x, mirroring
// pulseq's +mr/@Sequence/read.m (column layouts, unit scales, shape RLE).
// Older formats (< 1.4.0) are rejected with a clear message.
import { md5Hex } from './md5.ts'
import type {
  AdcEvent,
  ArbGradEvent,
  DefinitionValue,
  ExtensionRef,
  ExtensionSpec,
  GradEvent,
  ParsedSeq,
  RfEvent,
  SeqBlock,
  SeqVersion,
  Shape,
  TrapGradEvent,
} from './types.ts'

export class SeqParseError extends Error {}

const DEFAULT_RASTERS = {
  blockDuration: 1e-5,
  gradient: 1e-5,
  rf: 1e-6,
  adc: 1e-7,
}

function tokens(line: string): string[] {
  return line.trim().split(/\s+/)
}

function nums(line: string, context: string): number[] {
  const out = tokens(line).map(Number)
  if (out.some((v) => Number.isNaN(v))) {
    throw new SeqParseError(`malformed numeric line in ${context}: "${line.trim()}"`)
  }
  return out
}

/** pulseq shape decompression: RLE on the derivative (v v n => v x(n+2)), then cumsum. */
export function decompressShape(numSamples: number, data: number[]): Float64Array {
  if (data.length === numSamples) {
    // uncompressed (stored) samples — v1.4+ marks this by matching lengths
    return Float64Array.from(data)
  }
  const deriv: number[] = []
  let i = 0
  while (i < data.length) {
    if (i + 2 < data.length && data[i] === data[i + 1]) {
      const v = data[i]
      const rep = data[i + 2] + 2
      for (let k = 0; k < rep; k++) deriv.push(v)
      i += 3
    } else {
      deriv.push(data[i])
      i++
    }
  }
  if (deriv.length !== numSamples) {
    throw new SeqParseError(
      `shape decompression produced ${deriv.length} samples, expected ${numSamples}`,
    )
  }
  const out = new Float64Array(numSamples)
  let acc = 0
  for (let k = 0; k < numSamples; k++) {
    acc += deriv[k]
    out[k] = acc
  }
  return out
}

export function parseSeq(text: string): ParsedSeq {
  const lines = text.split(/\r?\n/)

  let version: SeqVersion | null = null
  const definitions = new Map<string, DefinitionValue>()
  const blocks: SeqBlock[] = []
  const rf = new Map<number, RfEvent>()
  const grads = new Map<number, GradEvent>()
  const adcs = new Map<number, AdcEvent>()
  const shapes = new Map<number, Shape>()
  const extensions = new Map<number, ExtensionRef>()
  const extensionSpecs: ExtensionSpec[] = []
  let signature: ParsedSeq['signature']

  // Raw [BLOCKS] rows; converted to seconds once rasters are known (the
  // [DEFINITIONS] section precedes [BLOCKS] in practice, but don't rely on it).
  const blockRows: number[][] = []

  type Section =
    | 'none'
    | 'version'
    | 'definitions'
    | 'blocks'
    | 'rf'
    | 'gradients'
    | 'trap'
    | 'adc'
    | 'shapes'
    | 'extensions'
    | 'extension-spec'
    | 'signature'
    | 'skip'
  let section: Section = 'none'
  let currentSpec: ExtensionSpec | null = null

  // [SHAPES] accumulation state
  let shapeId = -1
  let shapeNumSamples = -1
  let shapeData: number[] = []
  const flushShape = () => {
    if (shapeId < 0) return
    shapes.set(shapeId, {
      id: shapeId,
      numSamples: shapeNumSamples,
      samples: decompressShape(shapeNumSamples, shapeData),
      storedSamples: shapeData.length,
    })
    shapeId = -1
    shapeNumSamples = -1
    shapeData = []
  }

  const requireVersion = (): SeqVersion => {
    if (!version) throw new SeqParseError('file must declare [VERSION] before event sections')
    return version
  }

  for (const rawLine of lines) {
    const line = rawLine.trim()
    if (line === '' || line.startsWith('#')) continue

    if (line.startsWith('[') && line.endsWith(']')) {
      if (section === 'shapes') flushShape()
      currentSpec = null
      switch (line) {
        case '[VERSION]':
          section = 'version'
          version = { major: 0, minor: 0, revision: '0', combined: 0 }
          break
        case '[DEFINITIONS]':
          section = 'definitions'
          break
        case '[BLOCKS]':
          section = 'blocks'
          break
        case '[RF]':
          requireVersion()
          section = 'rf'
          break
        case '[GRADIENTS]':
          requireVersion()
          section = 'gradients'
          break
        case '[TRAP]':
          section = 'trap'
          break
        case '[ADC]':
          requireVersion()
          section = 'adc'
          break
        case '[SHAPES]':
          section = 'shapes'
          break
        case '[EXTENSIONS]':
          section = 'extensions'
          break
        case '[SIGNATURE]':
          section = 'signature'
          break
        case '[DELAYS]':
          throw new SeqParseError(
            'this file uses the pre-1.4.0 pulseq format ([DELAYS] section), which is not supported',
          )
        default:
          section = 'skip' // unknown section: ignore its lines
          break
      }
      continue
    }

    // `extension NAME typeId` opens a specification subsection
    if ((section === 'extensions' || section === 'extension-spec') && line.startsWith('extension ')) {
      const t = tokens(line)
      currentSpec = { name: t[1], typeId: Number(t[2]), rows: [] }
      extensionSpecs.push(currentSpec)
      section = 'extension-spec'
      continue
    }

    switch (section) {
      case 'version': {
        const [key, value] = tokens(line)
        const v = version!
        if (key === 'major') v.major = Number(value)
        else if (key === 'minor') v.minor = Number(value)
        else if (key === 'revision') v.revision = value
        v.combined = 1000000 * v.major + 1000 * v.minor + (parseInt(v.revision, 10) || 0)
        break
      }
      case 'definitions': {
        const t = tokens(line)
        const key = t[0]
        const rest = t.slice(1)
        const asNums = rest.map(Number)
        definitions.set(
          key,
          rest.length > 0 && asNums.every((v) => !Number.isNaN(v))
            ? asNums
            : line.trim().slice(key.length).trim(),
        )
        break
      }
      case 'blocks':
        blockRows.push(nums(line, '[BLOCKS]'))
        break
      case 'rf': {
        const v = requireVersion()
        if (v.combined >= 1005000) {
          // id amp mag_id phase_id time_id center(us) delay(us) freqPPM phasePPM freq phase use
          const t = tokens(line)
          const d = t.slice(0, 11).map(Number)
          if (d.some(Number.isNaN) || t.length < 12) {
            throw new SeqParseError(`malformed [RF] line: "${line}"`)
          }
          rf.set(d[0], {
            id: d[0],
            amp: d[1],
            magShapeId: d[2],
            phaseShapeId: d[3],
            timeShapeId: d[4],
            center: d[5] * 1e-6,
            delay: d[6] * 1e-6,
            freqPPM: d[7],
            phasePPM: d[8],
            freq: d[9],
            phase: d[10],
            use: t[11],
          })
        } else {
          // v1.4: id amp mag_id phase_id time_id delay(us) freq phase
          const d = nums(line, '[RF]')
          rf.set(d[0], {
            id: d[0],
            amp: d[1],
            magShapeId: d[2],
            phaseShapeId: d[3],
            timeShapeId: d[4],
            center: NaN,
            delay: d[5] * 1e-6,
            freqPPM: 0,
            phasePPM: 0,
            freq: d[6],
            phase: d[7],
            use: 'u',
          })
        }
        break
      }
      case 'gradients': {
        const v = requireVersion()
        const d = nums(line, '[GRADIENTS]')
        let ev: ArbGradEvent
        if (v.combined >= 1005000) {
          // id amp first last shape_id time_id delay(us)
          ev = {
            id: d[0],
            kind: 'grad',
            amp: d[1],
            first: d[2],
            last: d[3],
            shapeId: d[4],
            timeShapeId: d[5],
            delay: d[6] * 1e-6,
          }
        } else {
          // v1.4: id amp shape_id time_id delay(us); first/last derived on demand
          ev = {
            id: d[0],
            kind: 'grad',
            amp: d[1],
            first: NaN,
            last: NaN,
            shapeId: d[2],
            timeShapeId: d[3],
            delay: d[4] * 1e-6,
          }
        }
        grads.set(ev.id, ev)
        break
      }
      case 'trap': {
        // id amp rise(us) flat(us) fall(us) delay(us)
        const d = nums(line, '[TRAP]')
        const ev: TrapGradEvent = {
          id: d[0],
          kind: 'trap',
          amp: d[1],
          rise: d[2] * 1e-6,
          flat: d[3] * 1e-6,
          fall: d[4] * 1e-6,
          delay: d[5] * 1e-6,
        }
        grads.set(ev.id, ev)
        break
      }
      case 'adc': {
        const v = requireVersion()
        const d = nums(line, '[ADC]')
        if (v.combined >= 1005000) {
          // id num dwell(ns) delay(us) freqPPM phasePPM freq phase phase_id
          adcs.set(d[0], {
            id: d[0],
            num: d[1],
            dwell: d[2] * 1e-9,
            delay: d[3] * 1e-6,
            freqPPM: d[4],
            phasePPM: d[5],
            freq: d[6],
            phase: d[7],
            phaseShapeId: d[8] ?? 0,
          })
        } else {
          // v1.4: id num dwell(ns) delay(us) freq phase
          adcs.set(d[0], {
            id: d[0],
            num: d[1],
            dwell: d[2] * 1e-9,
            delay: d[3] * 1e-6,
            freqPPM: 0,
            phasePPM: 0,
            freq: d[4],
            phase: d[5],
            phaseShapeId: 0,
          })
        }
        break
      }
      case 'shapes': {
        const t = tokens(line)
        if (t[0] === 'shape_id') {
          flushShape()
          shapeId = Number(t[1])
        } else if (t[0] === 'num_samples') {
          shapeNumSamples = Number(t[1])
        } else {
          for (const tok of t) {
            const v = Number(tok)
            if (Number.isNaN(v)) throw new SeqParseError(`malformed shape sample: "${line}"`)
            shapeData.push(v)
          }
        }
        break
      }
      case 'extensions': {
        const d = nums(line, '[EXTENSIONS]')
        extensions.set(d[0], { id: d[0], type: d[1], ref: d[2], next: d[3] })
        break
      }
      case 'extension-spec':
        currentSpec!.rows.push(tokens(line))
        break
      case 'signature': {
        const [key, value] = tokens(line)
        signature = signature ?? { type: '', hash: '' }
        if (key === 'Type') signature.type = value
        else if (key === 'Hash') signature.hash = value
        break
      }
      case 'none':
      case 'skip':
        break
    }
  }
  if (section === 'shapes') flushShape()

  if (!version) throw new SeqParseError('not a pulseq file: no [VERSION] section found')
  if (version.combined < 1004000) {
    throw new SeqParseError(
      `pulseq format ${version.major}.${version.minor}.${version.revision} is not supported (1.4.0 or later required)`,
    )
  }

  const rasters = {
    blockDuration:
      (definitions.get('BlockDurationRaster') as number[])?.[0] ?? DEFAULT_RASTERS.blockDuration,
    gradient:
      (definitions.get('GradientRasterTime') as number[])?.[0] ?? DEFAULT_RASTERS.gradient,
    rf: (definitions.get('RadiofrequencyRasterTime') as number[])?.[0] ?? DEFAULT_RASTERS.rf,
    adc: (definitions.get('AdcRasterTime') as number[])?.[0] ?? DEFAULT_RASTERS.adc,
  }

  let totalDuration = 0
  for (const row of blockRows) {
    // v1.4+: id dur rf gx gy gz adc ext
    if (row.length < 8) throw new SeqParseError(`malformed [BLOCKS] row (${row.length} columns)`)
    const duration = row[1] * rasters.blockDuration
    blocks.push({
      id: row[0],
      duration,
      rfId: row[2],
      gxId: row[3],
      gyId: row[4],
      gzId: row[5],
      adcId: row[6],
      extId: row[7],
    })
    totalDuration += duration
  }

  if (signature?.type?.toLowerCase() === 'md5') {
    // The hash covers the file up to (excluding) the newline that precedes
    // [SIGNATURE] — that newline was added together with the section.
    const idx = text.indexOf('\n[SIGNATURE]')
    if (idx >= 0) {
      signature.valid = md5Hex(text.slice(0, idx)) === signature.hash.toLowerCase()
    }
  }

  return {
    version,
    definitions,
    rasters,
    blocks,
    rf,
    grads,
    adcs,
    shapes,
    extensions,
    extensionSpecs,
    signature,
    totalDuration,
  }
}