import os
import subprocess
import vllm

# Get vllm reasoning directory
vllm_path = os.path.dirname(vllm.__file__)
reasoning_dir = os.path.join(vllm_path, "reasoning")
init_file = os.path.join(reasoning_dir, "__init__.py")
parser_file = os.path.join(reasoning_dir, "greenmind_14b_r1_reasoning_parser.py")

# Download the GreenMind parser
url = "https://huggingface.co/GreenNode/GreenMind-Medium-14B-R1/resolve/main/greenmind_14b_r1_reasoning_parser.py"
subprocess.run(["wget", "-O", parser_file, url], check=True)

# Read __init__.py
with open(init_file, "r") as f:
    lines = f.readlines()

# Ensure import is present
import_line = "from .greenmind_14b_r1_reasoning_parser import GreenMind14bR1ReasoningParser\n"
if import_line not in lines:
    last_import_idx = None
    for i, line in enumerate(lines):
        if line.startswith("from ."):
            last_import_idx = i
    if last_import_idx is not None:
        lines.insert(last_import_idx + 1, import_line)

# Ensure __all__ contains the parser
all_entry = '    "GreenMind14bR1ReasoningParser",\n'
if all_entry not in lines:
    for i, line in enumerate(lines):
        if line.strip().startswith("__all__"):
            # find the next line with a string entry inside __all__
            for j in range(i + 1, len(lines)):
                if lines[j].strip().startswith('"') or lines[j].strip().startswith("'"):
                    lines.insert(j, all_entry)
                    break
            break

# Write back updated file
with open(init_file, "w") as f:
    f.writelines(lines)

print("✅ GreenMind14bR1ReasoningParser successfully added to vLLM.")