import re
import struct
from pathlib import Path

from ..helpers.types import FileExtension, FileFormat, FilePath, MagicBytes

# Known GGML header variants (older formats like GGMF and GGJT)
GGML_MAGIC_VARIANTS = {
    b"GGML",
    b"GGMF",
    b"GGJT",
    b"GGLA",
    b"GGSA",
}


def is_zipfile(path: str) -> bool:
    """Check if file is a ZIP by reading the signature."""
    file_path = Path(path)
    if not file_path.is_file():
        return False
    try:
        signature = read_magic_bytes(path, 4)
        return signature.startswith(b"PK")
    except OSError:
        return False


def read_magic_bytes(path: str, num_bytes: int = 8) -> bytes:
    with Path(path).open("rb") as f:
        return f.read(num_bytes)


def detect_format_from_magic_bytes(magic4: MagicBytes, magic8: MagicBytes, magic16: MagicBytes) -> FileFormat:
    """Detect file format using Python 3.10+ pattern matching on magic bytes."""
    # Use pattern matching for cleaner magic byte detection
    match magic4:
        case b"GGUF":
            return "gguf"
        case magic if magic in GGML_MAGIC_VARIANTS:
            return "ggml"
        case magic if magic.startswith(b"PK"):
            return "zip"
        case b"\x08\x01\x12\x00":  # ONNX protobuf magic
            return "onnx"
        case _:
            pass

    # Check longer magic sequences
    match magic8:
        case b"\x89HDF\r\n\x1a\n":  # HDF5 magic
            return "hdf5"
        case magic if magic.startswith(b"\x93NUMPY"):
            return "numpy"
        case _:
            pass

    # Check pickle magic bytes using pattern matching
    match magic4[:2]:
        case b"\x80\x02" | b"\x80\x03" | b"\x80\x04" | b"\x80\x05":
            return "pickle"
        case _:
            pass

    # Check for JSON-like formats (SafeTensors, etc.)
    match magic4[0:1]:
        case b"{":
            return "safetensors"
        case _:
            pass

    # Check for patterns in first 16 bytes
    if b"onnx" in magic16:
        return "onnx"
    if b'"__metadata__"' in magic16:
        return "safetensors"

    return "unknown"


def detect_file_format_from_magic(path: str) -> str:
    """Detect file format solely from magic bytes."""
    file_path = Path(path)
    if file_path.is_dir():
        if (file_path / "saved_model.pb").exists():
            return "tensorflow_directory"
        return "directory"

    if not file_path.is_file():
        return "unknown"

    try:
        size = file_path.stat().st_size
        if size < 4:
            return "unknown"

        with file_path.open("rb") as f:
            header = f.read(16)

            # Check for TAR format by looking for the "ustar" signature
            if size >= 262:
                f.seek(257)
                if f.read(5).startswith(b"ustar"):
                    return "tar"
            # Reset to read from header for further checks
            f.seek(0)

            magic4 = header[:4]
            magic8 = header[:8]
            magic16 = header[:16]

            # Try the new pattern matching approach first
            format_result = detect_format_from_magic_bytes(magic4, magic8, magic16)
            if format_result != "unknown":
                return format_result

            # Check for XML-based formats (OpenVINO and PMML)
            if magic16.startswith(b"<?xml"):
                # Read first 64 bytes to check for format-specific tags
                f.seek(0)
                xml_header = f.read(64)
                if b"<net" in xml_header:
                    return "openvino"
                if b"<PMML" in xml_header:
                    return "pmml"

            # SafeTensors format check: 8-byte length header + JSON metadata
            if size >= 12:  # Minimum: 8 bytes length + some JSON
                try:
                    json_length = struct.unpack("<Q", magic8)[0]
                    # Sanity check: JSON length should be reasonable
                    if 0 < json_length < size and json_length < 1024 * 1024:  # Max 1MB JSON
                        f.seek(8)
                        json_start = f.read(min(32, json_length))
                        if json_start.startswith(b"{") and b'"' in json_start:
                            return "safetensors"
                except (struct.error, OSError):
                    pass

    except OSError:
        return "unknown"

    # Fallback: check if it starts with JSON (for old safetensors or other JSON formats)
    magic4 = header[:4]
    magic8 = header[:8]
    magic16 = header[:16]

    if magic4[0:1] == b"{" or (size > 8 and b'"__metadata__"' in magic16):
        return "safetensors"

    if magic4 == b"\x08\x01\x12\x00" or b"onnx" in magic16:
        return "onnx"

    return "unknown"


def detect_file_format(path: str) -> str:
    """
    Attempt to identify the format:
    - TensorFlow SavedModel (directory with saved_model.pb)
    - Keras HDF5 (.h5 file with HDF5 magic bytes)
    - PyTorch ZIP (.pt/.pth file that's a ZIP)
    - Pickle (.pkl/.pickle or other files with pickle magic)
    - PyTorch binary (.bin files with various formats)
    - GGUF/GGML files with magic bytes
    - If extension indicates pickle/pt/h5/pb, etc.
    """
    file_path = Path(path)
    if file_path.is_dir():
        # We'll let the caller handle directory logic.
        # But we do a quick guess if there's a 'saved_model.pb'.
        if any(f.name == "saved_model.pb" for f in file_path.iterdir()):
            return "tensorflow_directory"
        return "directory"

    # Single file
    size = file_path.stat().st_size
    if size < 4:
        return "unknown"

    # Read first bytes for format detection using a single file handle
    with file_path.open("rb") as f:
        header = f.read(16)

    magic4 = header[:4]
    magic8 = header[:8]
    magic16 = header[:16]

    # Check first 8 bytes for HDF5 magic
    hdf5_magic = b"\x89HDF\r\n\x1a\n"
    if magic8 == hdf5_magic:
        return "hdf5"

    # Check for GGUF/GGML magic bytes
    if magic4 == b"GGUF":
        return "gguf"
    if magic4 in GGML_MAGIC_VARIANTS:
        return "ggml"

    ext = file_path.suffix.lower()

    # Check ZIP magic first (for .pt/.pth files that are actually zips)
    if magic4.startswith(b"PK"):
        return "zip"

    # Check pickle magic patterns
    pickle_magics = [
        b"\x80\x02",  # Protocol 2
        b"\x80\x03",  # Protocol 3
        b"\x80\x04",  # Protocol 4
        b"\x80\x05",  # Protocol 5
    ]
    if any(magic4.startswith(m) for m in pickle_magics):
        return "pickle"

    # For .bin files, do more sophisticated detection
    if ext == ".bin":
        # IMPORTANT: Check ZIP format first (PyTorch models saved with torch.save())
        if magic4.startswith(b"PK"):
            return "zip"
        # Check if it's a pickle file
        if any(magic4.startswith(m) for m in pickle_magics):
            return "pickle"
        # Check for safetensors format (starts with JSON header)
        if magic4[0:1] == b"{" or (size > 8 and b'"__metadata__"' in magic16):
            return "safetensors"

        # Check for ONNX format (protobuf)
        if magic4 == b"\x08\x01\x12\x00" or b"onnx" in magic16:
            return "onnx"

        # Otherwise, assume raw binary format (PyTorch weights)
        return "pytorch_binary"

    # Extension-based detection for non-.bin files
    # For .pt/.pth/.ckpt files, check if they're ZIP format first
    if ext in (".pt", ".pth", ".ckpt"):
        # These files can be either ZIP or pickle format
        if magic4.startswith(b"PK"):
            return "zip"
        # If not ZIP, assume pickle format
        return "pickle"
    if ext in (".ptl", ".pte"):
        if magic4.startswith(b"PK"):
            return "executorch"
        return "executorch"
    if ext in (".pkl", ".pickle", ".dill"):
        return "pickle"
    if ext == ".h5":
        return "hdf5"
    if ext == ".pb":
        return "protobuf"
    if ext == ".tflite":
        return "tflite"
    if ext == ".safetensors":
        return "safetensors"
    if ext in (".pdmodel", ".pdiparams"):
        return "paddle"
    if ext == ".msgpack":
        return "flax_msgpack"
    if ext == ".onnx":
        return "onnx"
    ggml_exts = {".ggml", ".ggmf", ".ggjt", ".ggla", ".ggsa"}
    if ext in (".gguf", *ggml_exts):
        # Check magic bytes first for accuracy
        if magic4 == b"GGUF":
            return "gguf"
        if magic4 in GGML_MAGIC_VARIANTS:
            return "ggml"
        # Fall back to extension-based detection
        return "gguf" if ext == ".gguf" else "ggml"
    if ext == ".npy":
        return "numpy"
    if ext == ".npz":
        return "zip"
    if ext == ".joblib":
        if magic4.startswith(b"PK"):
            return "zip"
        return "pickle"
    if ext in (
        ".tar",
        ".tar.gz",
        ".tgz",
        ".tar.bz2",
        ".tbz2",
        ".tar.xz",
        ".txz",
    ):
        return "tar"

    return "unknown"


def find_sharded_files(directory: str) -> list[str]:
    """
    Look for sharded model files like:
    pytorch_model-00001-of-00002.bin
    """
    dir_path = Path(directory).resolve()
    return sorted(
        [
            str(fname.resolve())
            for fname in dir_path.iterdir()
            if fname.is_file() and re.match(r"pytorch_model-\d{5}-of-\d{5}\.bin", fname.name)
        ],
    )


EXTENSION_FORMAT_MAP = {
    ".pt": "pickle",
    ".pth": "pickle",
    ".ckpt": "pickle",
    ".pkl": "pickle",
    ".pickle": "pickle",
    ".dill": "pickle",
    ".h5": "hdf5",
    ".hdf5": "hdf5",
    ".keras": "hdf5",
    ".pb": "protobuf",
    ".safetensors": "safetensors",
    ".onnx": "onnx",
    ".bin": "pytorch_binary",
    ".zip": "zip",
    ".gguf": "gguf",
    ".ggml": "ggml",
    ".ggmf": "ggml",
    ".ggjt": "ggml",
    ".ggla": "ggml",
    ".ggsa": "ggml",
    ".ptl": "executorch",
    ".pte": "executorch",
    ".tar": "tar",
    ".tar.gz": "tar",
    ".tgz": "tar",
    ".tar.bz2": "tar",
    ".tbz2": "tar",
    ".tar.xz": "tar",
    ".txz": "tar",
    ".npy": "numpy",
    ".npz": "zip",
    ".joblib": "pickle",  # joblib can be either zip or pickle format
    ".pdmodel": "paddle",
    ".pdiparams": "paddle",
    ".engine": "tensorrt",
    ".plan": "tensorrt",
    ".msgpack": "flax_msgpack",
}


def detect_format_from_extension_pattern_matching(extension: FileExtension) -> FileFormat:
    """Detect format using Python 3.10+ pattern matching for file extensions."""
    # Use pattern matching for more readable extension handling
    match extension.lower():
        # PyTorch/Pickle formats
        case ".pt" | ".pth" | ".ckpt" | ".pkl" | ".pickle" | ".dill":
            return "pickle"
        # HDF5 formats
        case ".h5" | ".hdf5" | ".keras":
            return "hdf5"
        # Archive formats
        case ".zip":
            return "zip"
        case ".tar" | ".tar.gz" | ".tgz":
            return "tar"
        # ML model formats
        case ".onnx":
            return "onnx"
        case ".safetensors":
            return "safetensors"
        case ".bin":
            return "pytorch_binary"
        # GGML/GGUF formats
        case ".gguf":
            return "gguf"
        case ".ggml" | ".ggmf" | ".ggjt" | ".ggla" | ".ggsa":
            return "ggml"
        # ExecutorTorch formats
        case ".ptl" | ".pte":
            return "executorch"
        # Other formats
        case ".pb":
            return "protobuf"
        case ".tflite":
            return "tflite"
        case ".engine":
            return "tensorrt"
        case ".pdmodel":
            return "paddle"
        case ".xml":
            return "openvino"
        case ".pmml":
            return "pmml"
        case ".npy" | ".npz":
            return "numpy"
        case ".msgpack":
            return "flax_msgpack"
        case ".7z":
            return "sevenzip"
        case _:
            return "unknown"


def detect_format_from_extension(path: FilePath) -> FileFormat:
    """Return a format string based solely on the file extension."""
    file_path = Path(path)
    if file_path.is_dir():
        if (file_path / "saved_model.pb").exists():
            return "tensorflow_directory"
        return "directory"

    # Use pattern matching for modern Python 3.10+ approach
    return detect_format_from_extension_pattern_matching(file_path.suffix)


def validate_file_type(path: str) -> bool:
    """Validate that a file's magic bytes match its extension-based format."""
    try:
        header_format = detect_file_format_from_magic(path)
        ext_format = detect_format_from_extension(path)

        # If extension format is unknown, we can't validate - assume valid
        if ext_format == "unknown":
            return True

        # Small files (< 4 bytes) are always valid - can't determine magic bytes reliably
        file_path = Path(path)
        if file_path.is_file() and file_path.stat().st_size < 4:
            return True

        # Handle special cases where different formats are compatible first
        # before doing the unknown header check

        # Pickle files can be stored in various ways
        if ext_format == "pickle" and header_format in {"pickle", "zip"}:
            return True

        # PyTorch binary files are flexible in format
        if ext_format == "pytorch_binary" and header_format in {
            "pytorch_binary",
            "pickle",
            "zip",
            "unknown",  # .bin files can contain arbitrary binary data
        }:
            return True

        # TensorFlow protobuf files (.pb extension)
        if ext_format == "protobuf" and header_format in {"protobuf", "unknown"}:
            return True

        # PMML files are XML-based with <PMML> tag detection
        if ext_format == "pmml" and header_format == "pmml":
            return True

        # ZIP files can have various extensions (.zip, .pt, .pth, .ckpt, .ptl, .pte)
        if header_format == "zip" and ext_format in {"zip", "pickle", "pytorch_binary", "executorch"}:
            return True

        # TAR files must match
        if ext_format == "tar" and header_format == "tar":
            return True

        # ExecuTorch files should be zip archives
        if ext_format == "executorch":
            return header_format == "zip"

        # HDF5 files should always match
        if ext_format == "hdf5":
            return header_format == "hdf5"

        # SafeTensors files should always match
        if ext_format == "safetensors":
            return header_format == "safetensors"

        # GGUF/GGML files should match their format
        if ext_format in {"gguf", "ggml"}:
            return header_format == ext_format

        # ONNX files (Protocol Buffer format - difficult to detect reliably)
        if ext_format == "onnx":
            return header_format in {"onnx", "unknown"}

        # NumPy files (.npy should match, .npz is ZIP by design)
        if ext_format == "numpy":
            # .npz files are ZIP archives containing multiple .npy files
            # This is the standard NumPy compressed format, not spoofing
            # Use case-insensitive suffix check to handle MODEL.NPZ, model.Npz, etc.
            file_path = Path(path)
            if file_path.suffix.lower() == ".npz":
                return header_format in {"zip", "numpy"}
            return header_format == "numpy"

        # Flax msgpack files (less strict validation)
        if ext_format == "flax_msgpack":
            return True  # Hard to validate msgpack format reliably

        # TensorFlow directories are special case
        if ext_format == "tensorflow_directory":
            return header_format == "tensorflow_directory"

        # TensorFlow Lite files
        if ext_format == "tflite":
            return True  # TFLite format can be complex to validate

        if ext_format == "tensorrt":
            return True  # TensorRT engine files have complex binary format

        # If header format is unknown but extension is known, this might be suspicious
        # unless the file is very small or empty (checked after format-specific rules)
        if header_format == "unknown":
            file_path = Path(path)
            return not (file_path.is_file() and file_path.stat().st_size >= 4)  # Small files are acceptable

        # Default: exact match required
        return header_format == ext_format

    except Exception:
        # If validation fails due to error, assume valid to avoid breaking scans
        return True
