#!/usr/bin/env bash

set -euo pipefail

# Validate that exactly 2 arguments are provided
if [ $# -ne 2 ]; then
  echo "Error: Exactly 2 arguments required."
  echo "Usage: $0 <llama_cpp_dir> <original_safetensors_path>"
  echo "Example: $0 ~/code/llama.cpp ../../Qwen/Qwen3.5-397B-A10B"
  exit 1
fi

# Assign arguments to variables for clarity
LLAMA_CPP_DIR="$1"
ORIGINAL_SAFETENSORS_PATH="$2"

# Validate that the llama.cpp directory exists
if [ ! -d "$LLAMA_CPP_DIR" ]; then
  echo "Error: llama.cpp directory not found: $LLAMA_CPP_DIR"
  exit 1
fi

# Construct the path to the conversion script
CONVERT_TO_GGUF_PATH="$LLAMA_CPP_DIR/convert_hf_to_gguf.py"

# Validate that the Python script exists
if [ ! -f "$CONVERT_TO_GGUF_PATH" ]; then
  echo "Error: Python script not found: $CONVERT_TO_GGUF_PATH"
  exit 1
fi

# Validate that the model path exists
if [ ! -e "$ORIGINAL_SAFETENSORS_PATH" ]; then
  echo "Error: Model path not found: $ORIGINAL_SAFETENSORS_PATH"
  exit 1
fi

# Get the directory where the script is located
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Resolve GGUF_OUT_DIR absolutely (parent of script directory)
# We cd into the parent and run pwd to get the absolute path
if [ ! -d "$SCRIPT_DIR/.." ]; then
  echo "Error: Could not resolve parent directory for output."
  exit 1
fi
GGUF_OUT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"

# Define BF16_DIR inside GGUF_OUT_DIR
BF16_DIR="$GGUF_OUT_DIR/BF16"
mkdir -pv "$BF16_DIR"

# Activate virtual environment inside llama.cpp directory
cd "$LLAMA_CPP_DIR"
source .venv/bin/activate
cd -

# # Execute the Python script with the provided arguments
# # Main conversion output goes to BF16_DIR
# "$CONVERT_TO_GGUF_PATH" --fuse-gate-up-exps --no-tensor-first-split --split-max-size 50G --outfile "$BF16_DIR/." "$ORIGINAL_SAFETENSORS_PATH"
#
# # mmproj output goes to GGUF_OUT_DIR
# "$CONVERT_TO_GGUF_PATH" --fuse-gate-up-exps --mmproj --outtype bf16 --outfile "$GGUF_OUT_DIR/mmproj-BF16.gguf" "$ORIGINAL_SAFETENSORS_PATH"
# # also do a FP16 version
# "$CONVERT_TO_GGUF_PATH" --fuse-gate-up-exps --mmproj --outtype f16 --outfile "$GGUF_OUT_DIR/mmproj-F16.gguf" "$ORIGINAL_SAFETENSORS_PATH"

"$CONVERT_TO_GGUF_PATH" --no-tensor-first-split --split-max-size 50G --outfile "$BF16_DIR/." "$ORIGINAL_SAFETENSORS_PATH"

# mmproj output goes to GGUF_OUT_DIR
"$CONVERT_TO_GGUF_PATH" --mmproj --outtype bf16 --outfile "$GGUF_OUT_DIR/mmproj-BF16.gguf" "$ORIGINAL_SAFETENSORS_PATH"
# also do a FP16 version
"$CONVERT_TO_GGUF_PATH" --mmproj --outtype f16 --outfile "$GGUF_OUT_DIR/mmproj-F16.gguf" "$ORIGINAL_SAFETENSORS_PATH"