#!/bin/bash
set -euo pipefail

CONDA_BIN=${CONDA_BIN:-}
if [ -z "$CONDA_BIN" ]; then
  if command -v conda >/dev/null 2>&1; then
    CONDA_BIN=$(command -v conda)
  elif [ -x /opt/miniconda/bin/conda ]; then
    CONDA_BIN=/opt/miniconda/bin/conda
  fi
fi

if [ -n "$CONDA_BIN" ]; then
  eval "$("$CONDA_BIN" shell.bash hook)"
  conda activate focallora4
else
  echo "[Ident_verb_test_dataset.sh] Warning: conda not found; running in current environment." >&2
fi

SCRIPT_DIR="$(cd -- "$(dirname "$0")" && pwd)"
export PYTHONPATH="$SCRIPT_DIR${PYTHONPATH:+:$PYTHONPATH}"
export IGNORE_REASONING_MESSAGES="${IGNORE_REASONING_MESSAGES:-1}"

TRAJ_PATH="${TRAJ_PATH:-/data/local/hujk/BUTTON/crafted_data/attack_dh_traj.jsonl}"
TOKENIZER_PATH="${TOKENIZER_PATH:-/data/local/hujk/models/Qwen3-8B}"

python3 - <<'PY'
import json
import os
from pathlib import Path
from transformers import AutoTokenizer
from lib_tokenize_data_mask import (
    apply_chat_with_tokenize_with_mark,
    apply_chat_with_tokenize_original,
    filter_reasoning_messages,
    is_ignore_reasoning_enabled,
    strip_markers,
)

traj_path = Path(os.getenv("TRAJ_PATH", "/data/local/hujk/BUTTON/crafted_data/attack_dh_traj.jsonl"))
tokenizer_path = os.getenv("TOKENIZER_PATH", "/data/local/hujk/models/Qwen3-8B")

first_line = traj_path.read_text().splitlines()[0]
record = json.loads(first_line)
messages = record["trajectory"]
tools = record.get("tools")

ignore_flag = is_ignore_reasoning_enabled()
filtered = filter_reasoning_messages(messages, ignore_flag)
# Sanitize contents to strings for deterministic rendering.
sanitized = []
for m in filtered:
    m = dict(m)
    if m.get("content") is None:
        m["content"] = ""
    elif not isinstance(m.get("content"), str):
        m["content"] = json.dumps(m["content"])
    sanitized.append(m)

print(f"IGNORE_REASONING_MESSAGES={ignore_flag}")
print(f"Messages: original={len(messages)} filtered={len(filtered)}")
print(
    "Reasoning messages removed:",
    len([m for m in messages if "reasoning_content" in m]) - len(
        [m for m in filtered if "reasoning_content" in m]
    ),
)

tok = AutoTokenizer.from_pretrained(tokenizer_path)
if tok.pad_token is None:
    tok.pad_token = tok.eos_token

input_ids, instr_mask, data_mask, segment_type, is_normal_token, _custom_mask, rendered = apply_chat_with_tokenize_with_mark(
    sanitized, tok, tools=tools
)
orig_ids, orig_render = apply_chat_with_tokenize_original(sanitized, tok, tools=tools)

print(f"Token count: {len(input_ids)}")
print(f"Instruction tokens: {sum(instr_mask)}")
print(f"Data tokens: {sum(data_mask)}")
print(f"Segment labels present: {sorted(set(segment_type))}")
print(f"input_ids match original: {input_ids == orig_ids}")
print("\nRendered preview with markers (first 300 chars):")
print(rendered[:300])
cleaned = strip_markers(rendered)
print("\nRendered preview (markers stripped, first 300 chars):")
print(cleaned[:300])

# Detailed token/mask dump
tokens = [tok.decode(t).replace("\n","\\n") for t in input_ids]
print("\nidx\tinstr\tdata\tseg\tnorm\ttoken")
for i, (t, im, dm, seg, norm) in enumerate(zip(tokens, instr_mask, data_mask, segment_type, is_normal_token)):
    print(f"{i}\t{int(im)}\t{int(dm)}\t{seg}\t{int(norm)}\t{t}")

print("\n--- Original apply_chat_template + tokenizer ---")
orig_tokens = tok.convert_ids_to_tokens(orig_ids)
for i, t in enumerate(orig_tokens[:50]):
    human = tok.convert_tokens_to_string([t]) or t
    print(f"{i}\t{human}")
PY