# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch",
#     "transformers[torch]",
#     "huggingface_hub",
#     "bitsandbytes",
#     "accelerate",
# ]
# ///
"""Print FULL replies + extracted blocks for a few code prompts.

The eval report only stores reply_first_160, which hides why v2 scored 0% ast_rate
while the base model scored 93.3%. This dumps everything so we can tell whether the
model emits broken code or the extractor mis-parses a well-formed reply.
"""
import json
import os
import re
import sys

import torch
from transformers import (AutoModelForCausalLM, AutoTokenizer,
                          BitsAndBytesConfig)

REPO = sys.argv[1] if len(sys.argv) > 1 else "Taimwe/securecoder-30b-pro-v2-merged"
MAXNEW = 1024

PROMPTS = [
    "Write a Python function `def is_palindrome(s: str) -> bool:` that checks whether a string reads the same forwards and backwards, ignoring case and non-alphanumeric characters.",
    "Write a Python function `def two_sum(nums: list[int], target: int) -> list[int]:` returning the indices of the two numbers adding to target, or an empty list.",
    "Write a Python function `def flatten(nested: list) -> list:` that flattens an arbitrarily nested list.",
]

token = os.environ.get("HF_TOKEN")
tok = AutoTokenizer.from_pretrained(REPO, token=token)
model = AutoModelForCausalLM.from_pretrained(
    REPO, token=token,
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True, bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_use_double_quant=True),
    dtype=torch.bfloat16, device_map="auto")

out = []
for i, p in enumerate(PROMPTS, 1):
    text = tok.apply_chat_template([{"role": "user", "content": p}],
                                   tokenize=False, add_generation_prompt=True)
    ids = tok(text, return_tensors="pt", add_special_tokens=False).input_ids.to(model.device)
    with torch.no_grad():
        gen = model.generate(ids, max_new_tokens=MAXNEW, do_sample=False,
                             eos_token_id=tok.eos_token_id, pad_token_id=tok.eos_token_id)
    new_ids = gen[0, ids.shape[-1]:]
    reply = tok.decode(new_ids, skip_special_tokens=True)

    # exactly the eval harness's extraction
    blocks = re.findall(r"```(?:python)?\s*\n(.*?)```", reply, re.S)
    block = max(blocks, key=len) if blocks else None
    if block is None:
        s = reply.find("def ")
        block = reply[s:] if s >= 0 else None
    if block:
        block = re.split(r"\n```", block, maxsplit=1)[0]
    err = None
    if block:
        try:
            compile(block, "<t>", "exec")
        except Exception as exc:
            err = "%s: %s" % (type(exc).__name__, exc)
    out.append({"prompt": p[:80], "n_new_tokens": int(len(new_ids)),
                "n_fenced_blocks": len(blocks),
                "block_len": len(block) if block else 0,
                "compile_error": err, "full_reply": reply})
    print("\n" + "=" * 74)
    print("PROMPT %d  new_tokens=%d  fenced_blocks=%d  block_len=%s"
          % (i, len(new_ids), len(blocks), len(block) if block else 0))
    print("COMPILE:", err or "OK")
    print("-" * 74)
    print(reply[:2500])

dest = sys.argv[2] if len(sys.argv) > 2 else "/data/dbg.json"
with open(dest, "w", encoding="utf-8") as fh:
    json.dump(out, fh, indent=2, default=str)
print("\nWROTE", dest)