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

MEGATRON_PATH="${MEGATRON_PATH:-/opt/Megatron-Bridge/3rdparty/Megatron-LM}"
SERVER_PY="${MEGATRON_PATH}/tools/run_text_generation_server.py"
ENGINE_PY="${MEGATRON_PATH}/megatron/core/inference/text_generation_server/run_mcore_engine.py"

python3 - <<'PY'
import importlib.util
missing = [name for name in ["flask", "flask_restful"] if importlib.util.find_spec(name) is None]
if missing:
    raise SystemExit(f"missing python packages: {missing}; install with: pip install flask-restful")
print("flask dependencies ok")
PY

python3 - "$SERVER_PY" "$ENGINE_PY" <<'PY'
from pathlib import Path
import sys

server = Path(sys.argv[1])
engine = Path(sys.argv[2])

text = server.read_text()
old = "inference_context = StaticInferenceContext(args.inference_max_requests, args.inference_max_sequence_length)"
new = "inference_context = StaticInferenceContext(args.inference_max_requests, getattr(args, 'inference_max_sequence_length', getattr(args, 'inference_max_seq_length', args.seq_length)))"
if old in text:
    text = text.replace(old, new)
    server.write_text(text)
    print(f"patched {server}: inference_max_sequence_length fallback")
elif new in text:
    print(f"already patched {server}: inference_max_sequence_length fallback")
else:
    print(f"warning: expected inference_context line not found in {server}")

text = engine.read_text()
old = '"tokens": [x.prompt_tokens + x.generated_tokens.tolist() for x in result],'
new = '"tokens": [x.prompt_tokens.tolist() + x.generated_tokens if hasattr(x.prompt_tokens, "tolist") else x.prompt_tokens + x.generated_tokens for x in result],'
if old in text:
    text = text.replace(old, new)
    engine.write_text(text)
    print(f"patched {engine}: token list concatenation")
elif "x.prompt_tokens.tolist() + x.generated_tokens" in text:
    print(f"already patched {engine}: token list concatenation")
else:
    print(f"warning: expected token concat line not found in {engine}")

text = engine.read_text()
old = "        skip_prompt_log_probs=False,"
new = "        skip_prompt_log_probs=True,"
if old in text:
    text = text.replace(old, new)
    engine.write_text(text)
    print(f"patched {engine}: skip prompt logprobs")
elif new in text:
    print(f"already patched {engine}: skip prompt logprobs")
else:
    print(f"warning: expected skip_prompt_log_probs line not found in {engine}")

text = engine.read_text()
old = "            response_logprobs = [x.prompt_log_probs + x.generated_log_probs for x in result]"
new = "            response_logprobs = [x.generated_log_probs for x in result]"
if old in text:
    text = text.replace(old, new)
    engine.write_text(text)
    print(f"patched {engine}: generated-only logprobs response")
elif new in text:
    print(f"already patched {engine}: generated-only logprobs response")
else:
    print(f"warning: expected logprobs response line not found in {engine}")
PY

echo "Megatron inference patch setup complete."