"""PE file parser for Resource Hacker.

This module provides functionality to parse Windows Portable Executable (PE) files
and extract resource information using the pefile library.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path

import pefile

logger = logging.getLogger(__name__)

# Resource type constants
RESOURCE_TYPES = {
    pefile.RESOURCE_TYPE["RT_CURSOR"]: "Cursor",
    pefile.RESOURCE_TYPE["RT_BITMAP"]: "Bitmap",
    pefile.RESOURCE_TYPE["RT_ICON"]: "Icon",
    pefile.RESOURCE_TYPE["RT_MENU"]: "Menu",
    pefile.RESOURCE_TYPE["RT_DIALOG"]: "Dialog",
    pefile.RESOURCE_TYPE["RT_STRING"]: "String Table",
    pefile.RESOURCE_TYPE["RT_FONTDIR"]: "Font Directory",
    pefile.RESOURCE_TYPE["RT_FONT"]: "Font",
    pefile.RESOURCE_TYPE["RT_ACCELERATOR"]: "Accelerator",
    pefile.RESOURCE_TYPE["RT_RCDATA"]: "Raw Data",
    pefile.RESOURCE_TYPE["RT_MESSAGETABLE"]: "Message Table",
    pefile.RESOURCE_TYPE["RT_GROUP_CURSOR"]: "Group Cursor",
    pefile.RESOURCE_TYPE["RT_GROUP_ICON"]: "Group Icon",
    pefile.RESOURCE_TYPE["RT_VERSION"]: "Version",
    pefile.RESOURCE_TYPE["RT_DLGINCLUDE"]: "Dialog Include",
    pefile.RESOURCE_TYPE["RT_PLUGPLAY"]: "Plug and Play",
    pefile.RESOURCE_TYPE["RT_VXD"]: "VXD",
    pefile.RESOURCE_TYPE["RT_ANICURSOR"]: "Animated Cursor",
    pefile.RESOURCE_TYPE["RT_ANIICON"]: "Animated Icon",
    pefile.RESOURCE_TYPE["RT_HTML"]: "HTML",
    pefile.RESOURCE_TYPE["RT_MANIFEST"]: "Manifest",
}


@dataclass
class ResourceEntry:
    """Represents a single resource entry."""

    type_id: int
    type_name: str
    name_id: int
    name: str
    language_id: int
    offset: int
    size: int
    data: bytes


@dataclass
class PEFileInfo:
    """Information about a PE file."""

    file_path: Path
    machine: str
    subsystem: str
    is_dll: bool
    is_64bit: bool
    image_base: int
    entry_point: int
    resources: list[ResourceEntry]


class PEParser:
    """Parser for Windows PE files."""

    def __init__(self):
        self.pe: pefile.PE | None = None
        self.file_info: PEFileInfo | None = None

    def load_file(self, file_path: Path) -> bool:
        """Load a PE file for parsing.

        Args:
            file_path: Path to the PE file

        Returns
        -------
            bool: True if file loaded successfully, False otherwise
        """
        try:
            logger.info(f"Loading PE file: {file_path}")
            self.pe = pefile.PE(str(file_path))
            self.file_info = self._extract_file_info(file_path)
            logger.info(f"Successfully loaded PE file: {file_path.name}")
            return True
        except Exception as e:
            logger.error(f"Failed to load PE file {file_path}: {e}")
            return False

    def _extract_file_info(self, file_path: Path) -> PEFileInfo:
        """Extract basic file information from loaded PE."""
        if not self.pe:
            raise ValueError("No PE file loaded")

        # Get machine type
        machine_types = {
            0x014C: "x86",
            0x8664: "x64",
            0x0200: "IA-64",
            0x01C0: "ARM",
            0xAA64: "ARM64",
        }
        machine = machine_types.get(self.pe.FILE_HEADER.Machine, "Unknown")

        # Get subsystem
        subsystems = {
            1: "Native",
            2: "Windows GUI",
            3: "Windows Console",
            7: "POSIX",
            9: "Windows CE GUI",
            10: "EFI Application",
            11: "EFI Boot Service Driver",
            12: "EFI Runtime Driver",
            13: "EFI ROM",
            14: "XBOX",
            16: "Windows Boot Application",
        }
        subsystem = subsystems.get(self.pe.OPTIONAL_HEADER.Subsystem, "Unknown")

        # Extract resources
        resources = self._extract_resources()

        return PEFileInfo(
            file_path=file_path,
            machine=machine,
            subsystem=subsystem,
            is_dll=self.pe.is_dll(),
            is_64bit=hex(self.pe.OPTIONAL_HEADER.Magic) == "0x20b",
            image_base=self.pe.OPTIONAL_HEADER.ImageBase,
            entry_point=self.pe.OPTIONAL_HEADER.AddressOfEntryPoint,
            resources=resources,
        )

    def _extract_resources(self) -> list[ResourceEntry]:
        """Extract all resources from the PE file."""
        resources = []

        if not self.pe or not hasattr(self.pe, "DIRECTORY_ENTRY_RESOURCE"):
            return resources

        try:
            for resource_type in self.pe.DIRECTORY_ENTRY_RESOURCE.entries:
                type_id = resource_type.id
                type_name = RESOURCE_TYPES.get(type_id, f"Unknown ({type_id})")

                if hasattr(resource_type, "directory"):
                    for resource_id in resource_type.directory.entries:
                        name_id = resource_id.id
                        name = str(name_id)

                        if hasattr(resource_id, "directory"):
                            for resource_lang in resource_id.directory.entries:
                                language_id = resource_lang.data.lang
                                offset = resource_lang.data.struct.OffsetToData
                                size = resource_lang.data.struct.Size

                                # Extract resource data
                                try:
                                    data = self.pe.get_data(offset, size)
                                except Exception:
                                    data = b""

                                resource = ResourceEntry(
                                    type_id=type_id,
                                    type_name=type_name,
                                    name_id=name_id,
                                    name=name,
                                    language_id=language_id,
                                    offset=offset,
                                    size=size,
                                    data=data,
                                )
                                resources.append(resource)

        except Exception as e:
            logger.error(f"Error extracting resources: {e}")

        return resources

    def get_resource_by_id(self, type_id: int, name_id: int, language_id: int = 0) -> ResourceEntry | None:
        """Get a specific resource by its identifiers.

        Args:
            type_id: Resource type identifier
            name_id: Resource name identifier
            language_id: Language identifier (default: 0)

        Returns
        -------
            ResourceEntry if found, None otherwise
        """
        if not self.file_info:
            return None

        for resource in self.file_info.resources:
            if resource.type_id == type_id and resource.name_id == name_id and resource.language_id == language_id:
                return resource

        return None

    def get_resources_by_type(self, type_id: int) -> list[ResourceEntry]:
        """Get all resources of a specific type.

        Args:
            type_id: Resource type identifier

        Returns
        -------
            List of ResourceEntry objects
        """
        if not self.file_info:
            return []

        return [resource for resource in self.file_info.resources if resource.type_id == type_id]

    def get_resource_types(self) -> dict[int, str]:
        """Get all resource types present in the file.

        Returns
        -------
            Dictionary mapping type IDs to type names
        """
        if not self.file_info:
            return {}

        types = {}
        for resource in self.file_info.resources:
            types[resource.type_id] = resource.type_name
        return types

    def close(self):
        """Close the PE file and free resources."""
        if self.pe:
            self.pe.close()
            self.pe = None
            self.file_info = None
            logger.info("PE file closed")

    def __enter__(self):
        """Enter the PE parser context."""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit the PE parser context and close the file."""
        self.close()
