"""Output module for the Excel Spreadsheet (XLSX) output format."""

import datetime
import os
import re

import pytz
import xlsxwriter

from plaso.output import dynamic
from plaso.output import interface
from plaso.output import manager


class XLSXOutputModule(interface.OutputModule):
    """Output module for the Excel Spreadsheet (XLSX) output format."""

    NAME = "xlsx"
    DESCRIPTION = "Excel Spreadsheet (XLSX) output"

    SUPPORTS_ADDITIONAL_FIELDS = True
    SUPPORTS_CUSTOM_FIELDS = True

    WRITES_OUTPUT_FILE = True

    _DEFAULT_FIELDS = [
        "datetime",
        "timestamp_desc",
        "source",
        "source_long",
        "message",
        "parser",
        "display_name",
        "tag",
    ]

    _DEFAULT_TIMESTAMP_FORMAT = "YYYY-MM-DD HH:MM:SS.000"

    _MAXIMUM_COLUMN_WIDTH = 50
    _MINIMUM_COLUMN_WIDTH = 6

    # Illegal XML string characters.
    _ILLEGAL_XML_RE = re.compile(
        (
            r"[\x00-\x08\x0b-\x1f\x7f-\x84\x86-\x9f\ud800-\udfff\ufdd0-\ufddf"
            r"\ufffe-\uffff]"
        )
    )

    def __init__(self):
        """Initializes an output module."""
        super().__init__()
        self._column_widths = []
        self._current_row = 0
        self._custom_fields = {}
        self._field_formatting_helper = dynamic.DynamicFieldFormattingHelper()
        self._field_names = self._DEFAULT_FIELDS
        self._sheet = None
        self._timestamp_format = self._DEFAULT_TIMESTAMP_FORMAT
        self._workbook = None

    def _FormatDateTime(
        self, output_mediator, event, event_data
    ):  # pylint: disable=missing-return-type-doc
        """Formats the date to a datetime object without timezone information.

        Note: timezone information must be removed due to lack of support
        by xlsxwriter and Excel.

        Args:
          output_mediator (OutputMediator): mediates interactions between output
              modules and other components, such as storage and dfVFS.
          event (EventObject): event.
          event_data (EventData): event data.

        Returns:
          datetime.datetime|str: date and time value or a string containing
              "ERROR" on OverflowError.
        """
        try:
            datetime_object = datetime.datetime(1970, 1, 1, 0, 0, 0, 0, tzinfo=pytz.UTC)
            datetime_object += datetime.timedelta(microseconds=event.timestamp)
            datetime_object.astimezone(output_mediator.time_zone)

            return datetime_object.replace(tzinfo=None)

        except (OSError, OverflowError, TypeError, ValueError) as exception:
            message = (
                f"unable to copy timestamp: {event.timestamp!s} to a human readable "
                f'date and time with error: {exception!s}. Defaulting to: "ERROR"'
            )
            self._ReportEventError(
                event,
                event_data,
                message,
            )
            return "ERROR"

    def _SanitizeField(self, field):
        """Sanitizes a field for output.

        This method replaces any illegal XML string characters with the Unicode
        replacement character (\ufffd).

        Args:
          field (str): value of the field to sanitize.

        Returns:
          str: sanitized value of the field.
        """
        return self._ILLEGAL_XML_RE.sub("\ufffd", field)

    def Close(self):
        """Closes the workbook."""
        for column_index, column_width in enumerate(self._column_widths):
            self._sheet.set_column(column_index, column_index, column_width)

        self._workbook.close()
        self._workbook = None

    def GetFieldValues(
        self, output_mediator, event, event_data, event_data_stream, event_tag
    ):
        """Retrieves the output field values.

        Args:
          output_mediator (OutputMediator): mediates interactions between output
              modules and other components, such as storage and dfVFS.
          event (EventObject): event.
          event_data (EventData): event data.
          event_data_stream (EventDataStream): event data stream.
          event_tag (EventTag): event tag.

        Returns:
          dict[str, str]: output field values per name.
        """
        field_values = {}
        for field_name in self._field_names:
            if field_name == "datetime":
                field_values["datetime"] = self._FormatDateTime(
                    output_mediator, event, event_data
                )
                continue

            field_value = self._field_formatting_helper.GetFormattedField(
                output_mediator,
                field_name,
                event,
                event_data,
                event_data_stream,
                event_tag,
            )
            if field_value is None and field_name in self._custom_fields:
                field_value = self._custom_fields.get(field_name)

            if field_value is None:
                field_value = "-"
            else:
                field_value = self._SanitizeField(field_value)

            field_values[field_name] = field_value

        return field_values

    def Open(self, path=None, **kwargs):  # pylint: disable=arguments-differ
        """Creates a new workbook.

        Args:
          path (Optional[str]): path of the output file.

        Raises:
          OSError: if the specified output file already exists.
          ValueError: if path is not set.
        """
        if not path:
            raise ValueError("Missing filename.")

        if os.path.isfile(path):
            raise OSError(
                f"Unable to use an already existing file for output [{path:s}]"
            )

        options = {
            "constant_memory": True,
            "strings_to_urls": False,
            "strings_to_formulas": False,
            "default_date_format": self._timestamp_format,
        }
        self._workbook = xlsxwriter.Workbook(path, options)
        self._sheet = self._workbook.add_worksheet("Sheet")
        self._current_row = 0

    def SetAdditionalFields(self, field_names):
        """Sets the names of additional fields to output.

        Args:
          field_names (list[str]): names of additional fields to output.
        """
        self._field_names.extend(field_names)

    def SetCustomFields(self, field_names_and_values):
        """Sets the names and values of custom fields to output.

        Args:
          field_names_and_values (list[tuple[str, str]]): names and values of
              custom fields to output.
        """
        self._custom_fields = dict(field_names_and_values)
        self._field_names.extend(self._custom_fields.keys())

    def SetFields(self, field_names):
        """Sets the names of the fields to output.

        Args:
          field_names (list[str]): names of the fields to output.
        """
        self._field_names = field_names

    def SetTimestampFormat(self, timestamp_format):
        """Set the timestamp format to use for the datetime column.

        Args:
          timestamp_format (str): format string of date and time values.
        """
        self._timestamp_format = timestamp_format

    def WriteFieldValues(self, output_mediator, field_values):
        """Writes field values to the output.

        Args:
          output_mediator (OutputMediator): mediates interactions between output
              modules and other components, such as storage and dfVFS.
          field_values (dict[str, str]): output field values per name.
        """
        for column_index, field_name in enumerate(self._field_names):
            field_value = field_values.get(field_name)
            if field_name == "datetime" and field_value:
                self._sheet.write_datetime(self._current_row, column_index, field_value)
                column_width = len(self._timestamp_format) + 2
            else:
                field_value = field_value or ""

                self._sheet.write(self._current_row, column_index, field_value)
                column_width = len(field_value) + 2

            # Auto adjust the column width based on the length of the output value.
            column_width = min(column_width, self._MAXIMUM_COLUMN_WIDTH)
            column_width = max(
                column_width,
                self._MINIMUM_COLUMN_WIDTH,
                self._column_widths[column_index],
            )
            self._column_widths[column_index] = column_width

        self._current_row += 1

    def WriteHeader(self, output_mediator):
        """Writes the header to the spreadsheet.

        Args:
          output_mediator (OutputMediator): mediates interactions between output
              modules and other components, such as storage and dfVFS.
        """
        cell_format = self._workbook.add_format({"bold": True})
        cell_format.set_align("center")

        self._column_widths = []
        for column_index, field_name in enumerate(self._field_names):
            self._sheet.write(0, column_index, field_name, cell_format)

            column_width = len(field_name) + 2
            self._column_widths.append(column_width)

        self._current_row = 1
        self._sheet.autofilter(0, len(self._field_names) - 1, 0, 0)
        self._sheet.freeze_panes(1, 0)


manager.OutputManager.RegisterOutput(XLSXOutputModule)
