"""ReVeR reference pipeline: write a rubric for an instruction, judge each response against it, score.

    rubric generator  instruction -> rubric; every item ends with [Hard Rule] or [Principle]
    judge             (instruction, one response, rubric) -> 0/1 per Hard Rule, 0-100 per Principle;
                      it may call a Python tool for mechanical checks
    score             number of satisfied Hard Rules + mean Principle score / 101

Both models are reached through OpenAI-compatible endpoints (see README.md for the vLLM commands).

    python rever.py --rubric-url http://localhost:8001/v1 --judge-url http://localhost:8002/v1 \
        --instruction "Write exactly two sentences about tea." \
        --response "Tea is a drink. It is made from leaves." --response "Tea is nice."
"""
from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Any

from openai import OpenAI

from python_tool import RestrictedPythonExecutor

HERE = Path(__file__).resolve().parent
RUBRIC_PROMPT = (HERE / "prompts" / "rubric_generator_prompt.txt").read_text(encoding="utf-8")
JUDGE_PROMPT = (HERE / "prompts" / "judge_system_prompt.txt").read_text(encoding="utf-8")

PYTHON_TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "python",
            "description": "Execute Python code for deterministic verification.",
            "parameters": {
                "type": "object",
                "properties": {"code": {"type": "string"}},
                "required": ["code"],
            },
        },
    }
]

# Decoding used for both models in all our evaluations. Thinking stays on.
SAMPLING = {
    "temperature": 1.0,
    "top_p": 0.95,
    "presence_penalty": 1.5,
    "max_tokens": 8192,
    "extra_body": {
        "top_k": 20,
        "min_p": 0.0,
        "repetition_penalty": 1.0,
        "chat_template_kwargs": {"enable_thinking": True},
    },
}
MAX_TOOL_CALLS = 3      # executed tool calls per judgment
MAX_JUDGE_TURNS = 4     # model calls per judgment; the last one is made without tools
TOOL_TIMEOUT_SECONDS = 3.0
TOOL_OUTPUT_LIMIT = 4000
TOOL_MEMORY_LIMIT_MB = 512

THINK_RE = re.compile(r"<think\b[^>]*>.*?</think\s*>", re.IGNORECASE | re.DOTALL)
RUBRIC_OPEN_RE = re.compile(r"<rubric\s*>", re.IGNORECASE)
RUBRIC_CLOSE_RE = re.compile(r"</rubric\s*>", re.IGNORECASE)
RUBRIC_LABEL_RE = re.compile(r"\[(Hard Rule|Principle)\]\s*$")


# --------------------------------------------------------------------------------------------- rubric

def build_rubric_prompt(instruction: str) -> str:
    return RUBRIC_PROMPT.format(instruction=instruction, context="")


def parse_rubric(text: str) -> tuple[list[dict[str, Any]], str | None]:
    """Items of the last <rubric> block -> ([{rubric_id, criterion, kind}], error)."""
    text = THINK_RE.sub("", text or "")
    blocks = []
    for opening in RUBRIC_OPEN_RE.finditer(text):
        closing = RUBRIC_CLOSE_RE.search(text, opening.end())
        if closing is not None:
            blocks.append(text[opening.end():closing.start()])
        elif RUBRIC_LABEL_RE.search(text[opening.end():]):      # block cut off right after its last item
            blocks.append(text[opening.end():])
    if not blocks:
        return [], "missing <rubric>...</rubric> block"
    block = blocks[-1].strip().replace("```", "")
    matches = list(re.finditer(r"(?m)^\s*(\d+)\s*[.)]\s+", block))
    if not matches:
        return [], "rubric block contains no numbered items"
    items = []
    for position, match in enumerate(matches, start=1):
        if int(match.group(1)) != position:
            return [], f"rubric item numbering must be consecutive from 1; got {match.group(1)}"
        end = matches[position].start() if position < len(matches) else len(block)
        criterion = re.sub(r"\s+", " ", block[match.end():end]).strip()
        label = RUBRIC_LABEL_RE.search(criterion)
        if label is None:
            return [], f"rubric item {position} must end with [Hard Rule] or [Principle]"
        kind = "hard_rule" if label.group(1) == "Hard Rule" else "principle"
        items.append({"rubric_id": position, "criterion": criterion, "kind": kind})
    return items, None


def generate_rubric(client: OpenAI, model: str, instruction: str, max_attempts: int = 3) -> list[dict[str, Any]]:
    """One rubric for the instruction. The prompt is a single user message; there is no system message."""
    error = None
    for _ in range(max_attempts):
        response = client.chat.completions.create(
            model=model, messages=[{"role": "user", "content": build_rubric_prompt(instruction)}], **SAMPLING
        )
        items, error = parse_rubric(response.choices[0].message.content or "")
        if error is None:
            return items
    raise RuntimeError(f"no parseable rubric after {max_attempts} attempts: {error}")


# ---------------------------------------------------------------------------------------------- judge

def _reasoning(message: Any) -> str | None:
    """The thinking of a turn: vLLM >= 0.11 returns it as `reasoning`, older servers as `reasoning_content`."""
    for key in ("reasoning_content", "reasoning"):
        value = getattr(message, key, None)
        if isinstance(value, str) and value:
            return value
    return None


def _assistant_turn(message: Any, call: Any = None) -> dict[str, Any]:
    """The assistant turn to send back. Its thinking goes back under both keys, so that the chat template renders
    the <think> block again; without it the judge answers after a tool call without its own earlier reasoning."""
    turn: dict[str, Any] = {"role": "assistant", "content": message.content or ""}
    reasoning = _reasoning(message)
    if reasoning:
        turn["reasoning_content"] = reasoning
        turn["reasoning"] = reasoning
    if call is not None:
        turn["tool_calls"] = [
            {"id": call.id, "type": "function", "function": {"name": call.function.name, "arguments": call.function.arguments}}
        ]
    return turn


def _run_tool(call: Any, context: dict[str, Any]) -> str:
    try:
        arguments = json.loads(call.function.arguments or "{}")
    except json.JSONDecodeError:
        arguments = {}
    code = arguments.get("code") if isinstance(arguments, dict) else None
    if call.function.name != "python":
        return json.dumps({"ok": False, "stdout": "", "result": None, "error": "unsupported tool"})
    if not isinstance(code, str):
        return json.dumps({"ok": False, "stdout": "", "result": None, "error": "missing string argument: code"})
    executor = RestrictedPythonExecutor(
        timeout_seconds=TOOL_TIMEOUT_SECONDS, output_limit=TOOL_OUTPUT_LIMIT, memory_limit_mb=TOOL_MEMORY_LIMIT_MB
    )
    return executor.run(code, context)


def _balanced_object(text: str) -> str | None:
    start = text.find("{")
    if start < 0:
        return None
    depth, in_string, escaped = 0, False, False
    for index, char in enumerate(text[start:], start=start):
        if in_string:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == '"':
                in_string = False
            continue
        if char == '"':
            in_string = True
        elif char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                return text[start:index + 1]
    return None


def parse_judgments(text: str, items: list[dict[str, Any]]) -> tuple[list[dict[str, Any]] | None, str | None]:
    """The judgment object of the final answer -> ([{rubric_id, explanation, score}], error)."""
    by_id = {item["rubric_id"]: item for item in items}
    cleaned = THINK_RE.sub("", text or "")
    blocks = re.findall(r"<judge\s*>(.*?)</judge\s*>", cleaned, flags=re.IGNORECASE | re.DOTALL)
    body = re.sub(r"```(?:json)?|```", "", blocks[-1] if blocks else cleaned, flags=re.IGNORECASE).strip()
    candidates = [body]
    balanced = _balanced_object(body)
    if balanced and balanced not in candidates:
        candidates.append(balanced)
    candidates += [c.replace("\\'", "'") for c in candidates if "\\'" in c]   # \' is not a JSON escape
    parsed, error = None, "no judge JSON object found"
    for candidate in candidates:
        try:
            parsed = json.loads(candidate)
            break
        except json.JSONDecodeError as exc:
            error = str(exc)
    if parsed is None:
        return None, error
    if not isinstance(parsed, dict) or set(parsed) != {"judgments"} or not isinstance(parsed["judgments"], list):
        return None, "the judge JSON must be an object with the single key 'judgments' holding a list"
    judgments = parsed["judgments"]
    if len(judgments) != len(items):
        return None, f"expected {len(items)} judgments, got {len(judgments)}"
    if all(isinstance(j, dict) and j.get("rubric_id") in (None, "") for j in judgments):
        judgments = [dict(j, rubric_id=item["rubric_id"]) for j, item in zip(judgments, items)]
    out, seen = [], set()
    for judgment in judgments:
        if not isinstance(judgment, dict):
            return None, "each judgment must be an object"
        rubric_id, score = judgment.get("rubric_id"), judgment.get("score")
        if isinstance(rubric_id, str) and re.fullmatch(r"\d+", rubric_id.strip()):
            rubric_id = int(rubric_id.strip())
        if isinstance(score, str) and re.fullmatch(r"\d+", score.strip()):
            score = int(score.strip())
        if type(rubric_id) is not int or rubric_id not in by_id or rubric_id in seen:
            return None, f"unknown or repeated rubric_id: {judgment.get('rubric_id')!r}"
        if not isinstance(judgment.get("explanation"), str):
            return None, f"explanation for rubric {rubric_id} must be a string"
        if type(score) is not int:
            return None, f"score for rubric {rubric_id} must be an integer"
        if by_id[rubric_id]["kind"] == "hard_rule" and score not in (0, 1):
            return None, f"Hard Rule rubric {rubric_id} score must be 0 or 1"
        if by_id[rubric_id]["kind"] != "hard_rule" and not 0 <= score <= 100:
            return None, f"Principle rubric {rubric_id} score must be between 0 and 100"
        seen.add(rubric_id)
        out.append({"rubric_id": rubric_id, "explanation": judgment["explanation"], "score": score})
    return sorted(out, key=lambda j: j["rubric_id"]), None


def judge_once(client: OpenAI, model: str, instruction: str, response: str, items: list[dict[str, Any]]) -> str:
    """One judgment episode; returns the text of the final answer."""
    rubrics = [{"rubric_id": item["rubric_id"], "criterion": item["criterion"]} for item in items]
    payload = {"instruction": instruction, "response": response, "rubrics": rubrics}
    conversation: list[dict[str, Any]] = [
        {"role": "system", "content": JUDGE_PROMPT},
        {"role": "user", "content": json.dumps(payload, ensure_ascii=False, indent=2)},
    ]
    # variables the judge's code can read
    context = {
        "conversation": [{"role": "user", "content": instruction}, {"role": "assistant", "content": response}],
        "instruction": instruction,
        "response": response,
        "rubrics": rubrics,
        "rubric_items": items,
        "last_assistant_response": response,
    }
    executed = 0
    for turn in range(1, MAX_JUDGE_TURNS + 1):
        offer_tools = executed < MAX_TOOL_CALLS and turn < MAX_JUDGE_TURNS
        extra = {"tools": PYTHON_TOOLS, "tool_choice": "auto"} if offer_tools else {}
        message = client.chat.completions.create(model=model, messages=conversation, **SAMPLING, **extra).choices[0].message
        calls = (message.tool_calls or []) if offer_tools else []
        if not calls:
            return message.content if (message.content or "").strip() else (_reasoning(message) or "")
        call = calls[0]                                   # one call is executed per turn
        conversation.append(_assistant_turn(message, call))
        conversation.append({"role": "tool", "tool_call_id": call.id, "content": _run_tool(call, context)})
        executed += 1
    raise RuntimeError("the judge loop ended without a final answer")


def score_of(items: list[dict[str, Any]], judgments: list[dict[str, Any]]) -> float:
    """Number of satisfied Hard Rules + mean Principle score / 101. The second term is below 1, so Principles only
    order responses that satisfy the same number of Hard Rules."""
    by_id = {j["rubric_id"]: j["score"] for j in judgments}
    hard = sum(by_id[i["rubric_id"]] for i in items if i["kind"] == "hard_rule")
    principles = [by_id[i["rubric_id"]] for i in items if i["kind"] != "hard_rule"]
    return hard + (sum(principles) / len(principles) if principles else 0.0) / 101


def judge(client: OpenAI, model: str, instruction: str, response: str, items: list[dict[str, Any]],
          max_attempts: int = 3) -> dict[str, Any]:
    """Judge one response against the rubric -> {score, judgments}. The judge never sees another response."""
    error = None
    for _ in range(max_attempts):
        judgments, error = parse_judgments(judge_once(client, model, instruction, response, items), items)
        if error is None:
            return {"score": score_of(items, judgments), "judgments": judgments}
    raise RuntimeError(f"no parseable judgment after {max_attempts} attempts: {error}")


# ------------------------------------------------------------------------------------------------ CLI

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--rubric-url", required=True, help="base URL of the rubric generator endpoint")
    parser.add_argument("--judge-url", required=True, help="base URL of the judge endpoint")
    parser.add_argument("--rubric-model", default="rever-rubric", help="served model name of the rubric generator")
    parser.add_argument("--judge-model", default="rever-judge", help="served model name of the judge")
    parser.add_argument("--instruction", required=True)
    parser.add_argument("--response", action="append", required=True, help="repeat for several responses")
    parser.add_argument("--api-key", default="EMPTY")
    args = parser.parse_args()
    rubric_client = OpenAI(base_url=args.rubric_url, api_key=args.api_key, timeout=600)
    judge_client = OpenAI(base_url=args.judge_url, api_key=args.api_key, timeout=600)

    items = generate_rubric(rubric_client, args.rubric_model, args.instruction)
    print("Rubric")
    for item in items:
        print(f"  {item['rubric_id']}. {item['criterion']}")
    for number, response in enumerate(args.response, start=1):
        result = judge(judge_client, args.judge_model, args.instruction, response, items)
        print(f"\nResponse {number}: score {result['score']:.4f}")
        for judgment in result["judgments"]:
            print(f"  {judgment['rubric_id']}. score {judgment['score']:>3}  {judgment['explanation']}")


if __name__ == "__main__":
    main()