import db from '@adonisjs/lucid/services/db'
import type { TransactionClientContract } from '@adonisjs/lucid/types/database'
import { UPLOAD_GAP_MS, type EdgeBatch, type ProbeCandidate, type ReleaseChunk } from '#lib/edge/index'
import ChunkPartitions from '#services/chunk_partitions'
import { decodeChunk } from '#lib/chunks/index'

export class StaleBatch extends Error {}

export interface StoredBatch {
  releases: number
  newReleases: number
  candidates: number
  fragments: number
}

/**
 * Stores one batch of a leased range (REBUILD.md, parts 3–7), in one
 * transaction: release aggregates in the catalog, one chunk per release,
 * probe candidates, rule counts and samples, edge fragments, and the
 * range's resume point. A batch that doesn't start at the resume point was
 * sent twice or too late, and is refused (StaleBatch), so the sums stay right.
 */
class BatchStore {
  async store(rangeId: number, nodeId: number, batch: EdgeBatch): Promise<StoredBatch> {
    await ChunkPartitions.ensure()
    return db.transaction(async (trx) => {
      // One batch of a range at a time, without holding the range's row for the whole transaction: the
      // prefetcher's and heartbeats' updates of it waited on it (0.5 s on average, NNTP connections idle;
      // LOG.md 2026-10-09 (13)). The resume point is checked again when it moves, at the end.
      await trx.rawQuery(`SELECT pg_advisory_xact_lock(hashtextextended('range:' || ?::text, 0))`, [rangeId])
      const range = await trx.from('ranges').where('id', rangeId).first()
      if (!range || range.status !== 'leased' || Number(range.node_id) !== nodeId)
        throw new StaleBatch(`Range ${rangeId} is not leased to this node`)
      if (Number(range.next_id) !== batch.from)
        throw new StaleBatch(
          `Range ${rangeId} is at ${range.next_id}, the batch starts at ${batch.from}`
        )
      const result: StoredBatch = {
        releases: batch.releases.length,
        newReleases: 0,
        candidates: 0,
        fragments: batch.fragments.length,
      }

      const written = await this.writeReleases(
        trx,
        rangeId,
        batch.from,
        batch.releases,
        batch.candidates
      )
      result.newReleases = written.newReleases
      result.candidates = written.candidates

      if (batch.counts.length > 0) {
        await trx.rawQuery(
          `INSERT INTO rule_counts (range_id, rule, key, articles, files)
           SELECT ?, v.rule, v.key, v.articles, v.files FROM jsonb_to_recordset(?::jsonb) AS v(rule text, key text, articles bigint, files bigint)
           ON CONFLICT (range_id, rule, key) DO UPDATE SET articles = rule_counts.articles + excluded.articles, files = rule_counts.files + excluded.files`,
          [rangeId, JSON.stringify(batch.counts)]
        )
      }

      for (const s of batch.samples) {
        await trx.insertQuery().table('rule_samples').insert({
          range_id: rangeId,
          rule: s.rule,
          rule_key: s.ruleKey,
          release_key: s.key,
          name: s.name,
          files: s.files,
          segments: s.segments,
          payload: s.payload,
        })
      }

      for (let i = 0; i < batch.fragments.length; i += 2000) {
        const rows = batch.fragments.slice(i, i + 2000).map((f) => ({
          range_id: rangeId,
          file_key: f.fileKey,
          // Not Math.min(...): a big file's fragment has more segments than the call stack takes arguments.
          article_number: f.segments.reduce((min, s) => (s[3] < min ? s[3] : min), Infinity),
          segments: f.segments.length,
          fragment: JSON.stringify(f),
        }))
        await trx.insertQuery().table('fragments').multiInsert(rows)
      }

      const moved = await trx
        .from('ranges')
        .where('id', rangeId)
        .where('status', 'leased')
        .where('node_id', nodeId)
        .where('next_id', batch.from)
        .update({
          next_id: batch.next,
          headers: db.raw('headers + ?', [batch.headers]),
          heartbeat_at: new Date(),
        })
      if (Number(moved) !== 1 && !(Array.isArray(moved) && moved.length === 1)) throw new StaleBatch(`Range ${rangeId} moved on meanwhile`)
      return result
    })
  }

  /**
   * An upload spread over groups (FTD: par2 and part1 to a.b.music.misc,
   * part2 to a.b.sounds.music, each article in one group only; LOG.md
   * 2026-10-11 (9)): a chunk without a release in its own group joins one of
   * its key in another group whose window is within UPLOAD_GAP_MS, when none
   * of its file names are in that release yet. A post of the same files to
   * another group (other message-ids, the same names) stays a release of its
   * own, or the NZB would hold every file twice. Reads the candidate's file
   * names from its chunks; candidates are few.
   */
  private async joinAcrossGroups(
    trx: TransactionClientContract,
    part: ReleaseChunk[],
    windows: { idx: number; key: string; first: number | null | undefined; last: number | null | undefined }[],
    groupId: number,
    joined: Map<number, number>
  ) {
    const open = windows.filter((w) => !joined.has(w.idx) && w.first)
    if (open.length === 0) return
    const { rows: candidates } = await trx.rawQuery(
      `SELECT v.idx, r.id
       FROM jsonb_to_recordset(?::jsonb) AS v(idx int, key text, first bigint, last bigint)
       JOIN releases r ON r.identity_key = v.key AND r.group_id <> ?
       WHERE r.posted_at <= to_timestamp((v.last + ?) / 1000.0) AND coalesce(r.last_posted_at, r.posted_at) >= to_timestamp((v.first - ?) / 1000.0)
       ORDER BY v.idx, r.last_posted_at DESC NULLS LAST, r.id DESC`,
      [JSON.stringify(open), groupId, UPLOAD_GAP_MS, UPLOAD_GAP_MS]
    )
    if (candidates.length === 0) return
    const ids: number[] = [...new Set<number>(candidates.map((c: { id: string }) => Number(c.id)))]
    const names = new Map<number, Set<string>>()
    const { rows: chunks } = await trx.rawQuery(`SELECT release_id, payload FROM release_chunks WHERE release_id = ANY(?)`, [ids])
    for (const c of chunks) {
      const set = names.get(Number(c.release_id)) ?? new Set<string>()
      for (const f of decodeChunk(c.payload)) set.add(f.name)
      names.set(Number(c.release_id), set)
    }
    for (const c of candidates as { idx: number; id: string }[]) {
      if (joined.has(c.idx)) continue
      const have = names.get(Number(c.id)) ?? new Set<string>()
      const incoming = decodeChunk(part[c.idx].payload).map((f) => f.name)
      if (incoming.every((name) => !have.has(name))) {
        joined.set(c.idx, Number(c.id))
        for (const name of incoming) have.add(name)
        names.set(Number(c.id), have)
      }
    }
  }

  /**
   * Release aggregates in the catalog, one write-once chunk per release, and
   * probe candidates (a batch's, or files the fragment merger completed).
   * A chunk joins the release of its key whose articles the range's group
   * owns and whose posting window is within UPLOAD_GAP_MS of its own, else
   * starts one: re-uploads are releases of their own (decision d). The
   * group is locked, so concurrent writers don't both start one.
   */
  async writeReleases(
    trx: TransactionClientContract,
    rangeId: number,
    batchFrom: number,
    releases: ReleaseChunk[],
    candidates: ProbeCandidate[]
  ) {
    const out = { newReleases: 0, candidates: 0 }
    const ids = new Map<string, number>()
    const slot = (key: string, upload: number) => `${key}\n${upload}`
    const { group_id: groupId } = await trx.from('ranges').where('id', rangeId).select('group_id').firstOrFail()
    // One writer per group at a time: a chunk only matches releases of its group, and a lock per key
    // ran out of shared memory (a picture group's batch has 50k keys; LOG.md 2026-10-09 (12)).
    await trx.rawQuery(`SELECT pg_advisory_xact_lock(hashtextextended('releases:' || ?::text, 0))`, [groupId])
    const sorted = [...releases].sort((a, b) => a.key.localeCompare(b.key) || a.upload - b.upload)
    for (let i = 0; i < sorted.length; i += 1000) {
      const part = sorted.slice(i, i + 1000)
      const windows = part.map((r, idx) => ({ idx, key: r.key, first: r.postedAt, last: r.postedLast ?? r.postedAt }))
      const { rows: matches } = await trx.rawQuery(
        `SELECT DISTINCT ON (v.idx) v.idx, r.id
         FROM jsonb_to_recordset(?::jsonb) AS v(idx int, key text, first bigint, last bigint)
         JOIN releases r ON r.identity_key = v.key AND r.group_id = ?
         WHERE v.first IS NULL OR r.posted_at IS NULL
            OR (r.posted_at <= to_timestamp((v.last + ?) / 1000.0) AND coalesce(r.last_posted_at, r.posted_at) >= to_timestamp((v.first - ?) / 1000.0))
         ORDER BY v.idx, r.last_posted_at DESC NULLS LAST, r.id DESC`,
        [JSON.stringify(windows), groupId, UPLOAD_GAP_MS, UPLOAD_GAP_MS]
      )
      const joined = new Map<number, number>(matches.map((m: { idx: number; id: string }) => [m.idx, Number(m.id)]))
      await this.joinAcrossGroups(trx, part, windows, groupId, joined)
      const row = (r: ReleaseChunk) => ({
        name: r.name,
        source: r.source,
        poster: r.poster,
        groups: r.groups,
        first: r.postedAt,
        last: r.postedLast ?? r.postedAt,
        file_total: r.fileTotal,
        files: r.files,
        segments: r.segments,
        bytes: r.bytes,
        first_article: r.firstArticle,
        last_article: r.lastArticle,
        readable: r.readable,
        has_par2: r.hasPar2,
      })
      const updates = part.flatMap((r, idx) => (joined.has(idx) ? [{ id: joined.get(idx)!, ...row(r) }] : []))
      if (updates.length > 0) {
        await trx.rawQuery(
          `WITH v AS (SELECT * FROM jsonb_to_recordset(?::jsonb) AS v(id bigint, first bigint, last bigint, file_total int, files int, segments bigint, bytes bigint, first_article bigint, last_article bigint, readable boolean, has_par2 boolean)),
                a AS (SELECT id, sum(files) AS files, sum(segments) AS segments, sum(bytes) AS bytes, count(*) AS chunks, min(first) AS first, max(last) AS last,
                             max(file_total) AS file_total, min(first_article) AS first_article, max(last_article) AS last_article, bool_or(readable) AS readable, bool_or(has_par2) AS has_par2
                      FROM v GROUP BY id)
           UPDATE releases r SET
             files = r.files + a.files,
             segments = r.segments + a.segments,
             bytes = r.bytes + a.bytes,
             chunks = r.chunks + a.chunks,
             posted_at = least(r.posted_at, to_timestamp(a.first / 1000.0)),
             last_posted_at = greatest(r.last_posted_at, to_timestamp(a.last / 1000.0)),
             file_total = coalesce(r.file_total, a.file_total),
             first_article = least(r.first_article, a.first_article),
             last_article = greatest(r.last_article, a.last_article),
             readable = r.readable OR a.readable,
             has_par2 = r.has_par2 OR a.has_par2,
             -- As it stands (the counter's files all found); settling makes it final.
             complete = CASE WHEN coalesce(r.file_total, a.file_total) IS NULL THEN NULL ELSE r.files + a.files >= coalesce(r.file_total, a.file_total) END,
             last_seen_at = now()
           FROM a WHERE r.id = a.id`,
          [JSON.stringify(updates)]
        )
        part.forEach((r, idx) => joined.has(idx) && ids.set(slot(r.key, r.upload), joined.get(idx)!))
      }
      const fresh = part.filter((_, idx) => !joined.has(idx))
      if (fresh.length > 0) {
        // The first upload of a key keeps the key; later ones are key:2, key:3, …
        const { rows: counts } = await trx.rawQuery(`SELECT identity_key, count(*) AS n FROM releases WHERE identity_key = ANY(?) GROUP BY 1`, [[...new Set(fresh.map((r) => r.key))]])
        const taken = new Map<string, number>(counts.map((c: { identity_key: string; n: string }) => [c.identity_key, Number(c.n)]))
        const numbered = (r: ReleaseChunk) => {
          const n = (taken.get(r.key) ?? 0) + 1
          taken.set(r.key, n)
          return { release_key: n === 1 ? r.key : `${r.key}:${n}`, key: r.key, upload: r.upload, ...row(r) }
        }
        let inserts = fresh.map(numbered)
        // Writers of other groups number the same key at the same time (a spam post crossposted to 67
        // groups; only the group is locked): a number taken meanwhile is skipped, and the next one tried.
        for (let attempt = 0; inserts.length > 0; attempt++) {
          if (attempt === 20) throw new Error(`No free release key for ${inserts[0].key}`)
          const { rows } = await trx.rawQuery(
            `INSERT INTO releases (release_key, identity_key, group_id, name, source, poster, groups, posted_at, last_posted_at, file_total, files, segments, bytes, chunks, first_article, last_article, readable, has_par2, complete)
             SELECT v.release_key, v.key, ?, v.name, v.source, v.poster, v.groups, to_timestamp(v.first / 1000.0), to_timestamp(v.last / 1000.0), v.file_total, v.files, v.segments, v.bytes, 1,
                    v.first_article, v.last_article, v.readable, v.has_par2, CASE WHEN v.file_total IS NULL THEN NULL ELSE v.files >= v.file_total END
             FROM jsonb_to_recordset(?::jsonb) AS v(release_key text, key text, name text, source text, poster text, groups jsonb, first bigint, last bigint, file_total int, files int, segments bigint, bytes bigint, first_article bigint, last_article bigint, readable boolean, has_par2 boolean)
             ON CONFLICT (release_key) DO NOTHING
             RETURNING id, release_key`,
            [groupId, JSON.stringify(inserts)]
          )
          const byReleaseKey = new Map<string, number>(rows.map((x: { id: string; release_key: string }) => [x.release_key, Number(x.id)]))
          const missed: typeof inserts = []
          for (const r of inserts) {
            const id = byReleaseKey.get(r.release_key)
            if (id === undefined) missed.push(r)
            else ids.set(slot(r.key, r.upload), id)
          }
          out.newReleases += rows.length
          const chunks = new Map(fresh.map((r) => [slot(r.key, r.upload), r]))
          inserts = missed.map((r) => numbered(chunks.get(slot(r.key, r.upload))!))
        }
      }
    }

    // One write-once chunk per release.
    for (let i = 0; i < sorted.length; i += 500) {
      const rows = sorted.slice(i, i + 500).map((r) => ({
        release_id: ids.get(slot(r.key, r.upload))!,
        range_id: rangeId,
        batch_from: batchFrom,
        files: r.files,
        segments: r.segments,
        payload: r.payload,
      }))
      if (rows.length > 0) await trx.insertQuery().table('release_chunks').multiInsert(rows)
    }

    const kept = candidates.filter((c) => ids.has(slot(c.releaseKey, c.upload ?? 0)))
    for (let i = 0; i < kept.length; i += 2000) {
      const rows = kept
        .slice(i, i + 2000)
        .map((c) => ({
          release_id: ids.get(slot(c.releaseKey, c.upload ?? 0))!,
          kind: c.kind,
          message_id: c.messageId,
          file: c.file,
        }))
      const added = await trx
        .knexQuery()
        .table('probe_candidates')
        .insert(rows)
        .onConflict(['release_id', 'kind', 'message_id'])
        .ignore()
        .returning('id')
      out.candidates += added.length
    }

    return out
  }
}

export default new BatchStore()