"""
Joblib file format parser.

Joblib files can be:
- Raw pickle (just a .joblib with pickle bytes)
- Compressed pickle with various backends: zlib, gzip, bz2, lzma, xz, lz4
"""

import io
import gzip
import bz2
import lzma
import zlib
from dataclasses import dataclass, field
from typing import List, Optional

from pickle_scanner.formats.pickle_parser import PickleParser, PickleStream


# Compression magic bytes for detection
ZLIB_MAGIC = b'\x78'        # zlib (0x78 0x01/9c/da/5e)
GZIP_MAGIC = b'\x1f\x8b'   # gzip
BZ2_MAGIC = b'BZ'           # bz2
LZMA_MAGIC = b'\xfd7zXZ\x00'  # xz
LZMA_ALONE_MAGIC = b'\x5d\x00\x00'  # lzma alone format
# lz4 magic: 0x04 0x22 0x4d 0x18
LZ4_MAGIC = b'\x04\x22\x4d\x18'


@dataclass
class JoblibResult:
    """Result of parsing a joblib file."""
    compression: str  # 'none', 'zlib', 'gzip', 'bz2', 'lzma', 'xz', 'lz4'
    pickle_stream: Optional[PickleStream] = None
    pickle_data: Optional[bytes] = None
    errors: List[str] = field(default_factory=list)


def detect_joblib_compression(data: bytes) -> str:
    """Detect the compression type of a joblib file."""
    if data[:2] == GZIP_MAGIC:
        return 'gzip'
    if data[:2] == BZ2_MAGIC:
        return 'bz2'
    if data[:6] == LZMA_MAGIC:
        return 'xz'
    if data[:3] == LZMA_ALONE_MAGIC:
        return 'lzma'
    if data[:4] == LZ4_MAGIC:
        return 'lz4'
    if len(data) > 0 and data[0:1] == ZLIB_MAGIC:
        # zlib header: first byte is 0x78, second varies
        if len(data) > 1 and data[1] in (0x01, 0x5e, 0x9c, 0xda):
            return 'zlib'
    return 'none'


class JoblibParser:
    """Parser for Joblib model files."""

    def __init__(self, data: bytes):
        self.data = data

    def parse(self) -> JoblibResult:
        """Parse a joblib file, decompress if needed, and extract pickle."""
        compression = detect_joblib_compression(self.data)
        result = JoblibResult(compression=compression)

        try:
            raw_pickle = self._decompress(compression)
            result.pickle_data = raw_pickle
            parser = PickleParser(raw_pickle)
            result.pickle_stream = parser.parse()
        except Exception as e:
            result.errors.append(f"Joblib parse error ({compression}): {e}")

        return result

    def _decompress(self, compression: str) -> bytes:
        """Decompress joblib data based on detected compression."""
        if compression == 'none':
            return self.data
        elif compression == 'gzip':
            return gzip.decompress(self.data)
        elif compression == 'bz2':
            return bz2.decompress(self.data)
        elif compression == 'zlib':
            return zlib.decompress(self.data)
        elif compression in ('lzma', 'xz'):
            return lzma.decompress(self.data)
        elif compression == 'lz4':
            try:
                import lz4.frame
                return lz4.frame.decompress(self.data)
            except ImportError:
                raise RuntimeError(
                    "lz4 package required to decompress lz4 joblib files"
                )
        else:
            return self.data
