#!/usr/bin/env bash
set -euo pipefail

########################################
# MetaAMRplus v1.4.1 (STABLE RELEASE)
# AMR, Metal & Colocalisation Analysis
# Author: Misheck Shawa
########################################

VERSION="${METAAMRPLUS_VERSION:-1.4.1}"
COLOC_DIST=10000
MIN_CONTIG_LEN=200

########################################
# Banner
########################################
show_banner() {
cat << EOF
========================================
 MetaAMRplus v${VERSION}
 AMR, Metal & Colocalisation Analysis
 Author: Misheck Shawa
========================================
EOF
}

########################################
# DB INIT (AUTO DOWNLOADER)
########################################
init_db() {
    echo "========================================"
    echo " Initializing MetaAMRplus database"
    echo "========================================"

    if [[ -n "${CONDA_PREFIX:-}" ]]; then
        DB_BASE="$CONDA_PREFIX/share/metaamrplus"
    else
        DB_BASE="$HOME/.metaamrplus"
    fi

    DB_DIR="$DB_BASE/db"
    mkdir -p "$DB_DIR"

    echo "Downloading MetaAMRplus DB into: $DB_DIR"

    BASE_URL="https://github.com/migshaw03/MetaAMRplus/releases/download/v1.4.1"

    curl -L "$BASE_URL/MetaAMRplus_DB_v1.0.phr" -o "$DB_DIR/MetaAMRplus_DB_v1.0.phr"
    curl -L "$BASE_URL/MetaAMRplus_DB_v1.0.pin" -o "$DB_DIR/MetaAMRplus_DB_v1.0.pin"
    curl -L "$BASE_URL/MetaAMRplus_DB_v1.0.psq" -o "$DB_DIR/MetaAMRplus_DB_v1.0.psq"
    curl -L "$BASE_URL/MetaAMRplus_v1.0.idmap.tsv" -o "$DB_DIR/MetaAMRplus_v1.0.idmap.tsv"
    curl -L "$BASE_URL/MetaAMRplus_v1.0.annotations.tsv" -o "$DB_DIR/MetaAMRplus_v1.0.annotations.tsv"

    echo "Database installation complete ✔"
    exit 0
}

########################################
# HELP / VERSION
########################################
case "${1:-}" in
  --version)
    show_banner
    exit 0
    ;;
  --help|-h)
    show_banner
    cat << EOF

Usage:
  metaamrplus genome.fasta

Options:
  --init-db     Download and install required database
  --help        Show help
  --version     Show version

Environment Variables:
  METAAMRPLUS_DB       Override DB path
  METAAMRPLUS_RESULTS  Output directory

EOF
    exit 0
    ;;
  --init-db)
    init_db
    ;;
esac

########################################
# ARG CHECK
########################################
if [[ $# -ne 1 ]]; then
    show_banner
    echo "ERROR: Provide exactly one genome FASTA"
    exit 1
fi

GENOME="$1"
[[ -f "$GENOME" ]] || { echo "ERROR: File not found: $GENOME"; exit 1; }

########################################
# PATHS
########################################
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
RESULTS_DIR="${METAAMRPLUS_RESULTS:-$PWD/metaamrplus_results}"

########################################
# DB RESOLUTION (ROBUST)
########################################
if [[ -n "${METAAMRPLUS_DB:-}" ]]; then
    DB_BASE="$METAAMRPLUS_DB"
elif [[ -n "${CONDA_PREFIX:-}" ]]; then
    DB_BASE="$CONDA_PREFIX/share/metaamrplus"
else
    DB_BASE="$HOME/.metaamrplus"
fi

DB_DIR="$DB_BASE/db"
DB_PREFIX="$DB_DIR/MetaAMRplus_DB_v1.0"

IDMAP="$DB_DIR/MetaAMRplus_v1.0.idmap.tsv"
ANNOT="$DB_DIR/MetaAMRplus_v1.0.annotations.tsv"

########################################
# VALIDATE DB
########################################
for f in "$IDMAP" "$ANNOT"; do
  [[ -f "$f" ]] || {
    echo "ERROR: Missing DB file: $f"
    echo "Run: metaamrplus --init-db"
    exit 1
  }
done

for ext in phr pin psq; do
  [[ -f "${DB_PREFIX}.${ext}" ]] || {
    echo "ERROR: Missing BLAST DB: ${DB_PREFIX}.${ext}"
    exit 1
  }
done

########################################
# DEPENDENCIES
########################################
for tool in prodigal blastp awk python3 sort grep; do
  command -v "$tool" >/dev/null 2>&1 || {
    echo "ERROR: Missing dependency: $tool"
    exit 1
  }
done

########################################
# OUTPUT SETUP
########################################
PREFIX="$(basename "$GENOME")"
PREFIX="${PREFIX%.*}"
OUTDIR="$RESULTS_DIR/$PREFIX"
mkdir -p "$OUTDIR"

show_banner

########################################
# STEP 0
########################################
echo "[0/8] FASTA validation"
grep -q '^>' "$GENOME" || { echo "ERROR: Invalid FASTA"; exit 1; }

########################################
# STEP 1
########################################
echo "[1/8] Filtering contigs"

FILTERED_FASTA="$OUTDIR/${PREFIX}.filtered.fasta"

awk -v min="$MIN_CONTIG_LEN" '
  /^>/ {
    if (seq && length(seq) >= min) {
      print header
      print seq
    }
    header=$0
    seq=""
    next
  }
  { seq = seq $0 }
  END {
    if (seq && length(seq) >= min) {
      print header
      print seq
    }
  }
' "$GENOME" > "$FILTERED_FASTA"

[[ -s "$FILTERED_FASTA" ]] || {
  echo "WARNING: No contigs >= ${MIN_CONTIG_LEN}"
  exit 0
}

########################################
# STEP 2
########################################
echo "[2/8] ORF prediction"
prodigal -i "$FILTERED_FASTA" \
  -a "$OUTDIR/${PREFIX}.proteins.faa" \
  -d "$OUTDIR/${PREFIX}.genes.fna" \
  -f gff -o "$OUTDIR/${PREFIX}.gff" \
  -p meta >/dev/null 2>&1 || {
    echo "WARNING: Prodigal failed"
    exit 0
}

########################################
# STEP 3
########################################
echo "[3/8] BLASTP"
blastp -query "$OUTDIR/${PREFIX}.proteins.faa" \
  -db "$DB_PREFIX" \
  -out "$OUTDIR/raw.tsv" \
  -evalue 1e-5 \
  -outfmt "6 qseqid sseqid pident length qlen evalue bitscore"

########################################
# STEP 4
########################################
echo "[4/8] Filtering hits"

awk '$3>=95 && ($4/$5)*100>=80' "$OUTDIR/raw.tsv" \
| sort -k1,1 -k7,7nr \
| awk '!seen[$1]++' \
> "$OUTDIR/filtered.tsv"

########################################
# STEP 5
########################################
echo "[5/8] Annotation"

python3 "$SCRIPT_DIR/annotate_metaamrplus.py" \
  "$OUTDIR/filtered.tsv" \
  "$IDMAP" \
  "$ANNOT" \
  > "$OUTDIR/annotated.tsv"

########################################
# STEP 6
########################################
echo "[6/8] Remove housekeeping"

grep -v -E '\t(gyrA|gyrB|parC|parE|rpoB|rpsL|rrs|folP|murA)\t' \
  "$OUTDIR/annotated.tsv" \
  > "$OUTDIR/acquired.tsv"

########################################
# STEP 7
########################################
echo "[7/8] Coordinates"

{
echo -e "query_protein\tmetaamrplus_id\tgene\ttype\tphenotype\tmechanism\tsource\tpident\tcoverage\tbitscore\tcontig\tstart\tend\tstrand"
awk -F'\t' -v OFS='\t' '
NR==FNR && $3=="CDS" {
  split($9,a,";")
  sub(/^ID=/,"",a[1])
  loc[a[1]]=$1"\t"$4"\t"$5"\t"$7
  next
}
NR>FNR && FNR>1 {
  if ($1 in loc) print $0, loc[$1]
}
' "$OUTDIR/${PREFIX}.gff" "$OUTDIR/acquired.tsv"
} > "$OUTDIR/final.tsv"

########################################
# STEP 8
########################################
echo "[8/8] Colocalisation"

{
echo -e "AMR_gene\tmetal_gene\tcontig\tAMR_start\tAMR_end\tmetal_start\tmetal_end\tdistance_bp"
awk -F'\t' -v OFS='\t' -v D="$COLOC_DIST" '
NR==FNR && FNR>1 {amr[$11,++n[$11]]=$3"\t"$12"\t"$13; next}
NR>FNR && FNR>1 {
  for(i=1;i<=n[$11];i++){
    split(amr[$11,i],a,"\t")
    dist = ($12>a[2] ? $12-a[2] : a[2]-$12)
    if(dist<=D) print a[1],$3,$11,a[2],a[3],$12,$13,dist
  }
}
' "$OUTDIR/final.tsv" "$OUTDIR/final.tsv"
} > "$OUTDIR/colocalisation.tsv"

echo
echo "✔ DONE"
echo "✔ Version: $VERSION"
echo "✔ Output: $OUTDIR"