#!/usr/bin/env bash
# Docker-based SGLang server lifecycle helpers for Kunlun P800.
# Usage: source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/../scripts/common/server_docker.sh"

set -Eeuo pipefail

_DOCKER_COMMON_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=/dev/null
source "${_DOCKER_COMMON_DIR}/lib.sh"
# shellcheck source=/dev/null
source "${_DOCKER_COMMON_DIR}/platform.sh"

# Required platform variables: DOCKER_IMAGE, CONTAINER_NAME, DEFAULT_PORT, PATCH_ROOT

docker_server_status() {
  docker inspect -f '{{.State.Status}}' "$CONTAINER_NAME" 2>/dev/null || echo "missing"
}

docker_server_stop() {
  local container="${1:-$CONTAINER_NAME}"
  log "stopping container ${container}"
  docker rm -f "$container" 2>/dev/null || true

  # Clean up leftover shared memory on the host.
  ipcs -m 2>/dev/null | awk '$4 == 666 {print $2}' | while read -r shmid; do
    ipcrm -m "$shmid" 2>/dev/null || true
  done

  # Kill lingering sglang processes just in case.
  pkill -9 -f 'sglang.launch_server' 2>/dev/null || true
  pkill -9 -f 'multiprocessing.spawn' 2>/dev/null || true
  pkill -9 -f 'multiprocessing.resource_tracker' 2>/dev/null || true
  sleep 2
}

# Build device args for docker run (8 XPU devices).
_docker_device_args() {
  local args=""
  for i in 0 1 2 3 4 5 6 7; do
    args="${args} --device /dev/xpu${i}:/dev/xpu${i}"
  done
  args="${args} --device /dev/xpuctrl:/dev/xpuctrl"
  echo "$args"
}

# Start the SGLang server container for DeepSeek-V4 on P800.
# Args:
#   $1: model path on host (e.g. /data1/models/DeepSeek-V4-Flash-INT8)
#   $2: port (optional, defaults to DEFAULT_PORT)
#   $3: server log path (optional)
#   $4: server mode (optional): "fp8" or "w8a8_int8" (default "fp8")
docker_server_start() {
  local model_path="$1"
  local port="${2:-$DEFAULT_PORT}"
  local server_log="${3:-${RESULT_ROOT:-/tmp}/logs/server.outer.log}"
  local mode="${4:-fp8}"
  local extra_args="${5:-${SGLANG_EXTRA_LAUNCH_ARGS:-}}"
  local container="${CONTAINER_NAME}"
  local image="${DOCKER_IMAGE}"

  mkdir -p "$(dirname "$server_log")"

  log "starting container ${container} from ${image}"
  log "model path: ${model_path} -> /models"
  log "port: ${port}"
  log "mode: ${mode}"
  log "server log: ${server_log}"

  docker_server_stop "$container"

  local device_args
  device_args="$(_docker_device_args)"

  # Ensure patch files exist under PATCH_ROOT. Fall back to /tmp if the repo
  # copy is not present (legacy path).
  local patch_root="${PATCH_ROOT:-/tmp}"

  # Common environment variables shared by all modes.
  local env_args=(
    -e XPU_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
    -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
    -e CUDA_DEVICE_ORDER=OAM_ID
    -e SGLANG_USE_TRANSFORMERS_V5_TOKENIZER=1
    -e XMLIR_FORCE_USE_XPU_GRAPH=1
    -e SGLANG_DSV4_MODE=2604
    -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
    -e SGLANG_NSA_DUAL_STREAM=true
    -e SGLANG_NSA_QUANT_WQ_B_WK=false
    -e SGLANG_FP8_PAGED_MQA_LOGITS_TORCH=1
    -e SGLANG_OPT_DEEPGEMM_HC_PRENORM=false
    -e SGLANG_OPT_USE_TILELANG_MHC_PRE=1
    -e SGLANG_OPT_USE_TILELANG_MHC_POST=1
    -e SGLANG_CLEAN_REQUEST_WHEN_RETRACT=1
    -e SGLANG_SET_CPU_AFFINITY=1
    -e SGLANG_OPT_USE_KLX_TOPK_KERNEL=1
    -e XSGL_INTERTYPE_BFP16=1
    -e ENABLE_FAST_BFP16_ATTN=1
    -e XSGL_USE_DEEP_GEMM_BMM=1
    -e XSGL_XDNN_QUANT=1
    -e XSGL_FUSE_RMS_NORM_QUANT=1
    -e XSGL_TRANSPOSE_MATMUL_WEIGHT=1
    -e XINFER_QUANT_SDNN=1
    -e XSGL_USE_MOE_SIGMOID_GROUP_TOPK_NORM=1
    -e XSGL_EARLY_FIRST_TOKEN=1
    -e XSGL_ENABLE_TGEMM_FP16=1
    -e SGLANG_ENABLE_SPEC_V2=True
    -e SGLANG_ENABLE_OVERLAP_PLAN_STREAM=1
    -e PYTHONDONTWRITEBYTECODE=1
    -e XTORCH_OPS_LIB_DIR=/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/xtorch_ops
    -e XPU_RUNTIME_LIB_DIR=/root/miniconda/envs/python310_torch25_cuda/xcudart/lib
    -e BKCL_TREE_THRESHOLD=1048576
    -e CUDA_ENABLE_P2P_NO_UVA=1
    -e NCCL_IB_GID_INDEX=3
    -e IS_DSV4=1
    -e MC_CUSTOM_TOPO_JSON=/workspace/nic_priority_matrix_test.json
  )

  # Mode-specific environment variables.
  if [[ "$mode" == "w8a8_int8" || "$mode" == "w8a8_int8_baseline" ]]; then
    env_args+=(
      -e SGLANG_DSV4_FP4_EXPERTS=false
      -e SGLANG_APPLY_CONFIG_BACKUP=auto
      -e BKCL_ENABLE_XDR=1
      -e BKCL_RDMA_NICS=eth1,eth1,eth3,eth3,eth5,eth5,eth7,eth7
      -e BKCL_RDMA_VERBS=1
      -e XSGL_INT8_LM_HEAD=1
      -e SGLANG_P800_ALL_GATHER_FALLBACK=0
    )
  else
    env_args+=(
      -e SGLANG_DSV4_FP4_EXPERTS=true
      -e SGLANG_OPT_FUSE_WQA_WKV=false
      -e SGLANG_DSV4_2604_SUBMODE=2604B
      -e SGLANG_APPLY_CONFIG_BACKUP=small
      -e BKCL_ENABLE_XDR=0
      -e BKCL_FORCE_PCIE_P2P=1
      -e BKCL_ENABLE_PCIE_P2P=1
      -e BKCL_SOCKET_IFNAME=bond0
      -e BKCL_NET_ENABLE_INTRA=1
      -e XSHMEM_MODE=0
      -e SGLANG_P800_ALL_GATHER_FALLBACK=1
    )
  fi

  # Mode-specific launch arguments.
  local launch_args
  if [[ "$mode" == "w8a8_int8" ]]; then
    launch_args="--host 0.0.0.0 --port ${port} --model-path /models --attention-backend nsa --nsa-prefill klxdsa --nsa-decode klxdsa --trust-remote-code --disable-custom-all-reduce --chunked-prefill-size 8192 --page-size 64 --mem-fraction-static 0.8 --max-prefill-tokens 16384 --max-running-requests 64 --tensor-parallel-size 8 --ep-size 8 --disable-shared-experts-fusion --quantization w8a8_int8 --kv-cache-dtype float16 --disable-piecewise-cuda-graph --cuda-graph-max-bs 32 --watchdog-timeout 3000000 --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --constrained-json-disable-any-whitespace --enable-metrics --enable-request-time-stats-logging"
  elif [[ "$mode" == "w8a8_int8_baseline" ]]; then
    # Baseline w8a8_int8 launch without speculative decoding and without
    # hard-coded chunked-prefill / memory-fraction settings. Matches the
    # long-context probe tuning used for P800.
    launch_args="--host 0.0.0.0 --port ${port} --model-path /models --attention-backend nsa --nsa-prefill klxdsa --nsa-decode klxdsa --trust-remote-code --disable-custom-all-reduce --tensor-parallel-size 8 --ep-size 8 --disable-shared-experts-fusion --quantization w8a8_int8 --kv-cache-dtype float16 --disable-piecewise-cuda-graph --cuda-graph-max-bs 32 --watchdog-timeout 3000000 --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 --constrained-json-disable-any-whitespace --enable-metrics --enable-request-time-stats-logging"
  else
    launch_args="--host 0.0.0.0 --port ${port} --model-path /models --quantization fp8 --attention-backend nsa --nsa-prefill klxdsa --nsa-decode klxdsa --trust-remote-code --chunked-prefill-size 8192 --page-size 64 --mem-fraction-static 0.85 --max-prefill-tokens 16384 --context-length 65536 --max-running-requests 8 --max-total-tokens 524288 --tensor-parallel-size 8 --ep-size 8 --moe-runner-backend deep_gemm --disable-shared-experts-fusion --kv-cache-dtype float16 --disable-piecewise-cuda-graph --disable-cuda-graph --watchdog-timeout 3000000 --tool-call-parser deepseekv4 --reasoning-parser deepseek-v4 --enable-metrics --enable-request-time-stats-logging"
  fi

  if [[ -n "$extra_args" ]]; then
    launch_args="${launch_args} ${extra_args}"
    log "extra launch args: ${extra_args}"
  fi

  # Build the full server bootstrap command and base64-encode it to avoid
  # host-shell quoting hell.
  local server_cmd
  server_cmd=$(cat <<EOF
cd /workspace
find /root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true
/root/miniconda/envs/python310_torch25_cuda/bin/pip install --upgrade safetensors -q
/root/miniconda/envs/python310_torch25_cuda/bin/pip install https://files.pythonhosted.org/packages/14/8b/2a1333a6455c6fad401c2285dee6f58016c55b1cb44cae3a31f8a9cc7d83/apache_tvm_ffi-0.1.0b2-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl -q
/root/miniconda/envs/python310_torch25_cuda/bin/python -c "import torch; torch.float8_e8m0fnu = torch.uint8; import runpy, sys; sys.argv[0] = 'sglang.launch_server'; runpy.run_module('sglang.launch_server', run_name='__main__')" ${launch_args}
EOF
)
  local server_cmd_b64
  server_cmd_b64=$(printf '%s' "$server_cmd" | base64 -w0)

  # Patch mounts. For w8a8_int8 the working container relies on patches baked
  # into the image, so we only mount the runtime config, nic matrix, and a
  # local dummy dataset for offline benchmarking.
  local patch_mounts=(
    -v "${patch_root}/nic_priority_matrix_test.json:/workspace/nic_priority_matrix_test.json:ro"
    -v "${patch_root}/dummy_sharegpt.json:/workspace/dummy_sharegpt.json:ro"
  )
  if [[ "$mode" == "fp8" ]]; then
    patch_mounts+=(
      -v "${patch_root}/sitecustomize_xpu.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sitecustomize.py:ro"
      -v "${patch_root}/hf_transformers_utils.py.patched:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/utils/hf_transformers_utils.py:ro"
      -v "${patch_root}/fp8_utils.py.patched:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/layers/quantization/fp8_utils.py:ro"
      -v "${patch_root}/config_backup_small_fp8.json:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/configs/config_backup_small.json:ro"
      -v "${patch_root}/parallel_state.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/distributed/parallel_state.py:ro"
      -v "${patch_root}/vocab_parallel_embedding.py:/root/miniconda/envs/python310_torch25_cuda/lib/python3.10/site-packages/sglang/srt/layers/vocab_parallel_embedding.py:ro"
    )
  fi

  docker run -d \
    --name "${container}" \
    --privileged \
    --network host \
    --ipc host \
    ${device_args} \
    -v "${model_path}:/models:ro" \
    -v "${model_path}:${model_path}:ro" \
    "${patch_mounts[@]}" \
    "${env_args[@]}" \
    "${image}" \
    bash -c "echo '${server_cmd_b64}' | base64 -d | bash" \
    >> "${server_log}" 2>&1

  local health_wait="${SGLANG_HEALTH_CHECK_TIMEOUT:-600}"
  log "container ${container} started, waiting for health (timeout=${health_wait}s)"
  if health_check 127.0.0.1 "$port" "$health_wait"; then
    log "container ${container} is healthy"
  else
    log "ERROR: container ${container} failed health check after ${health_wait}s"
    return 1
  fi
}