"""
Inference script for Engine B (Malimg).

This module provides functionality to read raw binary files, convert them
to 2D grayscale image tensors, and classify their malware family using
a pre-trained Convolutional Neural Network (CNN).
"""

import os
import torch
from torchvision import transforms
from PIL import Image
import numpy as np
import json
from src.engine_b.model import MalwareCNN
import math


def bytes_to_image(file_path):
    """
    Reads a raw binary file and converts its bytes into a square grayscale image.

    If the file is a mock JSON profile, it reads the byte array from the JSON.
    Otherwise, it reads the raw bytes of the executable and reshapes them into
    a square 2D matrix, padding the end with zeros if necessary.

    Args:
        file_path (str): Path to the target binary file.

    Returns:
        PIL.Image: A grayscale (mode 'L') Image object representing the binary.
    """
    # Check if this is our safe mock profile
    if file_path.endswith(".json"):
        try:
            with open(file_path, "r") as f:
                data = json.load(f)
            if data.get("is_mock_profile"):
                byte_array = np.array(data["malimg_bytes"], dtype=np.uint8)
            else:
                with open(file_path, "rb") as f:
                    binary_data = f.read()
                byte_array = np.frombuffer(binary_data, dtype=np.uint8)
        except:
            with open(file_path, "rb") as f:
                binary_data = f.read()
            byte_array = np.frombuffer(binary_data, dtype=np.uint8)
    else:
        # Read raw binary
        with open(file_path, "rb") as f:
            binary_data = f.read()
        byte_array = np.frombuffer(binary_data, dtype=np.uint8)

    # Calculate image dimensions (square)
    length = len(byte_array)
    if length == 0:
        return Image.new("L", (128, 128), color=0)

    width = int(math.ceil(math.sqrt(length)))
    height = width

    # Pad array to form a perfect square
    padded_length = width * height
    padded_array = np.pad(byte_array, (0, padded_length - length), mode="constant")

    # Reshape and create PIL Image
    image_2d = padded_array.reshape((height, width))
    img = Image.fromarray(image_2d, mode="L")
    return img


class EngineBInfer:
    """
    Inference Engine for Visual Malware Family Classification.

    Attributes:
        device (torch.device): CPU or CUDA device for inference.
        classes (list): Ordered list of malware family class names.
        model (MalwareCNN): The loaded PyTorch CNN.
        transform (transforms.Compose): Image preprocessing pipeline.
    """

    def __init__(
        self,
        model_path="models/engine_b_model.pth",
        classes_path="models/engine_b_classes.json",
    ):
        """
        Initializes the vision inference engine.

        Args:
            model_path (str): Path to the trained PyTorch state dictionary.
            classes_path (str): Path to the JSON list mapping indices to class names.
        """
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.classes = []
        if os.path.exists(classes_path):
            with open(classes_path, "r") as f:
                self.classes = json.load(f)
        else:
            # Fallback dummy classes if not trained yet
            self.classes = [f"Class_{i}" for i in range(24)]

        num_classes = len(self.classes)
        self.model = MalwareCNN(num_classes=num_classes)

        if os.path.exists(model_path):
            self.model.load_state_dict(
                torch.load(model_path, map_location=self.device, weights_only=True)
            )
        else:
            print(f"Warning: {model_path} not found. Using untrained weights.")

        self.model.to(self.device)
        self.model.eval()

        self.transform = transforms.Compose(
            [
                transforms.Resize((128, 128)),
                transforms.ToTensor(),
            ]
        )

    def predict(self, file_path):
        """
        Converts the target file to an image and runs CNN inference.

        Args:
            file_path (str): Path to the target file.

        Returns:
            dict: Contains 'family' (str), 'confidence' (float),
                  'all_probabilities' (dict), and 'image' (PIL.Image upscaled).
        """
        img = bytes_to_image(file_path)
        tensor = self.transform(img).unsqueeze(0).to(self.device)

        with torch.no_grad():
            outputs = self.model(tensor)
            probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]

            # Create a dictionary of all class probabilities
            all_probs = {
                self.classes[i]: probabilities[i].item()
                for i in range(len(self.classes))
            }

            top_prob, top_class = torch.max(probabilities, 0)
            class_name = self.classes[top_class.item()]

            return {
                "family": class_name,
                "confidence": top_prob.item(),
                "all_probabilities": all_probs,
                "image": img.resize((512, 512), resample=Image.NEAREST),
            }