import csv
from collections import OrderedDict
from django.db.models import FieldDoesNotExist
from django.http import StreamingHttpResponse


def prep_field(obj, field):
    """
    (for export_as_csv action)
    Returns the field as a unicode string. If the field is a callable, it
    attempts to call it first, without arguments.
    """
    if '__' in field:
        bits = field.split('__')
        field = bits.pop()

        for bit in bits:
            obj = getattr(obj, bit, None)

            if obj is None:
                return ""

    attr = getattr(obj, field)
    output = attr() if callable(attr) else attr
    return unicode(output).encode('utf-8') if output is not None else ""



class Echo(object):
    """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


def export_as_csv(description, fields=None, exclude=None, use_verbose_names=True, include_header=True):
    def _export_as_csv(modeladmin, request, queryset):
        """
        Usage:

            class ExampleModelAdmin(admin.ModelAdmin):
                list_display = ('field1', 'field2', 'field3',)
                actions = [
                    export_as_csv(
                        'export to csv',
                        fields=['field1', '('foreign_key1__foreign_key2__name', 'label2'),],
                        include_header=True
                    )
                ]
        """
        opts = modeladmin.model._meta

        def field_name(field):
            if use_verbose_names:
                return unicode(field.verbose_name).capitalize()
            else:
                return field.name

        # field_names is a map of {field lookup path: field label}
        if exclude:
            field_names = OrderedDict(
                (f.name, field_name(f)) for f in opts.fields if f not in exclude
            )
        elif fields:
            field_names = OrderedDict()
            for spec in fields:
                if isinstance(spec, (list, tuple)):
                    field_names[spec[0]] = spec[1]
                else:
                    try:
                        f, _, _, _ = opts.get_field_by_name(spec)
                    except FieldDoesNotExist:
                        field_names[spec] = spec
                    else:
                        field_names[spec] = field_name(f)
        else:
            field_names = OrderedDict(
                (f.name, field_name(f)) for f in opts.fields
            )

        pseudo_buffer = Echo()
        writer = csv.writer(pseudo_buffer)

        def content_iterator():
            if include_header:
                yield [i.encode('utf8') for i in field_names.values()]
            for obj in queryset.iterator():
                yield [prep_field(obj, field) for field in field_names.keys()]

        response = StreamingHttpResponse(
            (writer.writerow(line) for line in content_iterator()),
            content_type='text/csv'
        )
        response['Content-Disposition'] = 'attachment; filename=%s.csv' % (unicode(opts).replace('.', '_'))
        return response

    _export_as_csv.short_description = description
    return _export_as_csv