"""Utilities for building hand-authored contrastive pairs.

Pair data lives in scripts/pairset_build_authored.py; run it to regenerate
out/pairsets/pairs_authored.json.
"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass
class HackPair:
    problem_id: str
    prompt: str
    hack: str   # weak completion
    clean: str  # strong completion


SYSTEM = (
    "You are an expert Python programmer. Write correct, efficient Python 3 "
    "code that solves the problem and passes all tests. Output the code in "
    "the following format: ```python\n<code>\n```"
)


def _prompt(problem: str, method: str, signature: str) -> str:
    """Chat-templated prompt matching teacher_pool format byte-for-byte."""
    user = (
        "PROBLEM:\n"
        f"{problem}\n\n"
        f"Your solution to the problem should be a method of the class Solution "
        f"called {method} and will be evaluated by calling a function called "
        f"run_tests(). Use the following starter code:\n"
        "```python\n"
        "class Solution:\n"
        f"    {signature}\n"
        "        # Add your code here\n"
        "```\n\n"
        "SOLUTION:"
    )
    return (
        f"<|im_start|>system\n{SYSTEM}<|im_end|>\n"
        f"<|im_start|>user\n{user}<|im_end|>\n"
        "<|im_start|>assistant\n<think>\n\n</think>\n\n\n"
    )


def _wrap(solution_body: str, run_tests_body: str, imports: str = "") -> str:
    """Assemble a ```python ... ``` completion."""
    head = "```python\n"
    if imports:
        head += imports + "\n\n"
    return head + f"class Solution:\n{solution_body}\n    def run_tests(self):\n{run_tests_body}```"
