"""
PyTorch file format parser.

Handles both ZIP-based and legacy tar-based PyTorch formats,
with recursive archive extraction for nested containers.
"""

import zipfile
import tarfile
import io
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Dict, Optional, Tuple, BinaryIO

from pickle_scanner.formats.pickle_parser import PickleParser, PickleStream


# Maximum recursion depth for nested archives (prevents zip bombs)
MAX_RECURSION_DEPTH = 3

# PyTorch legacy format uses files named 'pickle', 'storages', 'tensors'
# (no .pkl extension)
LEGACY_PICKLE_NAMES = {'pickle', 'storages', 'tensors'}


@dataclass
class TorchPickleEntry:
    """A pickle entry within a PyTorch archive."""
    name: str
    data: bytes
    stream: Optional[PickleStream] = None


@dataclass
class TorchArchive:
    """Parsed PyTorch archive."""
    path: str
    format: str  # 'zip' or 'tar' or 'legacy'
    pickle_entries: List[TorchPickleEntry] = field(default_factory=list)
    tensor_entries: List[str] = field(default_factory=list)
    version: Optional[str] = None
    errors: List[str] = field(default_factory=list)


def _is_pickle_name(name: str) -> bool:
    """Check if a filename looks like a pickle file."""
    basename = name.rsplit('/', 1)[-1] if '/' in name else name
    if basename.endswith('.pkl') or basename.endswith('.pickle'):
        return True
    # PyTorch legacy format: files named exactly 'pickle', 'storages', 'tensors'
    if basename in LEGACY_PICKLE_NAMES:
        return True
    return False


def _is_zip_data(data: bytes) -> bool:
    """Check if bytes look like a ZIP file."""
    return data[:4] in (b'PK\x03\x04', b'PK\x05\x06')


def _is_tar_data(data: bytes) -> bool:
    """Check if bytes look like a tar file (including compressed)."""
    # gzip-compressed tar
    if data[:2] == b'\x1f\x8b':
        return True
    # bz2-compressed tar
    if data[:3] == b'BZh':
        return True
    # xz-compressed tar
    if data[:6] == b'\xfd7zXZ\x00':
        return True
    # raw tar
    if len(data) >= 512 and data[257:262] in (b'ustar', b'ustar '):
        return True
    return False


class TorchParser:
    """Parser for PyTorch model files."""

    def __init__(self, path: Optional[Path] = None, data: Optional[bytes] = None):
        self.path = path
        self.data = data
        self.errors: List[str] = []

    def parse(self) -> TorchArchive:
        """Parse a PyTorch file and extract pickle streams."""
        if self.path:
            with open(self.path, 'rb') as f:
                data = f.read()
        elif self.data:
            data = self.data
        else:
            raise ValueError("No path or data provided")

        archive = TorchArchive(
            path=str(self.path) if self.path else "<bytes>",
            format='unknown'
        )

        # Try ZIP format first (most common for modern PyTorch)
        if _is_zip_data(data):
            self._parse_zip(data, archive, depth=0)
        # Try tar format (legacy, including compressed)
        elif _is_tar_data(data):
            self._parse_tar(data, archive, depth=0)
        # Try as raw pickle
        else:
            self._parse_raw_pickle(data, archive)

        return archive

    def _looks_like_tar(self, data: bytes) -> bool:
        """Check if data might be a tar archive."""
        return _is_tar_data(data)

    def _parse_zip(self, data: bytes, archive: TorchArchive,
                   depth: int = 0) -> None:
        """Parse a ZIP-based PyTorch file with recursive extraction."""
        if archive.format == 'unknown':
            archive.format = 'zip'

        if depth >= MAX_RECURSION_DEPTH:
            archive.errors.append(
                f"Max recursion depth ({MAX_RECURSION_DEPTH}) reached "
                "parsing nested ZIP"
            )
            return

        try:
            with zipfile.ZipFile(io.BytesIO(data)) as zf:
                names = zf.namelist()

                # Read version if present
                version_files = [n for n in names if 'version' in n.lower()]
                if version_files:
                    try:
                        archive.version = zf.read(
                            version_files[0]
                        ).decode('utf-8').strip()
                    except Exception:
                        pass

                for name in names:
                    try:
                        member_data = zf.read(name)
                    except Exception as e:
                        archive.errors.append(f"Error reading {name}: {e}")
                        continue

                    # Check for pickle files
                    if _is_pickle_name(name):
                        try:
                            parser = PickleParser(member_data)
                            stream = parser.parse()
                            entry = TorchPickleEntry(name, member_data, stream)
                            archive.pickle_entries.append(entry)
                        except Exception as e:
                            archive.errors.append(
                                f"Error parsing {name}: {e}"
                            )

                    # Recursively handle nested ZIP files
                    elif _is_zip_data(member_data):
                        self._parse_zip(
                            member_data, archive, depth=depth + 1
                        )

                    # Recursively handle nested tar files
                    elif _is_tar_data(member_data):
                        self._parse_tar(
                            member_data, archive, depth=depth + 1
                        )

                    # Track tensor storage files
                    elif '/data/' in name or name.endswith('.storage'):
                        archive.tensor_entries.append(name)

        except zipfile.BadZipFile as e:
            archive.errors.append(f"Invalid ZIP file: {e}")
            archive.format = 'invalid_zip'
            # Fail-closed: attempt raw pickle scan on the full data
            self._parse_raw_pickle(data, archive)

    def _parse_tar(self, data: bytes, archive: TorchArchive,
                   depth: int = 0) -> None:
        """Parse a tar-based PyTorch file with recursive extraction.

        Handles plain tar and compressed tar (gz, bz2, xz) via 'r:*' mode.
        """
        if archive.format == 'unknown':
            archive.format = 'tar'

        if depth >= MAX_RECURSION_DEPTH:
            archive.errors.append(
                f"Max recursion depth ({MAX_RECURSION_DEPTH}) reached "
                "parsing nested tar"
            )
            return

        try:
            with tarfile.open(fileobj=io.BytesIO(data), mode='r:*') as tf:
                for member in tf.getmembers():
                    if not member.isfile():
                        continue

                    try:
                        f = tf.extractfile(member)
                        if not f:
                            continue
                        member_data = f.read()
                    except Exception as e:
                        archive.errors.append(
                            f"Error reading {member.name}: {e}"
                        )
                        continue

                    # Check for pickle files (including legacy names)
                    if _is_pickle_name(member.name):
                        try:
                            parser = PickleParser(member_data)
                            stream = parser.parse()
                            entry = TorchPickleEntry(
                                member.name, member_data, stream
                            )
                            archive.pickle_entries.append(entry)
                        except Exception as e:
                            archive.errors.append(
                                f"Error parsing {member.name}: {e}"
                            )

                    # Recursively handle nested ZIP files inside tar
                    elif _is_zip_data(member_data):
                        self._parse_zip(
                            member_data, archive, depth=depth + 1
                        )

                    # Recursively handle nested tar files inside tar
                    elif _is_tar_data(member_data):
                        self._parse_tar(
                            member_data, archive, depth=depth + 1
                        )

                    elif ('tensor' in member.name.lower() or
                          member.name.endswith('.data')):
                        archive.tensor_entries.append(member.name)

        except tarfile.TarError as e:
            archive.errors.append(f"Invalid tar file: {e}")
            archive.format = 'invalid_tar'
            # Fail-closed: attempt raw pickle scan on the full data
            self._parse_raw_pickle(data, archive)

    def _parse_raw_pickle(self, data: bytes, archive: TorchArchive) -> None:
        """Parse as raw pickle data."""
        if archive.format == 'unknown':
            archive.format = 'raw_pickle'

        try:
            parser = PickleParser(data)
            stream = parser.parse()
            entry = TorchPickleEntry('root.pkl', data, stream)
            archive.pickle_entries.append(entry)
        except Exception as e:
            archive.errors.append(f"Error parsing as raw pickle: {e}")


def parse_torch_file(path: Path) -> TorchArchive:
    """Convenience function to parse a PyTorch file."""
    parser = TorchParser(path=path)
    return parser.parse()


def parse_torch_bytes(data: bytes) -> TorchArchive:
    """Convenience function to parse PyTorch data from bytes."""
    parser = TorchParser(data=data)
    return parser.parse()
