from __future__ import absolute_import, unicode_literals

import csv

import django
import six
from django.http.response import HttpResponse
from openpyxl import Workbook
from openpyxl.writer.excel import save_virtual_workbook
from openpyxl.writer.write_only import WriteOnlyCell

if django.VERSION >= (1, 9):
    from django.db.models.query import QuerySet
else:
    from django.db.models.query import QuerySet, ValuesQuerySet


ROW_LIMIT = 1048576
COL_LIMIT = 16384


class ExcelResponse(HttpResponse):
    """
    This class provides an HTTP Response in the form of an Excel spreadsheet, or CSV file.
    """

    def __init__(self, data, output_filename='excel_data', worksheet_name=None, force_csv=False, header_font=None,
                 data_font=None, *args, **kwargs):
        # We do not initialize this with streaming_content, as that gets generated when needed
        self.output_filename = output_filename
        self.worksheet_name = worksheet_name or 'Sheet 1'
        self.header_font = header_font
        self.data_font = data_font
        self.force_csv = force_csv
        super(ExcelResponse, self).__init__(data, *args, **kwargs)

    @property
    def content(self):
        return b''.join(self._container)

    @content.setter
    def content(self, value):
        workbook = None
        if isinstance(value, list):
            workbook = self._serialize_list(value)
        elif isinstance(value, QuerySet):
            workbook = self._serialize_queryset(value)
        if django.VERSION < (1, 9):
            if isinstance(value, ValuesQuerySet):
                workbook = self._serialize_values_queryset(value)
        if workbook is None:
            raise ValueError('ExcelResponse accepts the following data types: list, dict, QuerySet, ValuesQuerySet')

        if self.force_csv:
            self['Content-Type'] = 'text/csv; charset=utf8'
            self['Content-Disposition'] = 'attachment;filename={}.csv'.format(self.output_filename)
            workbook.seek(0)
            workbook = self.make_bytes(workbook.getvalue())
        else:
            self['Content-Type'] = 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
            self['Content-Disposition'] = 'attachment; filename={}.xlsx'.format(self.output_filename)
            workbook = save_virtual_workbook(workbook)
        self._container = [self.make_bytes(workbook)]

    def _serialize_list(self, data):
        workbook = None
        if isinstance(data[0], dict):  # If we're dealing with a list of dicationaries, generate the headers
            headers = [key for key in data[0]]
        else:
            headers = data[0]
        if len(data) > ROW_LIMIT or len(headers) > COL_LIMIT or self.force_csv:
            self.force_csv = True
            workbook = six.StringIO()
            csvwriter = csv.writer(workbook, dialect='excel')
            append = getattr(csvwriter, 'writerow')
            write_header = append
        else:
            workbook = Workbook(write_only=True)
            workbook.guess_types = True
            worksheet = workbook.create_sheet(title=self.worksheet_name)

            # Define custom functions for appending so that we can handle any formatting
            def append(data):
                return self._append_excel_row(worksheet, data, header=False)

            def write_header(data):
                return self._append_excel_row(worksheet, data, header=True)

        if isinstance(data[0], dict):
            append(headers)
        for index, row in enumerate(data):
            if isinstance(row, dict):
                write_header([row.get(col, None) for col in headers])
            else:
                if index > 0:
                    append(row)
                else:
                    write_header(row)
        return workbook

    def _serialize_queryset(self, data):
        if isinstance(data[0], dict):  # .values() returns a list of dicts
            return self._serialize_list(list(data))
        else:
            return self._serialize_list(list(data.values()))

    def _serialize_values_queryset(self, data):
        return self._serialize_list(list(data))

    def _append_excel_row(self, worksheet, data, header=False):
        if header:
            font = self.header_font
        else:
            font = self.data_font

        if not font:
            row = data
        else:
            row = []
            for cell in data:
                cell = WriteOnlyCell(worksheet, cell)
                cell.font = font
                row.append(cell)

        worksheet.append(row)
