"""Tests for TorchServe .mar and NeMo .nemo format support."""

import io
import gzip
import zipfile
import tarfile

import pytest

from pickle_scanner import PickleScanner, RiskLevel
from pickle_scanner.formats.detector import detect_format_from_bytes, FileFormat


class TestTorchServeMAR:
    """Test TorchServe .mar format detection and scanning."""

    def test_detect_mar_format(self):
        """A ZIP file with .mar extension should be detected as TORCHSERVE_MAR."""
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, 'w') as zf:
            zf.writestr("MAR-INF/MANIFEST.json", b'{"model": "test"}')
            zf.writestr("model.pkl", b"\x80\x04\x95\x05\x00\x00\x00\x00\x00\x00\x00}\x94.")

        result = detect_format_from_bytes(buf.getvalue(), suffix=".mar")
        assert result.format == FileFormat.TORCHSERVE_MAR

    def test_mar_with_malicious_pickle(self):
        """MAR containing a malicious pickle should be flagged CRITICAL."""
        malicious_pickle = (
            b"cos\nsystem\n"
            b"(S'echo pwned'\n"
            b"tR."
        )

        buf = io.BytesIO()
        with zipfile.ZipFile(buf, 'w') as zf:
            zf.writestr("MAR-INF/MANIFEST.json", b'{"model": "evil"}')
            zf.writestr("model.pkl", malicious_pickle)

        scanner = PickleScanner()
        result = scanner.scan_bytes(buf.getvalue(), "evil_model.mar")
        assert result.report.risk_level == RiskLevel.CRITICAL
        assert result.report.format == "torchserve_mar"

    def test_mar_with_nested_pytorch_zip(self):
        """MAR containing a nested PyTorch .pt file should be recursively scanned."""
        malicious_pickle = (
            b"cos\nsystem\n"
            b"(S'echo nested'\n"
            b"tR."
        )

        # Create inner PyTorch-like ZIP
        inner_buf = io.BytesIO()
        with zipfile.ZipFile(inner_buf, 'w') as inner_zf:
            inner_zf.writestr("archive/data.pkl", malicious_pickle)

        # Create outer MAR ZIP
        outer_buf = io.BytesIO()
        with zipfile.ZipFile(outer_buf, 'w') as outer_zf:
            outer_zf.writestr("MAR-INF/MANIFEST.json", b'{}')
            outer_zf.writestr("model.pt", inner_buf.getvalue())

        scanner = PickleScanner()
        result = scanner.scan_bytes(outer_buf.getvalue(), "nested.mar")
        assert result.report.risk_level == RiskLevel.CRITICAL

    def test_mar_safe_model(self):
        """MAR with safe pickle should not flag as dangerous."""
        safe_pickle = b"\x80\x04\x95\x05\x00\x00\x00\x00\x00\x00\x00}\x94."

        buf = io.BytesIO()
        with zipfile.ZipFile(buf, 'w') as zf:
            zf.writestr("MAR-INF/MANIFEST.json", b'{}')
            zf.writestr("model.pkl", safe_pickle)

        scanner = PickleScanner()
        result = scanner.scan_bytes(buf.getvalue(), "safe.mar")
        assert result.report.risk_level == RiskLevel.SAFE
        assert result.report.format == "torchserve_mar"


class TestNeMo:
    """Test NeMo .nemo format detection and scanning."""

    def test_detect_nemo_format(self):
        """.nemo extension should be detected as NEMO format."""
        # NeMo files are tars, create a minimal one
        tar_buf = io.BytesIO()
        with tarfile.open(fileobj=tar_buf, mode='w:gz') as tf:
            data = b"model_config: {}"
            info = tarfile.TarInfo(name="model_config.yaml")
            info.size = len(data)
            tf.addfile(info, io.BytesIO(data))

        result = detect_format_from_bytes(tar_buf.getvalue(), suffix=".nemo")
        assert result.format == FileFormat.NEMO

    def test_nemo_with_malicious_pickle(self):
        """NeMo tar with malicious pickle should be flagged CRITICAL."""
        malicious_pickle = (
            b"cos\nsystem\n"
            b"(S'echo pwned'\n"
            b"tR."
        )

        tar_buf = io.BytesIO()
        with tarfile.open(fileobj=tar_buf, mode='w:gz') as tf:
            info = tarfile.TarInfo(name="model_weights.pkl")
            info.size = len(malicious_pickle)
            tf.addfile(info, io.BytesIO(malicious_pickle))

        scanner = PickleScanner()
        result = scanner.scan_bytes(tar_buf.getvalue(), "evil.nemo")
        assert result.report.risk_level == RiskLevel.CRITICAL
        assert result.report.format == "nemo"

    def test_nemo_with_nested_pytorch(self):
        """NeMo tar with nested PyTorch .pt file should be scanned."""
        malicious_pickle = (
            b"cos\nsystem\n"
            b"(S'echo nested'\n"
            b"tR."
        )

        # Create inner PyTorch ZIP
        inner_buf = io.BytesIO()
        with zipfile.ZipFile(inner_buf, 'w') as inner_zf:
            inner_zf.writestr("archive/data.pkl", malicious_pickle)
        inner_data = inner_buf.getvalue()

        # Create NeMo tar
        tar_buf = io.BytesIO()
        with tarfile.open(fileobj=tar_buf, mode='w:gz') as tf:
            info = tarfile.TarInfo(name="model.pt")
            info.size = len(inner_data)
            tf.addfile(info, io.BytesIO(inner_data))

        scanner = PickleScanner()
        result = scanner.scan_bytes(tar_buf.getvalue(), "nested.nemo")
        assert result.report.risk_level == RiskLevel.CRITICAL

    def test_nemo_with_joblib_inside(self):
        """NeMo tar with joblib-compressed pickle should be scanned."""
        malicious_pickle = (
            b"cos\nsystem\n"
            b"(S'echo joblib'\n"
            b"tR."
        )
        compressed = gzip.compress(malicious_pickle)

        tar_buf = io.BytesIO()
        with tarfile.open(fileobj=tar_buf, mode='w:gz') as tf:
            info = tarfile.TarInfo(name="weights.joblib")
            info.size = len(compressed)
            tf.addfile(info, io.BytesIO(compressed))

        scanner = PickleScanner()
        result = scanner.scan_bytes(tar_buf.getvalue(), "joblib_inside.nemo")
        assert result.report.risk_level == RiskLevel.CRITICAL

    def test_nemo_safe_model(self):
        """NeMo tar with safe content should not flag."""
        safe_pickle = b"\x80\x04\x95\x05\x00\x00\x00\x00\x00\x00\x00}\x94."

        tar_buf = io.BytesIO()
        with tarfile.open(fileobj=tar_buf, mode='w:gz') as tf:
            info = tarfile.TarInfo(name="model.pkl")
            info.size = len(safe_pickle)
            tf.addfile(info, io.BytesIO(safe_pickle))

        scanner = PickleScanner()
        result = scanner.scan_bytes(tar_buf.getvalue(), "safe.nemo")
        assert result.report.risk_level == RiskLevel.SAFE
        assert result.report.format == "nemo"


class TestLZ4OptionalDep:
    """Test lz4 optional dependency handling."""

    def test_lz4_import_error_handled(self):
        """When lz4 is not installed, joblib parser should report a clear error."""
        from pickle_scanner.formats.joblib_parser import JoblibParser, LZ4_MAGIC

        # Fake lz4-compressed data (lz4 magic + garbage)
        fake_lz4 = LZ4_MAGIC + b'\x00' * 100

        parser = JoblibParser(fake_lz4)
        result = parser.parse()
        # Should either succeed (if lz4 is installed) or have a clear error
        if result.errors:
            assert any("lz4" in e.lower() for e in result.errors)


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