from transformers import AutoTokenizer

REPO_ID = "Resad2173/BPE_Magibu"
CHAT_TEMPLATE_PATH = "chat_template.jinja"

NEW_SPECIAL_TOKENS = [
    "<|yilmaz_start|>",
    "<|yilmaz_end|>",
    "<|yilmaz_think|>",
    "<|yilmaz_think_end|>",
    "<|yilmaz_tool_call|>",
    "<|yilmaz_tool_call_end|>",
    "<|yilmaz_tool_result|>",
    "<|yilmaz_tool_result_end|>",
]

print(f"Tokenizer indiriliyor: {REPO_ID}")
tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
print(f"Mevcut vocab boyutu: {len(tokenizer)}")

num_added = tokenizer.add_special_tokens(
    {"additional_special_tokens": NEW_SPECIAL_TOKENS}
)
print(f"{num_added} yeni özel token eklendi.")
print(f"Yeni vocab boyutu: {len(tokenizer)}")

for token in NEW_SPECIAL_TOKENS:
    token_id = tokenizer.convert_tokens_to_ids(token)
    assert token_id != tokenizer.unk_token_id, f"{token} eklenemedi, hâlâ <unk>."
    print(f"  {token!r:32} -> id {token_id}")

with open(CHAT_TEMPLATE_PATH, "r", encoding="utf-8") as f:
    tokenizer.chat_template = f.read()
print("\nchat_template tokenizer'a bağlandı.")

# --- Gerçek tokenizer ile doğrulama (push'tan önce) ---
sample_messages = [
    {"role": "system", "content": "Sen yardımsever bir bankacılık asistanısın."},
    {"role": "user", "content": "Hesap bakiyemi nasıl görebilirim?"},
]

prompt_text = tokenizer.apply_chat_template(
    sample_messages, tokenize=False, add_generation_prompt=True
)
print("\n--- Üretilen metin (basit sohbet) ---")
print(prompt_text)

result = tokenizer.apply_chat_template(
    sample_messages, tokenize=True, add_generation_prompt=True
)
token_ids = result["input_ids"] if hasattr(result, "keys") and "input_ids" in result else result
tokens = tokenizer.convert_ids_to_tokens(token_ids)
print("\n--- Token'lara bölünmüş hâli ---")
print(tokens)

for always_present in ["<|yilmaz_start|>", "<|yilmaz_end|>"]:
    assert always_present in tokens, f"{always_present} tek token olarak bulunamadı."
print("\nstart/end token'ları tek parça olarak doğrulandı.")

# --- Thinking + tool-calling'i AYRI, o özellikleri gerçekten kullanan bir örnekle test edelim ---
tool_messages = [
    {"role": "user", "content": "Ankara'da hava nasıl?"},
    {
        "role": "assistant",
        "content": "",
        "reasoning_content": "Kullanıcı hava durumu soruyor, get_weather çağırmalıyım.",
        "tool_calls": [
            {"function": {"name": "get_weather", "arguments": {"city": "Ankara"}}}
        ],
    },
    {"role": "tool", "content": '{"city": "Ankara", "temp_c": 27}'},
]

tool_prompt_text = tokenizer.apply_chat_template(
    tool_messages, tokenize=False, add_generation_prompt=True, enable_thinking=True
)
print("\n--- Üretilen metin (thinking + tool-calling) ---")
print(tool_prompt_text)

tool_result = tokenizer.apply_chat_template(
    tool_messages, tokenize=True, add_generation_prompt=True, enable_thinking=True
)
tool_token_ids = (
    tool_result["input_ids"]
    if hasattr(tool_result, "keys") and "input_ids" in tool_result
    else tool_result
)
tool_tokens = tokenizer.convert_ids_to_tokens(tool_token_ids)

for special in NEW_SPECIAL_TOKENS:
    assert special in tool_tokens, f"{special} tek token olarak bulunamadı."
print("\nTüm 8 özel token (bu örnekte) tek parça olarak doğrulandı.")


local_dir = "BPE_Magibu_with_yilmaz_template"
tokenizer.save_pretrained(local_dir)
print(f"\nGüncellenmiş tokenizer yerel olarak kaydedildi: {local_dir}")

push = input(f"\n'{REPO_ID}' reposuna push edilsin mi? (e/h): ").strip().lower()
if push == "e":
    tokenizer.push_to_hub(REPO_ID)
    print(f"Push edildi: https://huggingface.co/{REPO_ID}")
else:
    print("Push atlandı.")