#!/usr/bin/env bash
# Ensures vendored mlx-audio-swift is present and patched for StreamingConfig.context.
set -euo pipefail

ROOT="$(cd "$(dirname "$0")/.." && pwd)"
VENDOR="$ROOT/ThirdParty/mlx-audio-swift"
PATCH_MARKER="$VENDOR/.osg-context-patch-v2-applied"
PINNED_COMMIT="d302a5c6080d2bb97bae38c7418f82abb76013b6"

if [[ ! -f "$VENDOR/Package.swift" ]]; then
  echo "Cloning pinned mlx-audio-swift into ThirdParty/..."
  git clone --filter=blob:none --no-checkout https://github.com/Blaizzy/mlx-audio-swift "$VENDOR"
  git -C "$VENDOR" fetch --depth 1 origin "$PINNED_COMMIT"
  git -C "$VENDOR" checkout --detach FETCH_HEAD
fi

ACTUAL_COMMIT="$(git -C "$VENDOR" rev-parse HEAD)"
if [[ "$ACTUAL_COMMIT" != "$PINNED_COMMIT" ]]; then
  echo "error: mlx-audio-swift is at $ACTUAL_COMMIT; expected $PINNED_COMMIT" >&2
  echo "Remove ThirdParty/mlx-audio-swift and rerun this script." >&2
  exit 1
fi

if [[ -f "$PATCH_MARKER" ]]; then
  echo "mlx-audio-swift context patch already applied."
  exit 0
fi

python3 <<'PY'
from pathlib import Path

root = Path("ThirdParty/mlx-audio-swift")
manifest = root / "Package.swift"
types = root / "Sources/MLXAudioSTT/Streaming/StreamingTypes.swift"
session = root / "Sources/MLXAudioSTT/Streaming/StreamingInferenceSession.swift"

# Pin mlx-swift to 0.31.3: 0.31.4+ adds an unconditional swift-argument-parser
# dependency (for its `encuda` executable) that breaks offline resolution here.
mtext = manifest.read_text()
if "mlx-swift.git\", exact:" not in mtext:
    mtext = mtext.replace(
        '.package(url: "https://github.com/ml-explore/mlx-swift.git", .upToNextMajor(from: "0.30.6")),',
        '.package(url: "https://github.com/ml-explore/mlx-swift.git", exact: "0.31.3"),',
    )
mtext = mtext.replace(
    '.package(url: "https://github.com/ml-explore/mlx-swift-lm.git", .upToNextMajor(from: "3.31.3")),',
    '.package(url: "https://github.com/ml-explore/mlx-swift-lm.git", exact: "3.31.3"),',
)
mtext = mtext.replace(
    '.package(url: "https://github.com/huggingface/swift-transformers.git", .upToNextMajor(from: "1.1.6")),',
    '.package(url: "https://github.com/huggingface/swift-transformers.git", exact: "1.1.9"),',
)
mtext = mtext.replace(
    '.package(url: "https://github.com/huggingface/swift-huggingface.git", .upToNextMajor(from: "0.8.1"))',
    '.package(url: "https://github.com/huggingface/swift-huggingface.git", exact: "0.8.1")',
)
manifest.write_text(mtext)

text = types.read_text()
if "public var context: String?" not in text:
    text = text.replace(
        "    public var language: String?\n",
        "    public var language: String?\n"
        "    /// Optional soft-prompt context (vocabulary / domain hints) injected into the system turn.\n"
        "    public var context: String?\n",
    )
    text = text.replace(
        "        language: String? = \"English\",\n        temperature: Float = 0.0,",
        "        language: String? = \"English\",\n        context: String? = nil,\n        temperature: Float = 0.0,",
    )
    text = text.replace(
        "        self.language = language\n        self.temperature = temperature",
        "        self.language = language\n        self.context = context\n        self.temperature = temperature",
    )
    types.write_text(text)

text = session.read_text()
if "context: params.config.context" not in text:
    text = text.replace(
        "            language: params.config.language\n        )",
        "            context: params.config.context ?? \"\",\n            language: params.config.language\n        )",
    )
if "context: config.context" not in text:
    text = text.replace(
        "            language: config.language\n        )",
        "            context: config.context ?? \"\",\n            language: config.language\n        )",
    )
session.write_text(text)
PY

touch "$PATCH_MARKER"
echo "Applied mlx-audio-swift context patch."