import json
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory

import torch

from speech_bridge_gemma.slice_longcat_codes import slice_codes


class SliceLongCatCodesTest(unittest.TestCase):
    def test_slice_uses_manifest_subset_index(self) -> None:
        with TemporaryDirectory() as tmp:
            root = Path(tmp)
            source = root / "source.pt"
            manifest = root / "manifest.jsonl"
            out = root / "out.pt"
            torch.save(
                {
                    "meta": {"codec": "longcat", "n_acoustic_codebooks": 1, "input_sample_rate": 16000},
                    "semantic_codes": [
                        torch.tensor([0, 0]),
                        torch.tensor([1, 1, 1]),
                        torch.tensor([2]),
                    ],
                    "acoustic_codes": [
                        torch.tensor([[10, 10]]),
                        torch.tensor([[11, 11, 11]]),
                        torch.tensor([[12]]),
                    ],
                },
                source,
            )
            rows = [
                {"id": "row_b", "text": "texto b", "prompt": "pergunta b", "answer": "resposta b", "subset_index": 2, "duration": 1.0},
                {"id": "row_a", "text": "texto a", "question": "pergunta a", "subset_index": 0, "duration_sec": 2.0},
            ]
            manifest.write_text("\n".join(json.dumps(row) for row in rows) + "\n", encoding="utf-8")

            summary = slice_codes(source, manifest, out, "subset_index")
            payload = torch.load(out, map_location="cpu", weights_only=False)

            self.assertEqual(summary["rows"], 2)
            self.assertEqual(payload["semantic_codes"][0].tolist(), [2])
            self.assertEqual(payload["semantic_codes"][1].tolist(), [0, 0])
            self.assertEqual(payload["acoustic_codes"][0].tolist(), [[12]])
            self.assertEqual(payload["meta"]["samples"][0]["id"], "row_b")
            self.assertEqual(payload["meta"]["samples"][0]["frames"], 1)
            self.assertEqual(payload["meta"]["samples"][1]["answer_text"], "texto a")


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