#!/usr/bin/env python3
"""Minimal patch to the eugr image's vendored deepseek_v4 tokenizer to add the
updated reasoning_effort (low/high/max) support WITHOUT rewriting the whole file.

We deliberately do NOT replace the files with the full 0731 model-card variant.
Instead this applies the smallest possible edits on top of whatever eugr ships:

  deepseek_v4.py (API layer):
    - let "low"/"high"/"max" pass through; unknown/garbage -> None (reverts to
      engine default) instead of coercing everything to "high".

  deepseek_v4_encoding.py:
    - replace the single REASONING_EFFORT_MAX constant with the per-level
      REASONING_EFFORT_PROMPTS dict + DEFAULT_REASONING_EFFORT="low"
      ("high" text == the old "max" text; a stronger "max" added; "low" = "").
    - in render_message: default None->"low", assert against the dict, and emit
      the per-level prompt at index 0 in thinking mode.

Everything else in eugr's files (tool handling, `reasoning` key, structure) is
left intact, so no independent updates are lost.

Usage:  python3 patch_deepseek_v4_effort.py [--check]
  --check : verify each file is in a CONSISTENT state (either pristine-old or
            already-patched) and exit 0 if so. Both before and after apply.
  (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")
F_API = os.path.join(BASE, "deepseek_v4.py")


# --- deepseek_v4.py: let low/high/max through; unknown -> None ---
OLD_API = (
    "            elif reasoning_effort in (\"max\", \"xhigh\"):\n"
    "                reasoning_effort = \"max\"\n"
    "            else:\n"
    "                reasoning_effort = \"high\""
)
NEW_API = (
    "            elif reasoning_effort in (\"max\", \"xhigh\"):\n"
    "                reasoning_effort = \"max\"\n"
    "            elif reasoning_effort not in (\"low\", \"high\", \"max\"):\n"
    "                reasoning_effort = None"
)

# --- deepseek_v4_encoding.py: constant -> dict + default ---
OLD_EFF_CONST = (
    "REASONING_EFFORT_MAX = (\n"
    "    \"Reasoning Effort: Absolute maximum with no shortcuts permitted.\\n\"\n"
    '    "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\\n"\n'
    '    "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\\n\\n"\n'
    ")"
)
NEW_EFF_CONST = (
    "REASONING_EFFORT_PROMPTS = {\n"
    "    \"low\": \"\",\n"
    "    \"high\": (\n"
    "        \"Reasoning Effort: Absolute maximum with no shortcuts permitted.\\n\"\n"
    "        \"You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\\n\"\n"
    "        \"Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\\n\\n\"\n"
    "    ),\n"
    "    \"max\": (\n"
    "        \"Reasoning Effort: Beyond maximum \\u2014 exhaustive, relentless, and uncompromising.\\n\"\n"
    "        \"You MUST reason with the utmost depth and rigor, leaving absolutely nothing to chance: exhaustively decompose the problem into its most fundamental components, trace every causal chain to its root, and resolve the underlying cause rather than any surface symptom.\\n\"\n"
    "        \"Do not stop reasoning until you have independently verified the solution from multiple angles and are certain that no assumption remains unchecked and no error remains undiscovered.\\n\\n\"\n"
    "    ),\n"
    "}\n"
    'DEFAULT_REASONING_EFFORT = "low"'
)

# --- deepseek_v4_encoding.py: render_message effort hunk ---
OLD_EFF_HUNK = (
    "    # Reasoning effort prefix (only at index 0 in thinking mode with max effort)\n"
    "    assert reasoning_effort in ['max', None, 'high'], f\"Invalid reasoning effort: {reasoning_effort}\"\n"
    "    if index == 0 and thinking_mode == \"thinking\" and reasoning_effort == 'max':\n"
    "        prompt += REASONING_EFFORT_MAX"
)
NEW_EFF_HUNK = (
    "    # Reasoning effort prefix (only at index 0 in thinking mode; \"low\" adds nothing)\n"
    "    reasoning_effort = reasoning_effort or DEFAULT_REASONING_EFFORT\n"
    "    assert reasoning_effort in REASONING_EFFORT_PROMPTS, \\\n"
    "        f\"Invalid reasoning effort: {reasoning_effort}, expected one of {list(REASONING_EFFORT_PROMPTS)}\"\n"
    "    if index == 0 and thinking_mode == \"thinking\":\n"
    "        prompt += REASONING_EFFORT_PROMPTS[reasoning_effort]"
)

# Each entry: (path, [(old_anchor, new_anchor), ...])
TARGETS = [
    (F_API, [(OLD_API, NEW_API)]),
    (F_ENC, [(OLD_EFF_CONST, NEW_EFF_CONST), (OLD_EFF_HUNK, NEW_EFF_HUNK)]),
]


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
    states = [(path, open(path).read()) for path, _ in TARGETS]

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

    # Apply: require exact old anchors, then replace and verify new.
    for path, src in states:
        for old, new in dict(TARGETS)[path]:
            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 dict(TARGETS)[path]:
            src = src.replace(old, new)
        for old, new in dict(TARGETS)[path]:
            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()