"""Message format inferrer class and convenience functions.

Provides the ``MessageFormatInferrer`` class that orchestrates boundary
detection, field classification, and schema generation.  Also exposes
the three module-level convenience functions (``infer_format``,
``detect_field_types``, ``find_dependencies``) used by downstream code.

References:
    IPART: IP Packet Analysis using Random Forests. IEEE ISSRE 2014.
    Discoverer: Automatic Protocol Reverse Engineering. USENIX Security 2007.
"""

from __future__ import annotations

from typing import Literal

import numpy as np
from numpy.typing import NDArray

from oscura.inference.message_format.advanced_detectors import (
    detect_enum_field,
    detect_float_field,
    detect_length_field,
    detect_reserved_field,
    detect_timestamp_field,
)
from oscura.inference.message_format.boundaries import (
    detect_boundaries_voting,
    detect_field_boundaries,
    expert_alignment,
    expert_distribution,
    expert_entropy,
    expert_ngrams,
    expert_variance,
)
from oscura.inference.message_format.entropy import (
    calculate_byte_entropy,
    calculate_entropy,
)
from oscura.inference.message_format.field_classification import (
    classify_field,
    classify_field_ensemble,
    detect_checksum_field,
    detect_counter_field,
    detect_type_entropy,
    detect_type_patterns,
    detect_type_statistics,
    extract_field_data,
    vote_field_type,
)
from oscura.inference.message_format.types import (
    FieldDetectionResult,
    InferredField,
    MessageSchema,
)


class MessageFormatInferrer:
    """Infer message format from samples.

    : Message format inference using entropy and variance analysis.

    Algorithm:
    1. Detect field boundaries using entropy transitions
    2. Classify fields based on statistical patterns
    3. Detect dependencies between fields
    4. Generate complete schema
    """

    def __init__(self, min_samples: int = 10) -> None:
        """Initialize inferrer.

        Args:
            min_samples: Minimum number of message samples required
        """
        self.min_samples = min_samples

    # -----------------------------------------------------------------
    # Public API
    # -----------------------------------------------------------------

    def infer_format(self, messages: list[bytes | NDArray[np.uint8]]) -> MessageSchema:
        """Infer message format from collection of similar messages.

        : Complete format inference workflow.

        Args:
            messages: List of message samples (bytes or np.ndarray)

        Returns:
            MessageSchema with inferred field structure

        Raises:
            ValueError: If insufficient samples or invalid input
        """
        if len(messages) < self.min_samples:
            raise ValueError(f"Need at least {self.min_samples} messages, got {len(messages)}")

        msg_arrays = _convert_to_arrays(messages)
        _validate_array_lengths(msg_arrays)
        msg_len = len(msg_arrays[0])

        boundaries = detect_field_boundaries(msg_arrays, method="combined")
        fields = self.detect_field_types(msg_arrays, boundaries)
        header_size = _estimate_header_size(fields)

        checksum_field = None
        length_field = None
        for f in fields:
            if f.field_type == "checksum":
                checksum_field = f
            elif f.field_type == "length":
                length_field = f

        return MessageSchema(
            total_size=msg_len,
            fields=fields,
            field_boundaries=boundaries,
            header_size=header_size,
            payload_offset=header_size,
            checksum_field=checksum_field,
            length_field=length_field,
        )

    def infer_format_ensemble(
        self,
        messages: list[bytes | NDArray[np.uint8]],
        min_field_confidence: float = 0.6,
        min_boundary_confidence: float = 0.6,
    ) -> MessageSchema:
        """Infer message format using ensemble of techniques.

        : Ensemble-based format inference with confidence scoring.

        Combines:
        - Voting expert for boundary detection
        - Multiple field type detectors
        - Confidence scoring for each field

        Args:
            messages: List of protocol messages
            min_field_confidence: Minimum confidence to include field
            min_boundary_confidence: Minimum confidence for boundaries

        Returns:
            Message schema with confidence-scored fields

        Raises:
            ValueError: If insufficient messages provided
        """
        if len(messages) < self.min_samples:
            raise ValueError(f"Need at least {self.min_samples} messages, got {len(messages)}")

        bytes_messages = _convert_messages_to_bytes(messages)
        msg_len = _validate_bytes_lengths(bytes_messages)
        boundaries = detect_boundaries_voting(
            bytes_messages, min_confidence=min_boundary_confidence
        )
        msg_arrays = [np.frombuffer(msg, dtype=np.uint8) for msg in bytes_messages]
        fields = _infer_fields_from_boundaries(
            boundaries, msg_arrays, msg_len, min_field_confidence
        )
        checksum_field, length_field = _find_special_fields(fields)
        header_size = _estimate_header_size(fields)

        return MessageSchema(
            total_size=msg_len,
            fields=fields,
            field_boundaries=boundaries,
            header_size=header_size,
            payload_offset=header_size,
            checksum_field=checksum_field,
            length_field=length_field,
        )

    def detect_field_types(
        self,
        messages: list[NDArray[np.uint8]],
        boundaries: list[int],
    ) -> list[InferredField]:
        """Classify field types based on value patterns.

        : Field type classification.

        Args:
            messages: List of message arrays
            boundaries: Field boundary offsets

        Returns:
            List of InferredField objects
        """
        fields: list[InferredField] = []
        msg_len = len(messages[0])

        for i in range(len(boundaries)):
            offset = boundaries[i]
            if i < len(boundaries) - 1:
                size = boundaries[i + 1] - offset
            else:
                size = msg_len - offset

            field_data = extract_field_data(messages, offset, size)
            field_type, confidence = classify_field(field_data["values"], offset, size, messages)

            field_obj = InferredField(
                name=f"field_{i}",
                offset=offset,
                size=size,
                field_type=field_type,  # type: ignore[arg-type]
                entropy=float(field_data["entropy"]),
                variance=float(field_data["variance"]),
                confidence=confidence,
                values_seen=field_data["values"][:5],
            )
            fields.append(field_obj)

        return fields

    def find_dependencies(
        self,
        messages: list[NDArray[np.uint8]],
        schema: MessageSchema,
    ) -> dict[str, str]:
        """Find dependencies between fields (e.g., length->payload).

        : Field dependency detection.

        Args:
            messages: List of message arrays
            schema: Inferred message schema

        Returns:
            Dictionary mapping field names to dependency descriptions
        """
        dependencies: dict[str, str] = {}

        for field in schema.fields:
            if field.field_type == "length":
                for msg in messages:
                    _length_val = _extract_field_value(msg, field)
                dependencies[field.name] = "Potential length indicator"

        return dependencies

    # -----------------------------------------------------------------
    # Delegating boundary detection methods (keep public interface)
    # -----------------------------------------------------------------

    def detect_field_boundaries(
        self,
        messages: list[NDArray[np.uint8]],
        method: Literal["entropy", "variance", "combined"] = "combined",
    ) -> list[int]:
        """Detect field boundaries using entropy transitions.

        Delegates to :func:`boundaries.detect_field_boundaries`.

        Args:
            messages: List of message arrays
            method: Detection method ('entropy', 'variance', or 'combined')

        Returns:
            List of byte offsets marking field starts (always includes 0)
        """
        return detect_field_boundaries(messages, method=method)

    def detect_boundaries_voting(
        self,
        messages: list[bytes],
        min_confidence: float = 0.6,
    ) -> list[int]:
        """Detect field boundaries using voting expert algorithm.

        Delegates to :func:`boundaries.detect_boundaries_voting`.

        Args:
            messages: List of protocol messages (bytes)
            min_confidence: Minimum vote fraction to accept boundary

        Returns:
            List of byte positions that are likely field boundaries
        """
        return detect_boundaries_voting(messages, min_confidence=min_confidence)

    # -----------------------------------------------------------------
    # Delegating advanced detector methods (keep public interface)
    # -----------------------------------------------------------------

    def detect_timestamp_field(
        self,
        messages: list[NDArray[np.uint8]],
        offset: int,
        size: int,
    ) -> FieldDetectionResult:
        """Detect timestamp fields in message data.

        Delegates to :func:`advanced_detectors.detect_timestamp_field`.
        """
        return detect_timestamp_field(messages, offset, size)

    def detect_float_field(
        self,
        messages: list[NDArray[np.uint8]],
        offset: int,
        size: int,
    ) -> FieldDetectionResult:
        """Detect IEEE 754 floating-point fields.

        Delegates to :func:`advanced_detectors.detect_float_field`.
        """
        return detect_float_field(messages, offset, size)

    def detect_length_field(
        self,
        messages: list[NDArray[np.uint8]],
        offset: int,
        size: int,
        msg_len: int,
    ) -> FieldDetectionResult:
        """Detect length fields with endianness detection.

        Delegates to :func:`advanced_detectors.detect_length_field`.
        """
        return detect_length_field(messages, offset, size, msg_len)

    def detect_enum_field(
        self,
        messages: list[NDArray[np.uint8]],
        offset: int,
        size: int,
    ) -> FieldDetectionResult:
        """Detect enumeration fields.

        Delegates to :func:`advanced_detectors.detect_enum_field`.
        """
        return detect_enum_field(messages, offset, size)

    def detect_reserved_field(
        self,
        messages: list[NDArray[np.uint8]],
        offset: int,
        size: int,
    ) -> FieldDetectionResult:
        """Detect reserved or padding fields.

        Delegates to :func:`advanced_detectors.detect_reserved_field`.
        """
        return detect_reserved_field(messages, offset, size)

    # -----------------------------------------------------------------
    # Private methods (for testing/internal use)
    # -----------------------------------------------------------------

    def _calculate_byte_entropy(
        self,
        messages: list[NDArray[np.uint8]],
        offset: int,
    ) -> float:
        """Calculate byte entropy at given offset.

        Delegates to :func:`entropy.calculate_byte_entropy`.
        """
        return calculate_byte_entropy(messages, offset)

    def _calculate_entropy(self, values: NDArray[np.uint8]) -> float:
        """Calculate Shannon entropy of values.

        Delegates to :func:`entropy.calculate_entropy`.
        """
        return calculate_entropy(values)

    def _detect_counter_field(self, values: NDArray[np.uint8] | list[int]) -> bool:
        """Detect if field values represent a counter.

        Delegates to :func:`field_classification.detect_counter_field`.
        """
        # Convert to list if needed for compatibility with the function signature
        if isinstance(values, np.ndarray):
            return detect_counter_field(values.tolist())
        return detect_counter_field(values)

    def _expert_entropy(self, messages: list[NDArray[np.uint8]]) -> set[int]:
        """Entropy-based boundary expert.

        Delegates to :func:`boundaries.expert_entropy`.
        """
        return expert_entropy(messages)

    def _expert_alignment(self, messages: list[NDArray[np.uint8]] | list[bytes]) -> set[int]:
        """Alignment-based boundary expert.

        Delegates to :func:`boundaries.expert_alignment`.
        """
        # Convert to bytes for compatibility with the function signature
        bytes_messages: list[bytes] = []
        for msg in messages:
            if isinstance(msg, bytes):
                bytes_messages.append(msg)
            else:
                bytes_messages.append(msg.tobytes())
        return expert_alignment(bytes_messages)

    def _expert_variance(self, messages: list[NDArray[np.uint8]]) -> set[int]:
        """Variance-based boundary expert.

        Delegates to :func:`boundaries.expert_variance`.
        """
        return expert_variance(messages)

    def _expert_distribution(self, messages: list[NDArray[np.uint8]]) -> set[int]:
        """Distribution-based boundary expert.

        Delegates to :func:`boundaries.expert_distribution`.
        """
        return expert_distribution(messages)

    def _expert_ngrams(self, messages: list[NDArray[np.uint8]], n: int = 2) -> set[int]:
        """N-gram based boundary expert.

        Delegates to :func:`boundaries.expert_ngrams`.
        """
        return expert_ngrams(messages, n)

    def _detect_checksum_field(
        self,
        messages: list[NDArray[np.uint8]],
        offset: int,
        size: int,
    ) -> bool:
        """Detect checksum fields.

        Delegates to :func:`field_classification.detect_checksum_field`.
        """
        return detect_checksum_field(messages, offset, size)

    def _estimate_header_size(self, fields: list[InferredField]) -> int:
        """Estimate header size from fields.

        Delegates to module-level :func:`_estimate_header_size`.
        """
        return _estimate_header_size(fields)

    def _extract_field_value(self, msg: NDArray[np.uint8], field: InferredField) -> int:
        """Extract field value from message.

        Delegates to module-level :func:`_extract_field_value`.
        """
        return _extract_field_value(msg, field)

    def _extract_field_data(
        self,
        messages: list[NDArray[np.uint8]],
        offset: int,
        size: int,
    ) -> dict[str, float | NDArray[np.uint8]]:
        """Extract field data for analysis.

        Delegates to :func:`field_classification.extract_field_data`.
        """
        return extract_field_data(messages, offset, size)

    def _detect_type_entropy(
        self, field_data: dict[str, float | NDArray[np.uint8]]
    ) -> tuple[str, float]:
        """Detect field type based on entropy.

        Delegates to :func:`field_classification.detect_type_entropy`.
        """
        return detect_type_entropy(field_data)

    def _detect_type_patterns(
        self,
        field_data: dict[str, float | NDArray[np.uint8]],
        offset: int,
        size: int,
        msg_len: int,
    ) -> tuple[str, float]:
        """Detect field type based on patterns.

        Delegates to :func:`field_classification.detect_type_patterns`.
        """
        return detect_type_patterns(field_data, offset, size, msg_len)

    def _detect_type_statistics(
        self, field_data: dict[str, float | NDArray[np.uint8]]
    ) -> tuple[str, float]:
        """Detect field type based on statistics.

        Delegates to :func:`field_classification.detect_type_statistics`.
        """
        return detect_type_statistics(field_data)

    def _vote_field_type(
        self,
        detections: list[tuple[str, float]],
    ) -> tuple[str, float, dict[str, bool]]:
        """Vote on field type from multiple detections.

        Delegates to :func:`field_classification.vote_field_type`.
        """
        return vote_field_type(detections)


# =====================================================================
# Module-private helpers (extracted from former class methods)
# =====================================================================


def _convert_to_arrays(
    messages: list[bytes | NDArray[np.uint8]],
) -> list[NDArray[np.uint8]]:
    """Convert heterogeneous message list to numpy arrays."""
    msg_arrays: list[NDArray[np.uint8]] = []
    for msg in messages:
        if isinstance(msg, bytes):
            msg_arrays.append(np.frombuffer(msg, dtype=np.uint8))
        elif isinstance(msg, np.ndarray):
            msg_arrays.append(msg.astype(np.uint8))
        else:
            raise ValueError(f"Invalid message type: {type(msg)}")
    return msg_arrays


def _validate_array_lengths(msg_arrays: list[NDArray[np.uint8]]) -> int:
    """Validate all arrays have the same length; return that length."""
    lengths = [len(m) for m in msg_arrays]
    if len(set(lengths)) > 1:
        raise ValueError(f"Messages have varying lengths: {set(lengths)}")
    return lengths[0]


def _convert_messages_to_bytes(
    messages: list[bytes | NDArray[np.uint8]],
) -> list[bytes]:
    """Convert messages to bytes format."""
    bytes_messages: list[bytes] = []
    for msg in messages:
        if isinstance(msg, bytes):
            bytes_messages.append(msg)
        elif isinstance(msg, np.ndarray):
            bytes_messages.append(msg.tobytes())
        else:
            raise ValueError(f"Invalid message type: {type(msg)}")
    return bytes_messages


def _validate_bytes_lengths(messages: list[bytes]) -> int:
    """Validate all byte messages have the same length; return it."""
    lengths = [len(m) for m in messages]
    if len(set(lengths)) > 1:
        raise ValueError(f"Messages have varying lengths: {set(lengths)}")
    return lengths[0]


def _infer_fields_from_boundaries(
    boundaries: list[int],
    msg_arrays: list[NDArray[np.uint8]],
    msg_len: int,
    min_confidence: float,
) -> list[InferredField]:
    """Infer fields from detected boundaries using ensemble classification."""
    fields: list[InferredField] = []
    for i in range(len(boundaries)):
        offset = boundaries[i]
        size = boundaries[i + 1] - offset if i < len(boundaries) - 1 else msg_len - offset
        field_data = extract_field_data(msg_arrays, offset, size)
        field_obj = classify_field_ensemble(field_data, offset, size, msg_len, min_confidence)
        if field_obj is not None:
            field_obj.name = f"field_{len(fields)}"
            fields.append(field_obj)
    return fields


def _find_special_fields(
    fields: list[InferredField],
) -> tuple[InferredField | None, InferredField | None]:
    """Find checksum and length fields."""
    checksum_field = None
    length_field = None
    for f in fields:
        if f.field_type == "checksum":
            checksum_field = f
        elif f.field_type == "length":
            length_field = f
    return checksum_field, length_field


def _estimate_header_size(fields: list[InferredField]) -> int:
    """Estimate header size from field patterns.

    Args:
        fields: List of inferred fields

    Returns:
        Estimated header size in bytes
    """
    for i, field in enumerate(fields):
        if field.field_type == "data" and field.entropy > 6.0:
            if i > 0:
                return field.offset

    if len(fields) >= 5:
        return fields[4].offset
    elif len(fields) >= 4:
        return fields[3].offset + fields[3].size
    elif fields:
        return min(16, fields[-1].offset)
    else:
        return 16


def _extract_field_value(msg: NDArray[np.uint8], field: InferredField) -> int:
    """Extract field value from message.

    Args:
        msg: Message array
        field: Field definition

    Returns:
        Field value as integer
    """
    if field.size == 1:
        return int(msg[field.offset])
    elif field.size == 2:
        return int(msg[field.offset]) << 8 | int(msg[field.offset + 1])
    elif field.size == 4:
        return (
            int(msg[field.offset]) << 24
            | int(msg[field.offset + 1]) << 16
            | int(msg[field.offset + 2]) << 8
            | int(msg[field.offset + 3])
        )
    else:
        return int(msg[field.offset])


# =====================================================================
# Module-level convenience functions
# =====================================================================


def infer_format(
    messages: list[bytes | NDArray[np.uint8]],
    min_samples: int = 10,
) -> MessageSchema:
    """Convenience function for format inference.

    : Top-level API for message format inference.

    Args:
        messages: List of message samples (bytes or np.ndarray)
        min_samples: Minimum required samples

    Returns:
        MessageSchema with inferred structure
    """
    inferrer = MessageFormatInferrer(min_samples=min_samples)
    return inferrer.infer_format(messages)


def detect_field_types_standalone(
    messages: list[bytes | NDArray[np.uint8]] | bytes | NDArray[np.uint8],
    boundaries: list[int] | None = None,
) -> list[InferredField]:
    """Detect field types at boundaries.

    : Field type detection.

    Args:
        messages: List of message samples OR a single message
        boundaries: Field boundary offsets (auto-detected if not provided)

    Returns:
        List of InferredField objects

    Raises:
        ValueError: If message type is invalid.
    """
    inferrer = MessageFormatInferrer()

    if isinstance(messages, (bytes, np.ndarray)):
        messages_list: list[bytes | NDArray[np.uint8]] = [messages]
    else:
        messages_list = messages

    msg_arrays = _convert_to_arrays(messages_list)

    if boundaries is None:
        boundaries = detect_field_boundaries(msg_arrays, method="combined")

    return inferrer.detect_field_types(msg_arrays, boundaries)


def find_dependencies_standalone(
    messages: list[bytes | NDArray[np.uint8]],
    schema: MessageSchema | None = None,
) -> dict[str, str]:
    """Find field dependencies.

    : Field dependency analysis.

    Args:
        messages: List of message samples
        schema: Message schema (auto-inferred if not provided)

    Returns:
        Dictionary of dependencies

    Raises:
        ValueError: If message type is invalid.
    """
    inferrer = MessageFormatInferrer()
    msg_arrays = _convert_to_arrays(messages)

    if schema is None:
        schema = inferrer.infer_format(msg_arrays)  # type: ignore[arg-type]

    return inferrer.find_dependencies(msg_arrays, schema)


# Keep the original public names via aliases so that __init__.py can
# export them with the same names as the old monolithic module.
detect_field_types_fn = detect_field_types_standalone
find_dependencies_fn = find_dependencies_standalone
