from forgeenv.roles.drift_generator import BaselineDriftGenerator
from forgeenv.roles.prompts import render_drift_generator_prompt
from forgeenv.tasks.task_sampler import TaskSampler

sampler = TaskSampler()
script = sampler.get_by_id("simple_regression").script_content

prompt = render_drift_generator_prompt(script, "ChangeTokenizerBehavior", {"transformers": "4.40"})
fence = "```python"
script_block = ""
if fence in prompt:
    script_block = prompt.split(fence, 1)[1].split("```", 1)[0]
print("script_block len:", len(script_block))
print("first 80 chars:", repr(script_block[:80]))

gen = BaselineDriftGenerator(seed=0)
spec = gen.propose(target_category="ChangeTokenizerBehavior", script=script_block)
print("spec:", spec)