"""Tests for BM25 pickle checkpoint utilities."""
import os
import pickle
import tempfile
import unittest

from app.utils.bm25_checkpoint import (
    atomic_save_pkl,
    backup_pkl,
    is_pkl_corrupted,
    list_backup_paths,
    load_pkl_dict,
    prune_all_backups,
    validate_pkl_data,
)


def _valid_payload(n: int = 2):
    return {
        "chunk_ids": [f"c{i}" for i in range(n)],
        "documents": [f"doc {i}" for i in range(n)],
        "metadatas": [{"file_id": f"f{i}"} for i in range(n)],
        "tokenized_corpus": [["doc", str(i)] for i in range(n)],
    }


class TestBm25Checkpoint(unittest.TestCase):
    def test_missing_not_corrupt(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "missing.pkl")
            corrupt, reason = is_pkl_corrupted(path)
            self.assertFalse(corrupt)
            self.assertEqual(reason, "")

    def test_zero_byte_corrupt(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "empty.pkl")
            open(path, "wb").close()
            corrupt, reason = is_pkl_corrupted(path)
            self.assertTrue(corrupt)
            self.assertIn("zero", reason.lower())

    def test_length_mismatch_corrupt(self):
        data = _valid_payload(2)
        data["documents"] = data["documents"][:1]
        ok, msg = validate_pkl_data(data)
        self.assertFalse(ok)
        self.assertIn("mismatch", msg)

    def test_valid_roundtrip(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "ok.pkl")
            payload = _valid_payload(3)
            atomic_save_pkl(path, payload)
            corrupt, _ = is_pkl_corrupted(path)
            self.assertFalse(corrupt)
            loaded = load_pkl_dict(path)
            self.assertEqual(len(loaded["chunk_ids"]), 3)

    def test_atomic_save_leaves_main_on_bad_temp(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "main.pkl")
            atomic_save_pkl(path, _valid_payload(1))
            size_before = os.path.getsize(path)
            # Simulate failed write by corrupting temp manually is hard; ensure replace works
            atomic_save_pkl(path, _valid_payload(2))
            self.assertGreater(os.path.getsize(path), 0)
            loaded = load_pkl_dict(path)
            self.assertEqual(len(loaded["chunk_ids"]), 2)

    def test_backup_rotation(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "idx_bm25.pkl")
            os.environ["BM25_BACKUP_KEEP"] = "2"
            try:
                from app.core.config import settings

                object.__setattr__(settings, "BM25_BACKUP_KEEP", 2)
                for i in range(4):
                    atomic_save_pkl(path, _valid_payload(i + 1))
                    backup_pkl(path)
                backups = list_backup_paths(path)
                self.assertLessEqual(len(backups), 2)
            finally:
                prune_all_backups(path)

    def test_truncated_pickle_corrupt(self):
        with tempfile.TemporaryDirectory() as tmp:
            path = os.path.join(tmp, "trunc.pkl")
            with open(path, "wb") as f:
                f.write(b"\x80\x04\x95")  # incomplete pickle header
            corrupt, reason = is_pkl_corrupted(path)
            self.assertTrue(corrupt)


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