#!/usr/bin/env python3
"""Optional patch to the eugr image's vendored deepseek_v4_encoding.py to read
the newer `reasoning_content` message key instead of the deprecated `reasoning`
key.

This is INDEPENDENT of the reasoning_effort patch and is DEFAULT-OFF so the
pure reasoning_effort (minimal) variant is the first-test behavior. Enable at
runtime with VLLM_DSV4_REASONING_CONTENT=1 (recipe [env] / -e). It stays
backward compatible: when enabled it uses `reasoning_content` only when the
message actually has that key, otherwise falls back to `reasoning`.

Three small edits on top of whatever eugr ships:
  1. add a module-level env-gated flag (_USE_REASONING_CONTENT_KEY) + import os
  2. render_message: use reasoning_content (fallback reasoning) when enabled
  3. drop_thinking: also pop reasoning_content (no-op when key unused)

Everything else in eugr's file is left intact, so no independent updates are
lost. Idempotent --check (pristine or already-patched) + apply + post-check,
mirroring patch_deepseek_v4_effort.py.

Usage:  python3 patch_deepseek_v4_reasoning_content.py [--check]
  --check : verify the file is in a consistent (pristine or patched) state.
  (apply) : replace old anchors with new (requires exact old-anchor match).
"""
import os
import sys

import vllm

BASE = os.path.join(vllm.__path__[0], "tokenizers")
F_ENC = os.path.join(BASE, "deepseek_v4_encoding.py")

# 1) import os + env-gated module flag (inserted after the base imports)
OLD_IMPORTS = (
    "from typing import Any, Dict, List, Union, Optional\n"
    "import copy\n"
    "import json"
)
NEW_IMPORTS = (
    "from typing import Any, Dict, List, Union, Optional\n"
    "import copy\n"
    "import json\n"
    "import os\n"
    "\n"
    "# Optional: read the newer `reasoning_content` message key instead of the\n"
    "# deprecated `reasoning` key. Disabled by default (preserves eugr\n"
    "# behavior); enable with VLLM_DSV4_REASONING_CONTENT=1.\n"
    "_USE_REASONING_CONTENT_KEY = bool(\n"
    "    os.environ.get(\"VLLM_DSV4_REASONING_CONTENT\", \"\").strip().lower()\n"
    "    in (\"1\", \"true\", \"yes\", \"on\")\n"
    ")"
)

# 2) render_message: use reasoning_content when enabled (fallback to reasoning)
OLD_READ = "    reasoning = msg.get(\"reasoning\")"
NEW_READ = (
    "    reasoning = msg.get(\"reasoning\")\n"
    "    if _USE_REASONING_CONTENT_KEY and \"reasoning_content\" in msg:\n"
    "        reasoning = msg.get(\"reasoning_content\")"
)

# 3) drop_thinking: also pop reasoning_content (no-op when key unused)
OLD_POP = "            msg.pop(\"reasoning\", None)"
NEW_POP = (
    "            msg.pop(\"reasoning_content\", None)\n"
    "            msg.pop(\"reasoning\", None)"
)

TARGETS = [(F_ENC, [(OLD_IMPORTS, NEW_IMPORTS), (OLD_READ, NEW_READ),
                    (OLD_POP, NEW_POP)])]


def _consistent(src, old, new):
    """True if the file holds exactly one of old|new (pristine or patched).
    Handles overlapping edits where NEW embeds OLD (e.g. an inserted import
    block that starts with the original lines)."""
    co, cn = src.count(old), src.count(new)
    if cn >= 1:
        # patched: new present once, and old does not appear beyond the copy
        # embedded inside new itself.
        return cn == 1 and co <= new.count(old)
    return co == 1


def main() -> None:
    check = "--check" in sys.argv
    path = TARGETS[0][0]
    src = open(path).read()

    if check:
        for old, new in TARGETS[0][1]:
            if not _consistent(src, old, new):
                raise SystemExit(
                    f"CHECK FAIL: {path} — neither pristine-old nor "
                    "already-patched for an anchor"
                )
        print("check OK: deepseek_v4 reasoning_content state is consistent")
        return

    for old, new in TARGETS[0][1]:
        if src.count(old) != 1:
            raise SystemExit(
                f"refusing: {path} old anchor matched {src.count(old)}x "
                "(expected 1); source changed, re-sync patch"
            )
    for old, new in TARGETS[0][1]:
        src = src.replace(old, new)
    for old, new in TARGETS[0][1]:
        if src.count(new) != 1:
            raise SystemExit(
                f"refusing: {path} new anchor present {src.count(new)}x "
                "after apply (internal error)"
            )
    with open(path, "w") as f:
        f.write(src)
    print(f"patched {path}")


if __name__ == "__main__":
    main()