import csv
import datetime
from collections import OrderedDict
from functools import partial
from io import BytesIO

from django.contrib.admin.utils import label_for_field
from django.core.exceptions import FieldDoesNotExist
from django.http import FileResponse, StreamingHttpResponse
from django.utils import timezone
from django.utils.dateformat import Formatter
from django.utils.encoding import force_str
from django.utils.formats import get_format
from django.utils.functional import cached_property
from django.utils.text import capfirst
from django.utils.translation import gettext as _

from wagtail.admin.widgets.button import Button
from wagtail.coreutils import multigetattr


class Echo:
    """An object that implements just the write method of the file-like interface."""

    def write(self, value):
        """Write the value by returning it, instead of storing in a buffer."""
        return value.encode("UTF-8")


def list_to_str(value):
    return force_str(", ".join(value))


class ExcelDateFormatter(Formatter):
    data = None

    # From: https://docs.djangoproject.com/en/stable/ref/templates/builtins/#date
    # To: https://support.microsoft.com/en-us/office/format-numbers-as-dates-or-times-418bd3fe-0577-47c8-8caa-b4d30c528309#bm2
    _formats = {
        # Day of the month, 2 digits with leading zeros.
        "d": "dd",
        # Day of the month without leading zeros.
        "j": "d",
        # Day of the week, textual, 3 letters.
        "D": "ddd",
        # Day of the week, textual, full.
        "l": "dddd",
        # English ordinal suffix for the day of the month, 2 characters.
        "S": "",  # Not supported in Excel
        # Day of the week, digits without leading zeros.
        "w": "",  # Not supported in Excel
        # Day of the year.
        "z": "",  # Not supported in Excel
        # ISO-8601 week number of year, with weeks starting on Monday.
        "W": "",  # Not supported in Excel
        # Month, 2 digits with leading zeros.
        "m": "mm",
        # Month without leading zeros.
        "n": "m",
        # Month, textual, 3 letters.
        "M": "mmm",
        # Month, textual, 3 letters, lowercase. (Not supported in Excel)
        "b": "mmm",
        # Month, locale specific alternative representation usually used for long date representation.
        "E": "mmmm",  # Not supported in Excel
        # Month, textual, full.
        "F": "mmmm",
        # Month abbreviation in Associated Press style. Proprietary extension.
        "N": "mmm.",  # Approximation, wrong for May
        # Number of days in the given month.
        "t": "",  # Not supported in Excel
        # Year, 2 digits with leading zeros.
        "y": "yy",
        # Year, 4 digits with leading zeros.
        "Y": "yyyy",
        # Whether it's a leap year.
        "L": "",  # Not supported in Excel
        # ISO-8601 week-numbering year.
        "o": "yyyy",  # Approximation, same as Y
        # Hour, 12-hour format without leading zeros.
        "g": "h",  # Only works when combined with AM/PM, 24-hour format is used otherwise
        # Hour, 24-hour format without leading zeros.
        "G": "hH",
        # Hour, 12-hour format with leading zeros.
        "h": "hh",  # Only works when combined with AM/PM, 24-hour format is used otherwise
        # Hour, 24-hour format with leading zeros.
        "H": "hh",
        # Minutes.
        "i": "mm",
        # Seconds.
        "s": "ss",
        # Microseconds.
        "u": ".00",  # Only works when combined with ss
        # 'a.m.' or 'p.m.'.
        "a": "AM/PM",  # Approximation, uses AM/PM and only works when combined with h/hh
        # AM/PM.
        "A": "AM/PM",  # Only works when combined with h/hh
        # Time, in 12-hour hours and minutes, with minutes left off if they’re zero.
        "f": "h:mm",  # Approximation, uses 24-hour format and minutes are never left off
        # Time, in 12-hour hours, minutes and ‘a.m.’/’p.m.’, with minutes left off if they’re zero and the special-case strings ‘midnight’ and ‘noon’ if appropriate.
        "P": "h:mm AM/PM",  # Approximation, minutes are never left off, no special case strings
        # Timezone name.
        "e": "",  # Not supported in Excel
        # Daylight saving time, whether it’s in effect or not.
        "I": "",  # Not supported in Excel
        # Difference to Greenwich time in hours.
        "O": "",  # Not supported in Excel
        # Time zone of this machine.
        "T": "",  # Not supported in Excel
        # Timezone offset in seconds.
        "Z": "",  # Not supported in Excel
        # ISO 8601 format.
        "c": "yyyy-mm-ddThh:mm:ss.00",
        # RFC 5322 formatted date.
        "r": "ddd, d mmm yyyy hh:mm:ss",
        # Seconds since the Unix epoch.
        "U": "",  # Not supported in Excel
    }

    def get(self):
        format = get_format("SHORT_DATETIME_FORMAT")
        return self.format(format)

    def __getattr__(self, name):
        if name in self._formats:
            return lambda: self._formats[name]
        raise AttributeError(
            f"'{type(self).__name__}' object has no attribute '{name}'"
        )


class SpreadsheetExportMixin:
    """A mixin for views, providing spreadsheet export functionality in csv and xlsx formats"""

    FORMAT_XLSX = "xlsx"
    FORMAT_CSV = "csv"
    FORMATS = (FORMAT_XLSX, FORMAT_CSV)

    # A list of fields or callables (without arguments) to export from each item in the queryset (dotted paths allowed)
    list_export = []
    # A dictionary of custom preprocessing functions by field and format (expected value would be of the form {field_name: {format: function}})
    # If a valid field preprocessing function is found, any applicable value preprocessing functions will not be used
    custom_field_preprocess = {}
    # A dictionary of preprocessing functions by value class and format
    custom_value_preprocess = {
        datetime.datetime: {
            FORMAT_XLSX: lambda value: (
                value
                if timezone.is_naive(value)
                else timezone.make_naive(value, datetime.timezone.utc)
            )
        },
        (datetime.date, datetime.time): {FORMAT_XLSX: None},
        list: {FORMAT_CSV: list_to_str, FORMAT_XLSX: list_to_str},
    }
    # A dictionary of column heading overrides in the format {field: heading}
    export_headings = {}

    export_buttons_template_name = "wagtailadmin/shared/export_buttons.html"

    export_filename = "spreadsheet-export"

    def setup(self, request, *args, **kwargs):
        super().setup(request, *args, **kwargs)
        self.is_export = request.GET.get("export") in self.FORMATS

    def get_paginate_by(self, queryset):
        if self.is_export:
            return None
        return super().get_paginate_by(queryset)

    def get_filename(self):
        """Gets the base filename for the exported spreadsheet, without extensions"""
        return self.export_filename

    def to_row_dict(self, item):
        """Returns an OrderedDict (in the order given by list_export) of the exportable information for a model instance"""
        row_dict = OrderedDict(
            (field, multigetattr(item, field)) for field in self.list_export
        )
        return row_dict

    def get_preprocess_function(self, field, value, export_format):
        """Returns the preprocessing function for a given field name, field value, and export format"""

        # Try to find a field specific function and return it
        format_dict = self.custom_field_preprocess.get(field, {})
        if export_format in format_dict:
            return format_dict[export_format]

        # Otherwise check for a value class specific function
        for value_classes, format_dict in self.custom_value_preprocess.items():
            if isinstance(value, value_classes) and export_format in format_dict:
                return format_dict[export_format]

        # Finally resort to force_str to prevent encoding errors
        return partial(force_str, strings_only=True)

    def preprocess_field_value(self, field, value, export_format):
        """Preprocesses a field value before writing it to the spreadsheet"""
        preprocess_function = self.get_preprocess_function(field, value, export_format)
        if preprocess_function is not None:
            return preprocess_function(value)
        else:
            return value

    def generate_xlsx_row(self, worksheet, row_dict, date_format=None):
        """Generate cells to append to the worksheet"""
        from openpyxl.cell import WriteOnlyCell

        for field, value in row_dict.items():
            cell = WriteOnlyCell(
                worksheet, self.preprocess_field_value(field, value, self.FORMAT_XLSX)
            )
            if date_format and isinstance(value, datetime.datetime):
                cell.number_format = date_format
            yield cell

    def write_csv_row(self, writer, row_dict):
        return writer.writerow(
            {
                field: self.preprocess_field_value(field, value, self.FORMAT_CSV)
                for field, value in row_dict.items()
            }
        )

    def get_heading(self, queryset, field):
        """Get the heading label for a given field for a spreadsheet generated from queryset"""
        heading_override = self.export_headings.get(field)
        if heading_override:
            return force_str(heading_override)
        try:
            return capfirst(force_str(label_for_field(field, queryset.model)))
        except (AttributeError, FieldDoesNotExist):
            return force_str(field)

    def stream_csv(self, queryset):
        """Generate a csv file line by line from queryset, to be used in a StreamingHTTPResponse"""
        writer = csv.DictWriter(Echo(), fieldnames=self.list_export)
        yield writer.writerow(
            {field: self.get_heading(queryset, field) for field in self.list_export}
        )

        for item in queryset:
            yield self.write_csv_row(writer, self.to_row_dict(item))

    def write_xlsx(self, queryset, output):
        """Write an xlsx workbook from a queryset"""
        from openpyxl import Workbook

        workbook = Workbook(write_only=True, iso_dates=True)

        worksheet = workbook.create_sheet(title="Sheet1")
        worksheet.append(
            self.get_heading(queryset, field) for field in self.list_export
        )

        date_format = ExcelDateFormatter().get()
        for item in queryset:
            worksheet.append(
                self.generate_xlsx_row(
                    worksheet, self.to_row_dict(item), date_format=date_format
                )
            )

        workbook.save(output)

    def write_xlsx_response(self, queryset):
        """Write an xlsx file from a queryset and return a FileResponse"""
        output = BytesIO()
        self.write_xlsx(queryset, output)
        output.seek(0)

        return FileResponse(
            output,
            as_attachment=True,
            content_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
            filename=f"{self.get_filename()}.xlsx",
        )

    def write_csv_response(self, queryset):
        stream = self.stream_csv(queryset)

        response = StreamingHttpResponse(stream, content_type="text/csv")
        response["Content-Disposition"] = 'attachment; filename="{}.csv"'.format(
            self.get_filename()
        )
        return response

    def as_spreadsheet(self, queryset, spreadsheet_format):
        """Return a response with a spreadsheet representing the exported data from queryset, in the format specified"""
        if spreadsheet_format == self.FORMAT_CSV:
            return self.write_csv_response(queryset)
        elif spreadsheet_format == self.FORMAT_XLSX:
            return self.write_xlsx_response(queryset)

    def get_export_url(self, format):
        params = self.request.GET.copy()
        params["export"] = format
        return self.request.path + "?" + params.urlencode()

    @property
    def xlsx_export_url(self):
        return self.get_export_url("xlsx")

    @property
    def csv_export_url(self):
        return self.get_export_url("csv")

    @cached_property
    def show_export_buttons(self):
        return bool(self.list_export)

    @cached_property
    def header_more_buttons(self):
        buttons = super().header_more_buttons.copy()
        if self.show_export_buttons:
            buttons.append(
                Button(
                    _("Download XLSX"),
                    url=self.xlsx_export_url,
                    icon_name="download",
                    priority=90,
                )
            )
            buttons.append(
                Button(
                    _("Download CSV"),
                    url=self.csv_export_url,
                    icon_name="download",
                    priority=100,
                )
            )

        return buttons