"""
Malware analyzer - processes JSON reports and binary files.
"""

import os
import sys
import logging
from tqdm import tqdm

from src.analyzers.base_analyzer import BaseAnalyzer
from src.config import Config
from src.utils.elf_utils import get_elf_info_with_pyelftools, get_elf_binary_info
from src.utils.packer_utils import (
    convert_to_one_line,
    get_family_using_avclass,
    run_diec_analysis
)


class MalwareAnalyzer(BaseAnalyzer):
    """Analyzer for malware datasets with JSON reports and binary files."""

    def __init__(self, config: Config):
        """
        Initialize the MalwareAnalyzer object.

        :param config: Configuration object containing input directory and output path.
        """
        super().__init__(config)

        # Validate malware mode requirements
        if not config.input_dir or not os.path.isdir(config.input_dir):
            print("Error: For malware mode, input_dir must be a valid directory containing JSON files.")
            sys.exit(1)

    def collect_files(self):
        """
        Recursively get all JSON files in the input directory and its subdirectories.
        """
        print(f"Searching for all JSON files in directory: {self.config.input_dir}...")

        # Use os.walk to recursively traverse the directory and all its subdirectories
        for root, dirs, files in tqdm(os.walk(self.config.input_dir),
                                     desc="Scanning directories", unit="dir"):
            for file in files:
                # Only process JSON files
                if file.endswith('.json'):
                    file_path = os.path.join(root, file)
                    self.file_list.append(file_path)

        print(f"Found {len(self.file_list)} JSON files")

    @staticmethod
    def process_single_file(json_file, binary_base_path=None):
        """
        Process a single JSON file, extracting required information.

        :param json_file: Path to the JSON file.
        :param binary_base_path: Base path for binary files.
        :return: Dictionary containing extracted information, or None if processing failed.
        """
        one_line_data, json_data = convert_to_one_line(json_file)

        # If JSON parsing failed, return None to indicate this record should be skipped
        if json_data is None:
            return None

        # Get file metadata
        sha256 = json_data.get('sha256', None)
        md5 = json_data.get('md5', None)
        size = json_data.get('size', 0)
        first_seen = json_data.get('first_seen', None)

        # Skip this record if sha256 is missing or empty
        if not sha256:
            return None

        # Initialize result dictionary
        result = {
            'file_name': sha256,
            'md5': md5,
            'label': None,
            'file_type': None,
            'CPU': None,
            'bits': None,
            'endianness': None,
            'load_segments': None,
            'is_stripped': None,
            'has_section_name': None,
            'family': None,
            'first_seen': first_seen,
            'size': size,
            'diec_is_packed': False,
            'diec_packer_info': None,
            'diec_packing_method': None
        }

        # Determine if it's malware
        positives = json_data.get('positives', 0)
        result['label'] = 'Malware' if positives > 0 else 'Benignware'

        # Extract endianness from gandelf information
        gandelf_info = json_data.get('additional_info', {}).get('gandelf', {}).get('header', {})
        result['endianness'] = gandelf_info.get('data', None)

        # Use AVClass to get the family
        result['family'] = get_family_using_avclass(json_file, one_line_data)

        # Build the potential path for the binary file
        # Get the first two characters of sha256 as a subdirectory
        # Check if sha256 is valid before trying to access it
        if sha256 and len(sha256) >= 2:
            prefix = sha256[:2]
            binary_path = os.path.join(binary_base_path, prefix, sha256)
        else:
            binary_path = None

        # Skip this record if binary file doesn't exist
        if not binary_path or not os.path.exists(binary_path):
            logging.debug(f"Binary file does not exist, skipping record for: {sha256}")
            return None

        # Use pyelftools to get CPU, endianness, file type, and stripped status
        elf_info = get_elf_info_with_pyelftools(binary_path)
        result['CPU'] = elf_info['cpu']
        result['endianness'] = elf_info['endianness']
        result['file_type'] = elf_info['file_type']
        result['is_stripped'] = elf_info['is_stripped']

        # Run diec analysis on the binary file
        diec_is_packed, diec_packer_info, diec_packing_method = run_diec_analysis(binary_path)
        result['diec_is_packed'] = diec_is_packed
        result['diec_packer_info'] = diec_packer_info
        result['diec_packing_method'] = diec_packing_method

        # Read ELF binary once to get bits, load segments, and section headers
        binary_info = get_elf_binary_info(binary_path)
        result['bits'] = binary_info['bits']
        result['load_segments'] = binary_info['load_segments']
        result['has_section_name'] = binary_info['has_section_name']

        return result
