import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

import os

from nzfc_gram_runtime import NZFCGramLongMemoryChat
from nzfc_gram_runtime.quality import attach_answer_quality_governor
from nzfc_gram_runtime.large_document import attach_large_document_memory
from nzfc_gram_runtime.diffusiongemma_adapter import attach_diffusiongemma_block_diffusion

MODEL_ID = 'google/diffusiongemma-26B-A4B-it'

if os.environ.get('LOAD_MODEL', '0') != '1':
    print('Set LOAD_MODEL=1 to run this optional model-load check.')
    raise SystemExit(0)

bot = NZFCGramLongMemoryChat(
    repo_dir=str(ROOT),
    model_id=MODEL_ID,
    memory_db_path='./optional_model_load_memory.sqlite3',
    load_model=False,
    require_model=False,
    preload_static_memory=False,
)

attach_large_document_memory(bot)
attach_answer_quality_governor(bot)

meta = attach_diffusiongemma_block_diffusion(
    bot,
    model_id=MODEL_ID,
    device_map='auto',
    dtype='auto',
)

print(meta)

out = bot.generate_answer(
    system_prompt='You are concise.',
    user_prompt='Say DIFFUSIONGEMMA_OK.',
    max_new_tokens=32,
)

print(out)