"""Tests for format detection and parsing."""

import tempfile
import os
import io
import gzip
import bz2
import lzma
import zipfile
import tarfile
from pathlib import Path

import pytest

from pickle_scanner.formats.detector import detect_format, detect_format_from_bytes, FileFormat
from pickle_scanner.formats.torch_parser import TorchParser, _is_pickle_name
from pickle_scanner.formats.joblib_parser import JoblibParser, detect_joblib_compression
from pickle_scanner import PickleScanner, RiskLevel


class TestFormatDetection:
    """Test format detection functionality."""

    def test_detect_raw_pickle(self):
        """Test detection of raw pickle format."""
        # Protocol 4 pickle header
        pickle_data = b"\x80\x04\x95\x00\x00\x00\x00\x00\x00\x00."

        with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f:
            f.write(pickle_data)
            temp_path = f.name

        try:
            result = detect_format(Path(temp_path))
            assert result.format == FileFormat.PICKLE
        finally:
            os.unlink(temp_path)

    def test_detect_pytorch_zip(self):
        """Test detection of PyTorch ZIP format."""
        # Create minimal PyTorch-like ZIP
        with tempfile.NamedTemporaryFile(suffix=".pt", delete=False) as f:
            temp_path = f.name
            with zipfile.ZipFile(f, 'w') as zf:
                zf.writestr("archive/data.pkl", b"\x80\x04\x95\x00\x00\x00\x00\x00\x00\x00.")

        try:
            result = detect_format(Path(temp_path))
            assert result.format == FileFormat.PYTORCH_ZIP
        finally:
            os.unlink(temp_path)

    def test_detect_safetensors(self):
        """Test detection of SafeTensors format (marked as safe)."""
        # SafeTensors header starts with 8-byte little-endian length
        safetensors_data = b"\x08\x00\x00\x00\x00\x00\x00\x00{\"test\":}"

        with tempfile.NamedTemporaryFile(suffix=".safetensors", delete=False) as f:
            f.write(safetensors_data)
            temp_path = f.name

        try:
            result = detect_format(Path(temp_path))
            assert result.format == FileFormat.SAFETENSORS
        finally:
            os.unlink(temp_path)


class TestJoblibDetection:
    """Test Joblib format detection and parsing."""

    def _make_pickle_bytes(self):
        """Create minimal pickle bytes (protocol 4, empty dict)."""
        return b"\x80\x04\x95\x05\x00\x00\x00\x00\x00\x00\x00}\x94."

    def test_detect_joblib_gzip(self):
        """Test detection of gzip-compressed joblib."""
        pkl = self._make_pickle_bytes()
        compressed = gzip.compress(pkl)
        assert detect_joblib_compression(compressed) == 'gzip'

    def test_detect_joblib_bz2(self):
        """Test detection of bz2-compressed joblib."""
        pkl = self._make_pickle_bytes()
        compressed = bz2.compress(pkl)
        assert detect_joblib_compression(compressed) == 'bz2'

    def test_detect_joblib_lzma(self):
        """Test detection of xz-compressed joblib."""
        pkl = self._make_pickle_bytes()
        compressed = lzma.compress(pkl, format=lzma.FORMAT_XZ)
        assert detect_joblib_compression(compressed) == 'xz'

    def test_detect_joblib_zlib(self):
        """Test detection of zlib-compressed joblib."""
        pkl = self._make_pickle_bytes()
        import zlib
        compressed = zlib.compress(pkl)
        assert detect_joblib_compression(compressed) == 'zlib'

    def test_detect_joblib_raw(self):
        """Test detection of uncompressed joblib (raw pickle)."""
        pkl = self._make_pickle_bytes()
        assert detect_joblib_compression(pkl) == 'none'

    def test_joblib_format_detection_by_extension(self):
        """Joblib file detected by extension."""
        pkl = self._make_pickle_bytes()
        result = detect_format_from_bytes(pkl, suffix=".joblib")
        assert result.format == FileFormat.JOBLIB

    def test_joblib_gzip_scanned(self):
        """Gzip-compressed joblib with malicious pickle is detected."""
        malicious_pickle = (
            b"cos\nsystem\n"
            b"(S'echo pwned'\n"
            b"tR."
        )
        compressed = gzip.compress(malicious_pickle)

        scanner = PickleScanner()
        # Detect as joblib by extension hint
        result = detect_format_from_bytes(compressed, suffix=".joblib")
        assert result.format == FileFormat.JOBLIB

        # Actually scan it
        parser = JoblibParser(compressed)
        jresult = parser.parse()
        assert jresult.compression == 'gzip'
        assert jresult.pickle_stream is not None

    def test_joblib_bz2_malicious_detected(self):
        """BZ2-compressed joblib with os.system should be flagged CRITICAL."""
        malicious_pickle = (
            b"cos\nsystem\n"
            b"(S'echo pwned'\n"
            b"tR."
        )
        compressed = bz2.compress(malicious_pickle)

        with tempfile.NamedTemporaryFile(suffix=".joblib", delete=False) as f:
            f.write(compressed)
            temp_path = f.name

        try:
            scanner = PickleScanner()
            result = scanner.scan_file(temp_path)
            assert result.report.risk_level == RiskLevel.CRITICAL
        finally:
            os.unlink(temp_path)


class TestKerasDetection:
    """Test Keras .keras format detection and scanning."""

    def test_detect_keras_format(self):
        """Keras .keras file (ZIP with weights.npz) detected."""
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, 'w') as zf:
            zf.writestr("config.json", b'{"class_name": "Sequential"}')
            # Create a minimal npz inside
            npz_buf = io.BytesIO()
            with zipfile.ZipFile(npz_buf, 'w') as npz_zf:
                npz_zf.writestr("arr_0.npy", b"\x93NUMPY\x01\x00v\x00{'descr': '<f4'}")
            zf.writestr("model.weights.npz", npz_buf.getvalue())
        data = buf.getvalue()

        result = detect_format_from_bytes(data, suffix=".keras")
        assert result.format == FileFormat.KERAS

    def test_keras_scanned(self):
        """Keras file scanning completes without errors."""
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, 'w') as zf:
            zf.writestr("config.json", b'{"class_name": "Sequential"}')
            npz_buf = io.BytesIO()
            with zipfile.ZipFile(npz_buf, 'w') as npz_zf:
                npz_zf.writestr("arr_0.npy", b"\x93NUMPY\x01\x00v\x00{'descr': '<f4'}")
            zf.writestr("model.weights.npz", npz_buf.getvalue())

        scanner = PickleScanner()
        result = scanner.scan_bytes(buf.getvalue(), "model.keras")
        # Should complete without crash
        assert result.report is not None
        assert result.report.format == "keras"


class TestNestedArchives:
    """Test recursive archive extraction."""

    def test_zip_in_zip_detected(self):
        """Pickle inside a ZIP inside a ZIP should be found."""
        malicious_pickle = (
            b"cos\nsystem\n"
            b"(S'echo pwned'\n"
            b"tR."
        )

        # Inner ZIP with pickle
        inner_buf = io.BytesIO()
        with zipfile.ZipFile(inner_buf, 'w') as inner_zf:
            inner_zf.writestr("model/data.pkl", malicious_pickle)

        # Outer ZIP containing inner ZIP
        outer_buf = io.BytesIO()
        with zipfile.ZipFile(outer_buf, 'w') as outer_zf:
            outer_zf.writestr("archive/data.pkl", b"\x80\x04\x95\x05\x00\x00\x00\x00\x00\x00\x00}\x94.")
            outer_zf.writestr("archive/inner.zip", inner_buf.getvalue())

        parser = TorchParser(data=outer_buf.getvalue())
        archive = parser.parse()
        # Should find pickle entries from both levels
        assert len(archive.pickle_entries) >= 2

    def test_tar_in_zip_detected(self):
        """Pickle inside a tar inside a ZIP should be found."""
        malicious_pickle = (
            b"cos\nsystem\n"
            b"(S'echo pwned'\n"
            b"tR."
        )

        # Create tar with pickle inside
        tar_buf = io.BytesIO()
        with tarfile.open(fileobj=tar_buf, mode='w') as tf:
            info = tarfile.TarInfo(name="model.pkl")
            info.size = len(malicious_pickle)
            tf.addfile(info, io.BytesIO(malicious_pickle))

        # Outer ZIP containing tar
        outer_buf = io.BytesIO()
        with zipfile.ZipFile(outer_buf, 'w') as outer_zf:
            outer_zf.writestr("data.pkl", b"\x80\x04\x95\x05\x00\x00\x00\x00\x00\x00\x00}\x94.")
            outer_zf.writestr("nested.tar", tar_buf.getvalue())

        parser = TorchParser(data=outer_buf.getvalue())
        archive = parser.parse()
        assert len(archive.pickle_entries) >= 2

    def test_compressed_tar_detected(self):
        """Pickle inside a gzip-compressed tar should be found."""
        malicious_pickle = (
            b"cos\nsystem\n"
            b"(S'echo pwned'\n"
            b"tR."
        )

        # Create gzip-compressed tar
        tar_buf = io.BytesIO()
        with tarfile.open(fileobj=tar_buf, mode='w:gz') as tf:
            info = tarfile.TarInfo(name="model.pkl")
            info.size = len(malicious_pickle)
            tf.addfile(info, io.BytesIO(malicious_pickle))

        parser = TorchParser(data=tar_buf.getvalue())
        archive = parser.parse()
        assert len(archive.pickle_entries) >= 1
        # Verify the malicious content was found
        found_pkl = False
        for entry in archive.pickle_entries:
            if entry.data == malicious_pickle:
                found_pkl = True
        assert found_pkl

    def test_recursion_depth_limit(self):
        """Deeply nested archives should hit the recursion limit."""
        pickle_data = b"\x80\x04\x95\x05\x00\x00\x00\x00\x00\x00\x00}\x94."

        # Build 5-level nested ZIPs (exceeds MAX_RECURSION_DEPTH=3)
        data = pickle_data
        for i in range(5):
            buf = io.BytesIO()
            with zipfile.ZipFile(buf, 'w') as zf:
                if i == 0:
                    zf.writestr("data.pkl", data)
                else:
                    zf.writestr("inner.zip", data)
            data = buf.getvalue()

        parser = TorchParser(data=data)
        archive = parser.parse()
        # Should have a recursion depth error
        assert any("recursion depth" in e.lower() or "Max recursion" in e
                    for e in archive.errors)


class TestLegacyPyTorchFormat:
    """Test PyTorch legacy tar format with non-standard filenames."""

    def test_legacy_pickle_file_detected(self):
        """Files named 'pickle' (no extension) in tar should be scanned."""
        pickle_data = (
            b"cos\nsystem\n"
            b"(S'echo test'\n"
            b"tR."
        )

        tar_buf = io.BytesIO()
        with tarfile.open(fileobj=tar_buf, mode='w') as tf:
            # Legacy PyTorch uses file named 'pickle' without extension
            info = tarfile.TarInfo(name="pickle")
            info.size = len(pickle_data)
            tf.addfile(info, io.BytesIO(pickle_data))

            info2 = tarfile.TarInfo(name="storages")
            info2.size = 4
            tf.addfile(info2, io.BytesIO(b"data"))

            info3 = tarfile.TarInfo(name="tensors")
            info3.size = 4
            tf.addfile(info3, io.BytesIO(b"data"))

        parser = TorchParser(data=tar_buf.getvalue())
        archive = parser.parse()
        # Should find 'pickle', 'storages', 'tensors' as pickle entries
        assert len(archive.pickle_entries) >= 1
        found_names = {e.name for e in archive.pickle_entries}
        assert 'pickle' in found_names

    def test_is_pickle_name(self):
        """Test the pickle name detection helper."""
        assert _is_pickle_name("data.pkl")
        assert _is_pickle_name("model.pickle")
        assert _is_pickle_name("archive/data.pkl")
        assert _is_pickle_name("pickle")
        assert _is_pickle_name("storages")
        assert _is_pickle_name("tensors")
        assert not _is_pickle_name("weights.bin")
        assert not _is_pickle_name("config.json")


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
