#!/bin/bash
# Argument parser for training qwen.

parse_qwen_args()
{
  parse_model_size "$@"
  parse_parallelism "$@"
  parse_other_args "$@"
}

parse_model_size()
{
  # default values
  model_size="8"

  # manual parsing
  while [[ $# -gt 0 ]]; do
    case $1 in
      --model_size)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        model_size="$2"
        shift 2
        ;;
      --*)
        # 未知的长选项：假设它带一个值，跳过两个参数
        if [[ -n "$2" ]] && [[ "$2" != -* ]]; then
          shift 2
        else
          # 如果下一个参数不存在或也是选项，只跳过当前
          shift 1
        fi
        ;;
      *)
        # 非 -- 开头的未知参数（如位置参数或其他），直接跳过一个
        shift 1
        ;;
    esac
  done
}

parse_parallelism()
{
  # default values
  tp="2"
  pp="1"
  cp="1"
  ep="0"
  etp="0"

  # manual parsing
  while [[ $# -gt 0 ]]; do
    case $1 in
      --tp)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        tp="$2"
        shift 2
        ;;
      --pp)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        pp="$2"
        shift 2
        ;;
      --cp)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        cp="$2"
        shift 2
        ;;
      --ep)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        ep="$2"
        shift 2
        ;;
      --etp)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        etp="$2"
        shift 2
        ;;
      --*)
        # 未知长选项：尝试跳过它和它的值（如果值存在且不是选项）
        if [[ -n "$2" ]] && [[ "$2" != -* ]]; then
          shift 2
        else
          shift 1
        fi
        ;;
      *)
        # 其他未知参数（如位置参数），跳过一个
        shift 1
        ;;
    esac
  done
}

parse_other_args()
{
  # default values
  mbs="1"
  mbn="64"
  seq_len="32768"
  train_iters="30"
  precision="bf16"
  checkpoint_path=""
  data_source=""

  # manual parsing
  while [[ $# -gt 0 ]]; do
    case $1 in
      --mbs)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        mbs="$2"
        shift 2
        ;;
      --mbn)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        mbn="$2"
        shift 2
        ;;
      --seq_len)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        seq_len="$2"
        shift 2
        ;;
      --train_iters)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        train_iters="$2"
        shift 2
        ;;
      --precision)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        precision="$2"
        shift 2
        ;;
      --checkpoint_path)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        checkpoint_path="$2"
        shift 2
        ;;
      --data_source)
        if [[ -z "$2" ]] || [[ "$2" == -* ]]; then
          echo "Error: '$1' requires a non-empty argument." >&2
          exit 1
        fi
        data_source="$2"
        shift 2
        ;;
      --*)
        # 未知长选项：跳过它和它的值（如果值存在且不是选项）
        if [[ -n "$2" ]] && [[ "$2" != -* ]]; then
          shift 2
        else
          shift 1
        fi
        ;;
      *)
        # 其他未知参数（如位置参数），跳过一个
        shift 1
        ;;
    esac
  done
}
