import { unzip, unzipChunkSlice } from '@gmod/bgzf-filehandle'
import crc32 from 'crc/calculators/crc32'
import { LocalFile, RemoteFile } from 'generic-filehandle2'

import BAI from './bai.ts'
import CSI from './csi.ts'
import NullFilehandle from './nullFilehandle.ts'
import BAMFeature from './record.ts'
import { parseHeaderText } from './sam.ts'
import { appendInRange, parseRefSeqs } from './util.ts'

import type Chunk from './chunk.ts'
import type { Bytes } from './record.ts'
import type { BamOpts, BaseOpts } from './util.ts'
import type { GenericFilehandle } from 'generic-filehandle2'

export interface BamRecordLike {
  ref_id: number
  start: number
  end: number
  name: string
  fileOffset: number
  next_pos: number
  next_refid: number
  flags: number
  tags: Record<string, unknown>
}

export type BamRecordClass<T extends BamRecordLike = BAMFeature> = new (args: {
  bytes: Bytes
  fileOffset: number
  dataView: DataView
}) => T

export const BAM_MAGIC = 21840194

const blockLen = 1 << 16

function resolveFilehandle(
  filehandle?: GenericFilehandle,
  path?: string,
  url?: string,
) {
  return (
    filehandle ??
    (path ? new LocalFile(path) : url ? new RemoteFile(url) : undefined)
  )
}

interface ChunkEntry<T> {
  // decompressed size of the chunk these features are views into
  bytes: number
  features: T[]
}

function chunkCacheKey(chunk: Chunk) {
  const { minv, maxv } = chunk
  return `${minv.blockPosition}:${minv.dataPosition}-${maxv.blockPosition}:${maxv.dataPosition}`
}

// Every record in an entry is a view into its chunk's decompressed buffer, so
// caching one entry pins that whole buffer — 8MB apiece on the nanopore and
// 2kb-read test files, and optimizeChunks merges spans up to 5MB *compressed*,
// so tens of MB is possible. A count-based LRU therefore gives no bound on
// memory at all, which is why this budgets by decompressed bytes instead. It is
// the *only* bound: a query keeps every chunk it parsed, since a chunk dropped
// here costs a re-download and re-decompress the next time the view moves.
export const DEFAULT_MAX_CACHE_BYTES = 100 * 1024 * 1024

class ChunkFeatureCache<T> {
  public maxBytes: number
  private entries = new Map<string, ChunkEntry<T>>()
  private bytes = 0

  constructor(maxBytes: number) {
    this.maxBytes = maxBytes
  }

  get size() {
    return this.entries.size
  }

  get byteSize() {
    return this.bytes
  }

  get(key: string) {
    const entry = this.entries.get(key)
    if (entry) {
      // re-insert so Map iteration order stays least-recently-used first
      this.entries.delete(key)
      this.entries.set(key, entry)
    }
    return entry
  }

  set(key: string, entry: ChunkEntry<T>) {
    this.delete(key)
    this.entries.set(key, entry)
    this.bytes += entry.bytes
    // Evict from the least-recently-used end. The size > 1 guard means a single
    // chunk larger than the whole budget is still kept: the caller needs it for
    // the query in flight, and dropping it would only force a re-decompress.
    const lru = this.entries.keys()
    while (this.bytes > this.maxBytes && this.entries.size > 1) {
      this.delete(lru.next().value!)
    }
  }

  delete(key: string) {
    const entry = this.entries.get(key)
    if (entry) {
      this.entries.delete(key)
      this.bytes -= entry.bytes
    }
  }

  clear() {
    this.entries.clear()
    this.bytes = 0
  }

  [Symbol.iterator]() {
    return this.entries[Symbol.iterator]()
  }
}

export default class BamFile<T extends BamRecordLike = BAMFeature> {
  public renameRefSeq: (a: string) => string
  public bam: GenericFilehandle
  public header?: string
  public chrToIndex?: Record<string, number>
  public indexToChr?: { refName: string; length: number }[]
  public index?: BAI | CSI
  public htsget = false
  public headerP?: ReturnType<BamFile<T>['getHeaderPre']>

  // Cache for parsed features by chunk, bounded by decompressed bytes
  public chunkFeatureCache: ChunkFeatureCache<T>

  private RecordClass: BamRecordClass<T>

  constructor({
    bamFilehandle,
    bamPath,
    bamUrl,
    baiPath,
    baiFilehandle,
    baiUrl,
    csiPath,
    csiFilehandle,
    csiUrl,
    htsget,
    renameRefSeqs = n => n,
    recordClass,
    maxCacheBytes = DEFAULT_MAX_CACHE_BYTES,
  }: {
    bamFilehandle?: GenericFilehandle
    bamPath?: string
    bamUrl?: string
    baiPath?: string
    baiFilehandle?: GenericFilehandle
    baiUrl?: string
    csiPath?: string
    csiFilehandle?: GenericFilehandle
    csiUrl?: string
    renameRefSeqs?: (a: string) => string
    htsget?: boolean
    recordClass?: BamRecordClass<T>
    /** budget for the parsed-chunk cache, in decompressed bytes */
    maxCacheBytes?: number
  }) {
    this.renameRefSeq = renameRefSeqs
    this.RecordClass = (recordClass ?? BAMFeature) as BamRecordClass<T>
    this.chunkFeatureCache = new ChunkFeatureCache<T>(maxCacheBytes)

    const bamFh = resolveFilehandle(bamFilehandle, bamPath, bamUrl)
    if (bamFh) {
      this.bam = bamFh
    } else if (htsget) {
      this.htsget = true
      this.bam = new NullFilehandle()
    } else {
      throw new Error(
        'no bam source: pass bamFilehandle, bamPath, bamUrl, or htsget: true',
      )
    }

    const csiFh = resolveFilehandle(csiFilehandle, csiPath, csiUrl)
    const baiFh =
      resolveFilehandle(baiFilehandle, baiPath, baiUrl) ??
      resolveFilehandle(
        undefined,
        bamPath ? `${bamPath}.bai` : undefined,
        bamUrl ? `${bamUrl}.bai` : undefined,
      )
    if (csiFh) {
      this.index = new CSI({ filehandle: csiFh })
    } else if (baiFh) {
      this.index = new BAI({ filehandle: baiFh })
    } else if (!htsget) {
      throw new Error(
        'no index source: pass csi*/bai* options or a bamPath/bamUrl so the .bai sibling can be inferred',
      )
    }
    // htsget mode operates without a parsed index
  }

  async getHeaderPre(opts: BaseOpts = {}) {
    if (!this.index) {
      return undefined
    }
    const indexData = await this.index.parse(opts)

    // firstDataLine is not defined in cases where there is no data in the file
    // (just bam header and nothing else)
    const readLen =
      indexData.firstDataLine === undefined
        ? undefined
        : indexData.firstDataLine.blockPosition + blockLen

    const buffer =
      readLen === undefined
        ? await this.bam.readFile()
        : await this.bam.read(readLen, 0)
    let uncba = await unzip(buffer)
    const dataView = new DataView(
      uncba.buffer,
      uncba.byteOffset,
      uncba.byteLength,
    )

    if (dataView.getInt32(0, true) !== BAM_MAGIC) {
      throw new Error('Not a BAM file')
    }
    const headLen = dataView.getInt32(4, true)
    this.header = new TextDecoder('utf8').decode(uncba.subarray(8, 8 + headLen))

    // BAM files with many reference sequences may need more data than the
    // initial read covers. If the first attempt comes up short, fall back to
    // reading the whole file (the index's firstDataLine is just an
    // optimization hint, not a guaranteed cap on the ref-seq table size).
    const refSeqStart = headLen + 8
    let parsed = parseRefSeqs(uncba, refSeqStart, this.renameRefSeq)
    if (!parsed) {
      uncba = await unzip(await this.bam.readFile())
      parsed = parseRefSeqs(uncba, refSeqStart, this.renameRefSeq)
    }
    if (!parsed) {
      throw new Error('Insufficient data for reference sequences')
    }
    this.chrToIndex = parsed.chrToIndex
    this.indexToChr = parsed.indexToChr
    return parseHeaderText(this.header)
  }

  getHeader(opts?: BaseOpts) {
    if (!this.headerP) {
      this.headerP = this.getHeaderPre(opts).catch((e: unknown) => {
        this.headerP = undefined
        throw e
      })
    }
    return this.headerP
  }

  async getHeaderText(opts: BaseOpts = {}) {
    await this.getHeader(opts)
    return this.header
  }

  // Resolve a reference name to its numeric id, ensuring the header (which
  // populates chrToIndex) has been parsed first.
  private async getSeqId(seqName: string, opts?: BaseOpts) {
    await this.getHeader(opts)
    return this.chrToIndex?.[seqName]
  }

  async getRecordsForRange(
    chr: string,
    min: number,
    max: number,
    opts?: BamOpts,
  ) {
    const chrId = await this.getSeqId(chr, opts)
    if (chrId === undefined || !this.index) {
      return []
    }
    const chunks = await this.index.blocksForRange(chrId, min - 1, max, opts)
    return this._fetchChunkFeatures(chunks, chrId, min, max, opts)
  }

  // Parsed records for a chunk, reading and decompressing it only on a miss.
  // Every path that wants a chunk's features goes through here — mate lookups
  // included, since a viewAsPairs query revisits the same mate chunks each time
  // the view moves.
  private async _cachedChunkFeatures(chunk: Chunk, opts: BaseOpts) {
    const cacheKey = chunkCacheKey(chunk)
    let entry = this.chunkFeatureCache.get(cacheKey)
    if (!entry) {
      entry = await this._readChunkFeatures(chunk, opts)
      this.chunkFeatureCache.set(cacheKey, entry)
    }
    return entry.features
  }

  private async _fetchChunkFeatures(
    chunks: Chunk[],
    chrId: number,
    min: number,
    max: number,
    opts: BamOpts = {},
  ) {
    const { viewAsPairs, onProgress } = opts
    const result: T[] = []

    let totalBytes = 0
    for (let ci = 0, cl = chunks.length; ci < cl; ci++) {
      totalBytes += chunks[ci]!.fetchedSize()
    }
    let downloadedBytes = 0
    onProgress?.(0, totalBytes)

    for (let ci = 0, cl = chunks.length; ci < cl; ci++) {
      const chunk = chunks[ci]!
      const features = await this._cachedChunkFeatures(chunk, opts)

      downloadedBytes += chunk.fetchedSize()
      onProgress?.(downloadedBytes, totalBytes)
      appendInRange(features, chrId, min, max, result)
    }

    if (viewAsPairs) {
      const pairs = await this.fetchPairs(chrId, result, opts)
      for (let i = 0, l = pairs.length; i < l; i++) {
        result.push(pairs[i]!)
      }
    }

    return result
  }

  async fetchPairs(chrId: number, records: T[], opts: BamOpts) {
    const { pairAcrossChr, maxInsertSize = 200000 } = opts
    // Map, not a plain object: read names come from the file, and on a plain
    // object a read named "constructor" would read back Object.prototype's and
    // make the count NaN.
    const readNameCounts = new Map<string, number>()
    const readIds = new Set<number>()

    for (let i = 0, l = records.length; i < l; i++) {
      const r = records[i]!
      const name = r.name
      readNameCounts.set(name, (readNameCounts.get(name) ?? 0) + 1)
      readIds.add(r.fileOffset)
    }

    const matePromises: Promise<Chunk[]>[] = []
    for (let i = 0, l = records.length; i < l; i++) {
      const f = records[i]!
      const name = f.name
      if (
        this.index &&
        readNameCounts.get(name) === 1 &&
        (pairAcrossChr ||
          (f.next_refid === chrId &&
            Math.abs(f.start - f.next_pos) < maxInsertSize))
      ) {
        matePromises.push(
          this.index.blocksForRange(
            f.next_refid,
            f.next_pos,
            f.next_pos + 1,
            opts,
          ),
        )
      }
    }

    const map = new Map<string, Chunk>()
    const res = await Promise.all(matePromises)
    for (let i = 0, l = res.length; i < l; i++) {
      const chunks = res[i]!
      for (let j = 0, jl = chunks.length; j < jl; j++) {
        const m = chunks[j]!
        map.set(m.toString(), m)
      }
    }

    const mateFeatLists = await Promise.all(
      [...map.values()].map(async c => {
        const features = await this._cachedChunkFeatures(c, opts)
        const mateRecs = [] as T[]
        for (let i = 0, l = features.length; i < l; i++) {
          const feature = features[i]!
          if (
            readNameCounts.get(feature.name) === 1 &&
            !readIds.has(feature.fileOffset)
          ) {
            mateRecs.push(feature)
          }
        }
        return mateRecs
      }),
    )
    return mateFeatLists.flat()
  }

  async _readChunkFeatures(chunk: Chunk, opts: BaseOpts) {
    // Don't forward onProgress to the inner read: getRecordsForRange already
    // reports progress at chunk granularity (downloadedBytes/totalBytes). If the
    // filehandle's own streaming onProgress also fired this callback it would
    // report a different `total` (this chunk's size, not the whole query),
    // making the determinate bar jump around.
    const buf = await this.bam.read(
      chunk.fetchedSize(),
      chunk.minv.blockPosition,
      {
        signal: opts.signal,
      },
    )
    const {
      buffer: data,
      cpositions,
      dpositions,
    } = await unzipChunkSlice(buf, chunk)
    return {
      features: this.readBamFeatures(data, cpositions, dpositions, chunk),
      bytes: data.byteLength,
    }
  }

  readBamFeatures(
    ba: Uint8Array,
    cpositions: number[],
    dpositions: number[],
    chunk: Chunk,
  ) {
    let blockStart = 0
    const sink = [] as T[]
    let pos = 0

    const dataView = new DataView(ba.buffer, ba.byteOffset, ba.byteLength)
    const hasDpositions = dpositions.length > 0
    const hasCpositions = cpositions.length > 0

    while (blockStart + 4 < ba.length) {
      const blockSize = dataView.getInt32(blockStart, true)
      const blockEnd = blockStart + 4 + blockSize - 1

      if (hasDpositions) {
        const target = blockStart + chunk.minv.dataPosition
        while (pos < dpositions.length && target >= dpositions[pos]!) {
          pos++
        }
      }

      if (blockEnd < ba.length) {
        const feature = new this.RecordClass({
          bytes: {
            byteArray: ba,
            start: blockStart,
            end: blockEnd,
          },
          fileOffset: hasCpositions
            ? cpositions[pos]! * (1 << 8) +
              (blockStart - dpositions[pos]!) +
              chunk.minv.dataPosition +
              1
            : crc32(ba.subarray(blockStart, blockEnd)) >>> 0,
          dataView,
        })

        sink.push(feature)
      }

      blockStart = blockEnd + 1
    }
    return sink
  }

  async hasRefSeq(seqName: string, opts?: BaseOpts) {
    const seqId = await this.getSeqId(seqName, opts)
    return !this.index || seqId === undefined
      ? false
      : this.index.hasRefSeq(seqId)
  }

  async lineCount(seqName: string, opts?: BaseOpts) {
    const seqId = await this.getSeqId(seqName, opts)
    return !this.index || seqId === undefined ? 0 : this.index.lineCount(seqId)
  }

  async indexCov(seqName: string, start?: number, end?: number) {
    const seqId = await this.getSeqId(seqName)
    return !this.index || seqId === undefined
      ? []
      : this.index.indexCov(seqId, start, end)
  }

  async blocksForRange(
    seqName: string,
    start: number,
    end: number,
    opts?: BaseOpts,
  ) {
    const seqId = await this.getSeqId(seqName, opts)
    return !this.index || seqId === undefined
      ? []
      : this.index.blocksForRange(seqId, start, end, opts)
  }

  clearFeatureCache() {
    this.chunkFeatureCache.clear()
  }

  async estimatedBytesForRegions(
    regions: { refName: string; start: number; end: number }[],
    opts?: BaseOpts,
  ) {
    if (!this.index) {
      return 0
    }
    await this.getHeader(opts)
    const mapped = regions.flatMap(r => {
      const refId = this.chrToIndex?.[r.refName]
      return refId === undefined ? [] : [{ refId, start: r.start, end: r.end }]
    })
    return this.index.estimatedBytesForRegions(mapped, opts)
  }
}