#!/usr/bin/env bash
# Usage: query_one.sh KIND MODE SPECIMEN
# KIND = presence | count
# MODE = dense | sparse — `merge` packs sparse by default (2026-08-28), so
#        "sparse" is global_index_KIND itself, unmodified; "dense" is the
#        explicitly repacked global_index_KIND_dense (see pack_dense.sh).
# SPECIMEN = "species--strain" (Make pattern stem), reads from query_data/
# Outputs:
#   query_KIND_MODE/SPECIMEN.fasta.gz     (obikmer query output, --count-missing)
#   stats/query_KIND_MODE/SPECIMEN.stats  (one CSV data row, no header)
set -euo pipefail

KIND="$1"
MODE="$2"
SPECIMEN="$3"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BINARY="${SCRIPT_DIR}/../src/target/release/obikmer"

case "${KIND}" in
    presence|count) ;;
    *) echo "ERROR: unknown kind '${KIND}' (expected presence|count)" >&2; exit 1 ;;
esac

case "${MODE}" in
    sparse) INDEX="${SCRIPT_DIR}/global_index_${KIND}" ;;
    dense)  INDEX="${SCRIPT_DIR}/global_index_${KIND}_dense" ;;
    *) echo "ERROR: unknown mode '${MODE}' (expected dense|sparse)" >&2; exit 1 ;;
esac

species="${SPECIMEN%%--*}"
strain="${SPECIMEN#*--}"

READS_DIR="${SCRIPT_DIR}/query_data/${species}/${strain}"
OUT_DIR="${SCRIPT_DIR}/query_${KIND}_${MODE}"
STATS_DIR="${SCRIPT_DIR}/stats/query_${KIND}_${MODE}"
OUT_FILE="${OUT_DIR}/${SPECIMEN}.fasta.gz"
STATS_FILE="${STATS_DIR}/${SPECIMEN}.stats"

mkdir -p "${OUT_DIR}" "${STATS_DIR}"

r1="${READS_DIR}/reads_R1.fastq.gz"
r2="${READS_DIR}/reads_R2.fastq.gz"
if [[ ! -f "${r1}" || ! -f "${r2}" ]]; then
    echo "ERROR: reads not found in ${READS_DIR}" >&2
    exit 1
fi

echo "[${SPECIMEN}] query (${KIND}/${MODE}) → ${OUT_FILE}"

STDERR_LOG=$(mktemp)
trap 'rm -f "${STDERR_LOG}"' EXIT

"${BINARY}" query \
    --count-missing \
    "${INDEX}" "${r1}" "${r2}" \
    2>"${STDERR_LOG}" \
    | gzip >"${OUT_FILE}"

cat "${STDERR_LOG}" >&2

python3 - "${species}" "${strain}" "${STDERR_LOG}" <<'PYEOF' >"${STATS_FILE}"
import sys, re

species, strain, logfile = sys.argv[1], sys.argv[2], sys.argv[3]

def strip_ansi(s):
    return re.sub(r'\x1b\[[\x30-\x3f]*[\x20-\x2f]*[\x40-\x7e]', '', s)

def parse_wall(s):
    s = s.strip()
    if s.endswith('ms'): return float(s[:-2]) / 1000.0
    if s.endswith('s'):  return float(s[:-1])
    return 0.0

def parse_rss(s):
    m = re.match(r'([\d.]+)\s*(GB|MB|KB|B)', s.strip())
    if not m: return 0
    return int(float(m.group(1)) * {'GB': 1<<30, 'MB': 1<<20, 'KB': 1024, 'B': 1}[m.group(2)])

def is_sep(s):
    return bool(s) and not re.search(r'[A-Za-z0-9]', s)

stats = {}
state = 'scan'
with open(logfile, errors='replace') as fh:
    for raw in fh:
        line = strip_ansi(raw.rstrip('\n'))
        s    = line.strip()
        if state == 'scan':
            if re.search(r'\bstage\b.*\bwall\b', line):
                state = 'in_header'
        elif state == 'in_header':
            if is_sep(s): state = 'rows'
        elif state == 'rows':
            if is_sep(s): state = 'total'
            elif s:
                parts = re.split(r'  +', s)
                if len(parts) >= 4:
                    stats[parts[0]] = (parse_wall(parts[1]), parse_rss(parts[3]))
        elif state == 'total':
            if s:
                parts = re.split(r'  +', s)
                if len(parts) >= 3:
                    stats['TOTAL'] = (parse_wall(parts[1]),
                                      parse_rss(parts[3]) if len(parts) > 3 else 0)
            break

qw, qr = stats.get('query', ('', ''))
tw, tr = stats.get('TOTAL', ('', ''))
row = [species, strain,
       f'{qw:.3f}' if isinstance(qw, float) else '', str(qr),
       f'{tw:.3f}' if isinstance(tw, float) else '', str(tr)]
print(','.join(row))
PYEOF