#!/bin/zsh
set -euo pipefail

SCRIPT_DIR="${0:A:h}"
REPO_DIR="${SCRIPT_DIR:h}"
BASE="${OMLX_MLX_LM_DIR:-/Applications/oMLX.app/Contents/Resources/Python/framework-mlx-base/lib/python3.11/site-packages/mlx_lm}"
PY="${OMLX_PYTHON:-/Applications/oMLX.app/Contents/Resources/Python/cpython-3.11/bin/python3.11}"
OMLX_RESOURCES="${OMLX_RESOURCES:-/Applications/oMLX.app/Contents/Resources}"
PREFIX_CACHE="$OMLX_RESOURCES/omlx/cache/prefix_cache.py"
STAMP="$(date +%Y%m%d-%H%M%S)"

if [[ "$(id -u)" -ne 0 ]]; then
  echo "Run this installer with sudo so it can patch the OMLX app bundle."
  exit 1
fi

test -f "$REPO_DIR/solar_open2.py"
test -f "$SCRIPT_DIR/solar_open2_tool_parser.py"
test -f "$BASE/tokenizer_utils.py"
test -f "$PREFIX_CACHE"

mkdir -p "$BASE/models" "$BASE/tool_parsers"

cp "$BASE/tokenizer_utils.py" "$BASE/tokenizer_utils.py.bak-$STAMP"
cp "$PREFIX_CACHE" "$PREFIX_CACHE.bak-$STAMP"
if [[ -f "$BASE/models/solar_open2.py" ]]; then
  cp "$BASE/models/solar_open2.py" "$BASE/models/solar_open2.py.bak-$STAMP"
fi
if [[ -f "$BASE/tool_parsers/solar_open2.py" ]]; then
  cp "$BASE/tool_parsers/solar_open2.py" "$BASE/tool_parsers/solar_open2.py.bak-$STAMP"
fi

install -m 0644 "$REPO_DIR/solar_open2.py" "$BASE/models/solar_open2.py"
install -m 0644 "$SCRIPT_DIR/solar_open2_tool_parser.py" "$BASE/tool_parsers/solar_open2.py"

TOKENIZER_UTILS="$BASE/tokenizer_utils.py" "$PY" - <<'PY'
import os
from pathlib import Path

path = Path(os.environ["TOKENIZER_UTILS"])
text = path.read_text(encoding="utf-8")

think_marker = '        ("<|think:start|>", "<|think:end|>"),\n'
if think_marker not in text:
    anchor = '        ("<longcat_think>", "</longcat_think>"),\n'
    if anchor not in text:
        raise SystemExit("Could not find the thinking-token list in tokenizer_utils.py")
    text = text.replace(anchor, anchor + think_marker, 1)

parser_marker = '        return "solar_open2"\n'
if parser_marker not in text:
    anchor = (
        '    elif "<|tool_list_start|>" in chat_template:\n'
        '        return "pythonic"\n'
    )
    insert = (
        '    elif "<|tool_call:start|>" in chat_template and "<|tool_arg:start|>" in chat_template:\n'
        '        return "solar_open2"\n'
    )
    if anchor not in text:
        raise SystemExit("Could not find the tool-parser inference block in tokenizer_utils.py")
    text = text.replace(anchor, anchor + insert, 1)

solar_apply_marker = 'is_solar_open2 = "<|tool_call:start|>" in template and "<|tool_arg:start|>" in template and "<|think:start|>" in template\n'
if solar_apply_marker not in text:
    old = (
        '    def apply_chat_template(self, *args, tokenize=True, **kwargs):\n'
        '        if "enable_thinking" not in kwargs:\n'
        '            kwargs["enable_thinking"] = self.has_thinking\n'
    )
    new = (
        '    def apply_chat_template(self, *args, tokenize=True, **kwargs):\n'
        '        template = self._chat_template or getattr(self._tokenizer, "chat_template", None) or ""\n'
        '        is_solar_open2 = "<|tool_call:start|>" in template and "<|tool_arg:start|>" in template and "<|think:start|>" in template\n'
        '        if is_solar_open2 and "reasoning_effort" not in kwargs:\n'
        '            enable_thinking = kwargs.pop("enable_thinking", False)\n'
        '            kwargs["reasoning_effort"] = "high" if enable_thinking else "none"\n'
        '        elif "enable_thinking" not in kwargs:\n'
        '            kwargs["enable_thinking"] = self.has_thinking\n'
    )
    if old not in text:
        raise SystemExit("Could not patch TokenizerWrapper.apply_chat_template in tokenizer_utils.py")
    text = text.replace(old, new, 1)

path.write_text(text, encoding="utf-8")
PY

PREFIX_CACHE="$PREFIX_CACHE" "$PY" - <<'PY'
import os
from pathlib import Path

path = Path(os.environ["PREFIX_CACHE"])
text = path.read_text(encoding="utf-8")

marker = "Solar Open2 variable ArraysCache compatibility"
if marker not in text:
    old = """                    cache = marker_handler.deserialize_state(
                        tuple(elements), meta_state
                    )
"""
    new = """                    # Solar Open2 variable ArraysCache compatibility:
                    # preserve all recurrent-state tensors and their token count.
                    if marker_handler.is_variable_length_state():
                        cache = marker_handler.reconstruct_cache(
                            {
                                \"states\": list(elements),
                                \"cache_type\": marker_class,
                            },
                            meta_state,
                            token_count=valid_token_count,
                        )
                    else:
                        cache = marker_handler.deserialize_state(
                            tuple(elements), meta_state
                        )
"""
    if old not in text:
        raise SystemExit("Could not patch variable ArraysCache reconstruction in prefix_cache.py")
    text = text.replace(old, new, 1)

path.write_text(text, encoding="utf-8")
PY

"$PY" -m py_compile \
  "$BASE/models/solar_open2.py" \
  "$BASE/tool_parsers/solar_open2.py" \
  "$BASE/tokenizer_utils.py" \
  "$PREFIX_CACHE"

echo "Installed Solar Open2 OMLX compatibility patch. Backups use stamp $STAMP."