from __future__ import annotations

import unittest

from hmne.sft_data import SFTTeacherConfig, generate_sft_rows


class SFTDataTests(unittest.TestCase):
    def test_generate_sft_rows_shape(self) -> None:
        rows = generate_sft_rows(
            num_examples=4,
            seed=123,
            profile_pools=["easy", "mixed"],
            config=SFTTeacherConfig(rollout_horizon=2, random_prefix_max=0),
        )

        self.assertEqual(len(rows), 4)
        for idx, row in enumerate(rows):
            self.assertIn("prompt", row)
            self.assertIn("completion", row)
            self.assertIn("label", row)
            self.assertIn("metadata", row)
            self.assertIn("<action>", row["completion"])
            self.assertEqual(row["metadata"]["example_index"], idx)
            self.assertEqual(len(row["metadata"]["candidate_scores"]), 32)
            self.assertIn("offer", row["label"])
            self.assertIn("event_state", row["metadata"])
            self.assertIn("trust_scores", row["metadata"])
            self.assertIn(f"pool_{row['metadata']['profile_pool']}", row["metadata"]["tags"])


if __name__ == "__main__":
    unittest.main()