# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

from collections import namedtuple, OrderedDict
import binascii
import json
import os
import random
import tempfile

import numpy as np

from .util import frombytes, tobytes, random_bytes, random_utf8


def metadata_key_values(pairs):
    return [{'key': k, 'value': v} for k, v in pairs]


class Field(object):

    def __init__(self, name, *, nullable=True, metadata=None):
        self.name = name
        self.nullable = nullable
        self.metadata = metadata or []

    def get_json(self):
        entries = [
            ('name', self.name),
            ('type', self._get_type()),
            ('nullable', self.nullable),
            ('children', self._get_children()),
        ]

        dct = self._get_dictionary()
        if dct:
            entries.append(('dictionary', dct))

        if self.metadata is not None and len(self.metadata) > 0:
            entries.append(('metadata', metadata_key_values(self.metadata)))

        return OrderedDict(entries)

    def _get_dictionary(self):
        return None

    def _make_is_valid(self, size, null_probability=0.4):
        if self.nullable:
            return (np.random.random_sample(size) > null_probability
                    ).astype(np.int8)
        else:
            return np.ones(size, dtype=np.int8)


class Column(object):

    def __init__(self, name, count):
        self.name = name
        self.count = count

    def __len__(self):
        return self.count

    def _get_children(self):
        return []

    def _get_buffers(self):
        return []

    def get_json(self):
        entries = [
            ('name', self.name),
            ('count', self.count)
        ]

        buffers = self._get_buffers()
        entries.extend(buffers)

        children = self._get_children()
        if len(children) > 0:
            entries.append(('children', children))

        return OrderedDict(entries)


class PrimitiveField(Field):

    def _get_children(self):
        return []


class PrimitiveColumn(Column):

    def __init__(self, name, count, is_valid, values):
        super().__init__(name, count)
        self.is_valid = is_valid
        self.values = values

    def _encode_value(self, x):
        return x

    def _get_buffers(self):
        return [
            ('VALIDITY', [int(v) for v in self.is_valid]),
            ('DATA', list([self._encode_value(x) for x in self.values]))
        ]


class NullColumn(Column):
    # This subclass is for readability only
    pass


class NullField(PrimitiveField):

    def __init__(self, name, metadata=None):
        super().__init__(name, nullable=True,
                         metadata=metadata)

    def _get_type(self):
        return OrderedDict([('name', 'null')])

    def generate_column(self, size, name=None):
        return NullColumn(name or self.name, size)


TEST_INT_MAX = 2 ** 31 - 1
TEST_INT_MIN = ~TEST_INT_MAX


class IntegerField(PrimitiveField):

    def __init__(self, name, is_signed, bit_width, *, nullable=True,
                 metadata=None,
                 min_value=TEST_INT_MIN,
                 max_value=TEST_INT_MAX):
        super().__init__(name, nullable=nullable,
                         metadata=metadata)
        self.is_signed = is_signed
        self.bit_width = bit_width
        self.min_value = min_value
        self.max_value = max_value

    def _get_generated_data_bounds(self):
        if self.is_signed:
            signed_iinfo = np.iinfo('int' + str(self.bit_width))
            min_value, max_value = signed_iinfo.min, signed_iinfo.max
        else:
            unsigned_iinfo = np.iinfo('uint' + str(self.bit_width))
            min_value, max_value = 0, unsigned_iinfo.max

        lower_bound = max(min_value, self.min_value)
        upper_bound = min(max_value, self.max_value)
        return lower_bound, upper_bound

    def _get_type(self):
        return OrderedDict([
            ('name', 'int'),
            ('isSigned', self.is_signed),
            ('bitWidth', self.bit_width)
        ])

    def generate_column(self, size, name=None):
        lower_bound, upper_bound = self._get_generated_data_bounds()
        return self.generate_range(size, lower_bound, upper_bound,
                                   name=name, include_extremes=True)

    def generate_range(self, size, lower, upper, name=None,
                       include_extremes=False):
        values = np.random.randint(lower, upper, size=size, dtype=np.int64)
        if include_extremes and size >= 2:
            values[:2] = [lower, upper]
        values = list(map(int if self.bit_width < 64 else str, values))

        is_valid = self._make_is_valid(size)

        if name is None:
            name = self.name
        return PrimitiveColumn(name, size, is_valid, values)


class DateField(IntegerField):

    DAY = 0
    MILLISECOND = 1

    # 1/1/1 to 12/31/9999
    _ranges = {
        DAY: [-719162, 2932896],
        MILLISECOND: [-62135596800000, 253402214400000]
    }

    def __init__(self, name, unit, *, nullable=True, metadata=None):
        bit_width = 32 if unit == self.DAY else 64

        min_value, max_value = self._ranges[unit]
        super().__init__(
            name, True, bit_width,
            nullable=nullable, metadata=metadata,
            min_value=min_value, max_value=max_value
        )
        self.unit = unit

    def _get_type(self):
        return OrderedDict([
            ('name', 'date'),
            ('unit', 'DAY' if self.unit == self.DAY else 'MILLISECOND')
        ])


TIMEUNIT_NAMES = {
    's': 'SECOND',
    'ms': 'MILLISECOND',
    'us': 'MICROSECOND',
    'ns': 'NANOSECOND'
}


class TimeField(IntegerField):

    BIT_WIDTHS = {
        's': 32,
        'ms': 32,
        'us': 64,
        'ns': 64
    }

    _ranges = {
        's': [0, 86400],
        'ms': [0, 86400000],
        'us': [0, 86400000000],
        'ns': [0, 86400000000000]
    }

    def __init__(self, name, unit='s', *, nullable=True,
                 metadata=None):
        min_val, max_val = self._ranges[unit]
        super().__init__(name, True, self.BIT_WIDTHS[unit],
                         nullable=nullable, metadata=metadata,
                         min_value=min_val, max_value=max_val)
        self.unit = unit

    def _get_type(self):
        return OrderedDict([
            ('name', 'time'),
            ('unit', TIMEUNIT_NAMES[self.unit]),
            ('bitWidth', self.bit_width)
        ])


class TimestampField(IntegerField):

    # 1/1/1 to 12/31/9999
    _ranges = {
        's': [-62135596800, 253402214400],
        'ms': [-62135596800000, 253402214400000],
        'us': [-62135596800000000, 253402214400000000],

        # Physical range for int64, ~584 years and change
        'ns': [np.iinfo('int64').min, np.iinfo('int64').max]
    }

    def __init__(self, name, unit='s', tz=None, *, nullable=True,
                 metadata=None):
        min_val, max_val = self._ranges[unit]
        super().__init__(name, True, 64,
                         nullable=nullable,
                         metadata=metadata,
                         min_value=min_val,
                         max_value=max_val)
        self.unit = unit
        self.tz = tz

    def _get_type(self):
        fields = [
            ('name', 'timestamp'),
            ('unit', TIMEUNIT_NAMES[self.unit])
        ]

        if self.tz is not None:
            fields.append(('timezone', self.tz))

        return OrderedDict(fields)


class DurationIntervalField(IntegerField):

    def __init__(self, name, unit='s', *, nullable=True,
                 metadata=None):
        min_val, max_val = np.iinfo('int64').min, np.iinfo('int64').max,
        super().__init__(
            name, True, 64,
            nullable=nullable, metadata=metadata,
            min_value=min_val, max_value=max_val)
        self.unit = unit

    def _get_type(self):
        fields = [
            ('name', 'duration'),
            ('unit', TIMEUNIT_NAMES[self.unit])
        ]

        return OrderedDict(fields)


class YearMonthIntervalField(IntegerField):
    def __init__(self, name, *, nullable=True, metadata=None):
        min_val, max_val = [-10000*12, 10000*12]  # +/- 10000 years.
        super().__init__(
            name, True, 32,
            nullable=nullable, metadata=metadata,
            min_value=min_val, max_value=max_val)

    def _get_type(self):
        fields = [
            ('name', 'interval'),
            ('unit', 'YEAR_MONTH'),
        ]

        return OrderedDict(fields)


class DayTimeIntervalField(PrimitiveField):
    def __init__(self, name, *, nullable=True, metadata=None):
        super().__init__(name,
                         nullable=True,
                         metadata=metadata)

    @property
    def numpy_type(self):
        return object

    def _get_type(self):

        return OrderedDict([
            ('name', 'interval'),
            ('unit', 'DAY_TIME'),
        ])

    def generate_column(self, size, name=None):
        min_day_value, max_day_value = -10000*366, 10000*366
        values = [{'days': random.randint(min_day_value, max_day_value),
                   'milliseconds': random.randint(-86400000, +86400000)}
                  for _ in range(size)]

        is_valid = self._make_is_valid(size)
        if name is None:
            name = self.name
        return PrimitiveColumn(name, size, is_valid, values)


class FloatingPointField(PrimitiveField):

    def __init__(self, name, bit_width, *, nullable=True,
                 metadata=None):
        super().__init__(name,
                         nullable=nullable,
                         metadata=metadata)

        self.bit_width = bit_width
        self.precision = {
            16: 'HALF',
            32: 'SINGLE',
            64: 'DOUBLE'
        }[self.bit_width]

    @property
    def numpy_type(self):
        return 'float' + str(self.bit_width)

    def _get_type(self):
        return OrderedDict([
            ('name', 'floatingpoint'),
            ('precision', self.precision)
        ])

    def generate_column(self, size, name=None):
        values = np.random.randn(size) * 1000
        values = np.round(values, 3)

        is_valid = self._make_is_valid(size)
        if name is None:
            name = self.name
        return PrimitiveColumn(name, size, is_valid, values)


DECIMAL_PRECISION_TO_VALUE = {
    key: (1 << (8 * i - 1)) - 1 for i, key in enumerate(
        [1, 3, 5, 7, 10, 12, 15, 17, 19, 22, 24, 27, 29, 32, 34, 36,
         40, 42, 44, 50, 60, 70],
        start=1,
    )
}


def decimal_range_from_precision(precision):
    assert 1 <= precision <= 76
    try:
        max_value = DECIMAL_PRECISION_TO_VALUE[precision]
    except KeyError:
        return decimal_range_from_precision(precision - 1)
    else:
        return ~max_value, max_value


class DecimalField(PrimitiveField):
    def __init__(self, name, precision, scale, bit_width, *,
                 nullable=True, metadata=None):
        super().__init__(name, nullable=True,
                         metadata=metadata)
        self.precision = precision
        self.scale = scale
        self.bit_width = bit_width

    @property
    def numpy_type(self):
        return object

    def _get_type(self):
        return OrderedDict([
            ('name', 'decimal'),
            ('precision', self.precision),
            ('scale', self.scale),
            ('bitWidth', self.bit_width),
        ])

    def generate_column(self, size, name=None):
        min_value, max_value = decimal_range_from_precision(self.precision)
        values = [random.randint(min_value, max_value) for _ in range(size)]

        is_valid = self._make_is_valid(size)
        if name is None:
            name = self.name
        return DecimalColumn(name, size, is_valid, values, self.bit_width)


class DecimalColumn(PrimitiveColumn):

    def __init__(self, name, count, is_valid, values, bit_width):
        super().__init__(name, count, is_valid, values)
        self.bit_width = bit_width

    def _encode_value(self, x):
        return str(x)


class BooleanField(PrimitiveField):
    bit_width = 1

    def _get_type(self):
        return OrderedDict([('name', 'bool')])

    @property
    def numpy_type(self):
        return 'bool'

    def generate_column(self, size, name=None):
        values = list(map(bool, np.random.randint(0, 2, size=size)))
        is_valid = self._make_is_valid(size)
        if name is None:
            name = self.name
        return PrimitiveColumn(name, size, is_valid, values)


class FixedSizeBinaryField(PrimitiveField):

    def __init__(self, name, byte_width, *, nullable=True,
                 metadata=None):
        super().__init__(name, nullable=nullable,
                         metadata=metadata)
        self.byte_width = byte_width

    @property
    def numpy_type(self):
        return object

    @property
    def column_class(self):
        return FixedSizeBinaryColumn

    def _get_type(self):
        return OrderedDict([('name', 'fixedsizebinary'),
                            ('byteWidth', self.byte_width)])

    def generate_column(self, size, name=None):
        is_valid = self._make_is_valid(size)
        values = []

        for i in range(size):
            values.append(random_bytes(self.byte_width))

        if name is None:
            name = self.name
        return self.column_class(name, size, is_valid, values)


class BinaryField(PrimitiveField):

    @property
    def numpy_type(self):
        return object

    @property
    def column_class(self):
        return BinaryColumn

    def _get_type(self):
        return OrderedDict([('name', 'binary')])

    def _random_sizes(self, size):
        return np.random.exponential(scale=4, size=size).astype(np.int32)

    def generate_column(self, size, name=None):
        is_valid = self._make_is_valid(size)
        values = []

        sizes = self._random_sizes(size)

        for i, nbytes in enumerate(sizes):
            if is_valid[i]:
                values.append(random_bytes(nbytes))
            else:
                values.append(b"")

        if name is None:
            name = self.name
        return self.column_class(name, size, is_valid, values)


class StringField(BinaryField):

    @property
    def column_class(self):
        return StringColumn

    def _get_type(self):
        return OrderedDict([('name', 'utf8')])

    def generate_column(self, size, name=None):
        K = 7
        is_valid = self._make_is_valid(size)
        values = []

        for i in range(size):
            if is_valid[i]:
                values.append(tobytes(random_utf8(K)))
            else:
                values.append(b"")

        if name is None:
            name = self.name
        return self.column_class(name, size, is_valid, values)


class LargeBinaryField(BinaryField):

    @property
    def column_class(self):
        return LargeBinaryColumn

    def _get_type(self):
        return OrderedDict([('name', 'largebinary')])


class LargeStringField(StringField):

    @property
    def column_class(self):
        return LargeStringColumn

    def _get_type(self):
        return OrderedDict([('name', 'largeutf8')])


class Schema(object):

    def __init__(self, fields, metadata=None):
        self.fields = fields
        self.metadata = metadata

    def get_json(self):
        entries = [
            ('fields', [field.get_json() for field in self.fields])
        ]

        if self.metadata is not None and len(self.metadata) > 0:
            entries.append(('metadata', metadata_key_values(self.metadata)))

        return OrderedDict(entries)


class _NarrowOffsetsMixin:

    def _encode_offsets(self, offsets):
        return list(map(int, offsets))


class _LargeOffsetsMixin:

    def _encode_offsets(self, offsets):
        # 64-bit offsets have to be represented as strings to roundtrip
        # through JSON.
        return list(map(str, offsets))


class _BaseBinaryColumn(PrimitiveColumn):

    def _encode_value(self, x):
        return frombytes(binascii.hexlify(x).upper())

    def _get_buffers(self):
        offset = 0
        offsets = [0]

        data = []
        for i, v in enumerate(self.values):
            if self.is_valid[i]:
                offset += len(v)
            else:
                v = b""

            offsets.append(offset)
            data.append(self._encode_value(v))

        return [
            ('VALIDITY', [int(x) for x in self.is_valid]),
            ('OFFSET', self._encode_offsets(offsets)),
            ('DATA', data)
        ]


class _BaseStringColumn(_BaseBinaryColumn):

    def _encode_value(self, x):
        return frombytes(x)


class BinaryColumn(_BaseBinaryColumn, _NarrowOffsetsMixin):
    pass


class StringColumn(_BaseStringColumn, _NarrowOffsetsMixin):
    pass


class LargeBinaryColumn(_BaseBinaryColumn, _LargeOffsetsMixin):
    pass


class LargeStringColumn(_BaseStringColumn, _LargeOffsetsMixin):
    pass


class FixedSizeBinaryColumn(PrimitiveColumn):

    def _encode_value(self, x):
        return frombytes(binascii.hexlify(x).upper())

    def _get_buffers(self):
        data = []
        for i, v in enumerate(self.values):
            data.append(self._encode_value(v))

        return [
            ('VALIDITY', [int(x) for x in self.is_valid]),
            ('DATA', data)
        ]


class ListField(Field):

    def __init__(self, name, value_field, *, nullable=True,
                 metadata=None):
        super().__init__(name, nullable=nullable,
                         metadata=metadata)
        self.value_field = value_field

    @property
    def column_class(self):
        return ListColumn

    def _get_type(self):
        return OrderedDict([
            ('name', 'list')
        ])

    def _get_children(self):
        return [self.value_field.get_json()]

    def generate_column(self, size, name=None):
        MAX_LIST_SIZE = 4

        is_valid = self._make_is_valid(size)
        list_sizes = np.random.randint(0, MAX_LIST_SIZE + 1, size=size)
        offsets = [0]

        offset = 0
        for i in range(size):
            if is_valid[i]:
                offset += int(list_sizes[i])
            offsets.append(offset)

        # The offset now is the total number of elements in the child array
        values = self.value_field.generate_column(offset)

        if name is None:
            name = self.name
        return self.column_class(name, size, is_valid, offsets, values)


class LargeListField(ListField):

    @property
    def column_class(self):
        return LargeListColumn

    def _get_type(self):
        return OrderedDict([
            ('name', 'largelist')
        ])


class _BaseListColumn(Column):

    def __init__(self, name, count, is_valid, offsets, values):
        super().__init__(name, count)
        self.is_valid = is_valid
        self.offsets = offsets
        self.values = values

    def _get_buffers(self):
        return [
            ('VALIDITY', [int(v) for v in self.is_valid]),
            ('OFFSET', self._encode_offsets(self.offsets))
        ]

    def _get_children(self):
        return [self.values.get_json()]


class ListColumn(_BaseListColumn, _NarrowOffsetsMixin):
    pass


class LargeListColumn(_BaseListColumn, _LargeOffsetsMixin):
    pass


class MapField(Field):

    def __init__(self, name, key_field, item_field, *, nullable=True,
                 metadata=None, keys_sorted=False, entries_name='entries'):
        super().__init__(name, nullable=nullable,
                         metadata=metadata)

        assert not key_field.nullable
        self.key_field = key_field
        self.item_field = item_field
        self.pair_field = StructField(entries_name, [key_field, item_field],
                                      nullable=False)
        self.keys_sorted = keys_sorted

    def _get_type(self):
        return OrderedDict([
            ('name', 'map'),
            ('keysSorted', self.keys_sorted)
        ])

    def _get_children(self):
        return [self.pair_field.get_json()]

    def generate_column(self, size, name=None):
        MAX_MAP_SIZE = 4

        is_valid = self._make_is_valid(size)
        map_sizes = np.random.randint(0, MAX_MAP_SIZE + 1, size=size)
        offsets = [0]

        offset = 0
        for i in range(size):
            if is_valid[i]:
                offset += int(map_sizes[i])
            offsets.append(offset)

        # The offset now is the total number of elements in the child array
        pairs = self.pair_field.generate_column(offset)
        if name is None:
            name = self.name

        return MapColumn(name, size, is_valid, offsets, pairs)


class MapColumn(Column):

    def __init__(self, name, count, is_valid, offsets, pairs):
        super().__init__(name, count)
        self.is_valid = is_valid
        self.offsets = offsets
        self.pairs = pairs

    def _get_buffers(self):
        return [
            ('VALIDITY', [int(v) for v in self.is_valid]),
            ('OFFSET', list(self.offsets))
        ]

    def _get_children(self):
        return [self.pairs.get_json()]


class FixedSizeListField(Field):

    def __init__(self, name, value_field, list_size, *, nullable=True,
                 metadata=None):
        super().__init__(name, nullable=nullable,
                         metadata=metadata)
        self.value_field = value_field
        self.list_size = list_size

    def _get_type(self):
        return OrderedDict([
            ('name', 'fixedsizelist'),
            ('listSize', self.list_size)
        ])

    def _get_children(self):
        return [self.value_field.get_json()]

    def generate_column(self, size, name=None):
        is_valid = self._make_is_valid(size)
        values = self.value_field.generate_column(size * self.list_size)

        if name is None:
            name = self.name
        return FixedSizeListColumn(name, size, is_valid, values)


class FixedSizeListColumn(Column):

    def __init__(self, name, count, is_valid, values):
        super().__init__(name, count)
        self.is_valid = is_valid
        self.values = values

    def _get_buffers(self):
        return [
            ('VALIDITY', [int(v) for v in self.is_valid])
        ]

    def _get_children(self):
        return [self.values.get_json()]


class StructField(Field):

    def __init__(self, name, fields, *, nullable=True,
                 metadata=None):
        super().__init__(name, nullable=nullable,
                         metadata=metadata)
        self.fields = fields

    def _get_type(self):
        return OrderedDict([
            ('name', 'struct')
        ])

    def _get_children(self):
        return [field.get_json() for field in self.fields]

    def generate_column(self, size, name=None):
        is_valid = self._make_is_valid(size)

        field_values = [field.generate_column(size) for field in self.fields]
        if name is None:
            name = self.name
        return StructColumn(name, size, is_valid, field_values)


class _BaseUnionField(Field):

    def __init__(self, name, fields, type_ids=None, *, nullable=True,
                 metadata=None):
        super().__init__(name, nullable=nullable, metadata=metadata)
        if type_ids is None:
            type_ids = list(range(fields))
        else:
            assert len(fields) == len(type_ids)
        self.fields = fields
        self.type_ids = type_ids
        assert all(x >= 0 for x in self.type_ids)

    def _get_type(self):
        return OrderedDict([
            ('name', 'union'),
            ('mode', self.mode),
            ('typeIds', self.type_ids),
        ])

    def _get_children(self):
        return [field.get_json() for field in self.fields]

    def _make_type_ids(self, size):
        return np.random.choice(self.type_ids, size)


class SparseUnionField(_BaseUnionField):
    mode = 'SPARSE'

    def generate_column(self, size, name=None):
        array_type_ids = self._make_type_ids(size)
        field_values = [field.generate_column(size) for field in self.fields]

        if name is None:
            name = self.name
        return SparseUnionColumn(name, size, array_type_ids, field_values)


class DenseUnionField(_BaseUnionField):
    mode = 'DENSE'

    def generate_column(self, size, name=None):
        # Reverse mapping {logical type id => physical child id}
        child_ids = [None] * (max(self.type_ids) + 1)
        for i, type_id in enumerate(self.type_ids):
            child_ids[type_id] = i

        array_type_ids = self._make_type_ids(size)
        offsets = []
        child_sizes = [0] * len(self.fields)

        for i in range(size):
            child_id = child_ids[array_type_ids[i]]
            offset = child_sizes[child_id]
            offsets.append(offset)
            child_sizes[child_id] = offset + 1

        field_values = [
            field.generate_column(child_size)
            for field, child_size in zip(self.fields, child_sizes)]

        if name is None:
            name = self.name
        return DenseUnionColumn(name, size, array_type_ids, offsets,
                                field_values)


class Dictionary(object):

    def __init__(self, id_, field, size, name=None, ordered=False):
        self.id_ = id_
        self.field = field
        self.values = field.generate_column(size=size, name=name)
        self.ordered = ordered

    def __len__(self):
        return len(self.values)

    def get_json(self):
        dummy_batch = RecordBatch(len(self.values), [self.values])
        return OrderedDict([
            ('id', self.id_),
            ('data', dummy_batch.get_json())
        ])


class DictionaryField(Field):

    def __init__(self, name, index_field, dictionary, *, nullable=True,
                 metadata=None):
        super().__init__(name, nullable=nullable,
                         metadata=metadata)
        assert index_field.name == ''
        assert isinstance(index_field, IntegerField)
        assert isinstance(dictionary, Dictionary)

        self.index_field = index_field
        self.dictionary = dictionary

    def _get_type(self):
        return self.dictionary.field._get_type()

    def _get_children(self):
        return self.dictionary.field._get_children()

    def _get_dictionary(self):
        return OrderedDict([
            ('id', self.dictionary.id_),
            ('indexType', self.index_field._get_type()),
            ('isOrdered', self.dictionary.ordered)
        ])

    def generate_column(self, size, name=None):
        if name is None:
            name = self.name
        return self.index_field.generate_range(size, 0, len(self.dictionary),
                                               name=name)


ExtensionType = namedtuple(
    'ExtensionType', ['extension_name', 'serialized', 'storage_field'])


class ExtensionField(Field):

    def __init__(self, name, extension_type, *, nullable=True, metadata=None):
        metadata = (metadata or []) + [
            ('ARROW:extension:name', extension_type.extension_name),
            ('ARROW:extension:metadata', extension_type.serialized),
        ]
        super().__init__(name, nullable=nullable, metadata=metadata)
        self.extension_type = extension_type

    def _get_type(self):
        return self.extension_type.storage_field._get_type()

    def _get_children(self):
        return self.extension_type.storage_field._get_children()

    def _get_dictionary(self):
        return self.extension_type.storage_field._get_dictionary()

    def generate_column(self, size, name=None):
        if name is None:
            name = self.name
        return self.extension_type.storage_field.generate_column(size, name)


class StructColumn(Column):

    def __init__(self, name, count, is_valid, field_values):
        super().__init__(name, count)
        self.is_valid = is_valid
        self.field_values = field_values

    def _get_buffers(self):
        return [
            ('VALIDITY', [int(v) for v in self.is_valid])
        ]

    def _get_children(self):
        return [field.get_json() for field in self.field_values]


class SparseUnionColumn(Column):

    def __init__(self, name, count, type_ids, field_values):
        super().__init__(name, count)
        self.type_ids = type_ids
        self.field_values = field_values

    def _get_buffers(self):
        return [
            ('TYPE_ID', [int(v) for v in self.type_ids])
        ]

    def _get_children(self):
        return [field.get_json() for field in self.field_values]


class DenseUnionColumn(Column):

    def __init__(self, name, count, type_ids, offsets, field_values):
        super().__init__(name, count)
        self.type_ids = type_ids
        self.offsets = offsets
        self.field_values = field_values

    def _get_buffers(self):
        return [
            ('TYPE_ID', [int(v) for v in self.type_ids]),
            ('OFFSET', [int(v) for v in self.offsets]),
        ]

    def _get_children(self):
        return [field.get_json() for field in self.field_values]


class RecordBatch(object):

    def __init__(self, count, columns):
        self.count = count
        self.columns = columns

    def get_json(self):
        return OrderedDict([
            ('count', self.count),
            ('columns', [col.get_json() for col in self.columns])
        ])


class File(object):

    def __init__(self, name, schema, batches, dictionaries=None,
                 skip=None, path=None):
        self.name = name
        self.schema = schema
        self.dictionaries = dictionaries or []
        self.batches = batches
        self.skip = set()
        self.path = path
        if skip:
            self.skip.update(skip)

    def get_json(self):
        entries = [
            ('schema', self.schema.get_json())
        ]

        if len(self.dictionaries) > 0:
            entries.append(('dictionaries',
                            [dictionary.get_json()
                             for dictionary in self.dictionaries]))

        entries.append(('batches', [batch.get_json()
                                    for batch in self.batches]))
        return OrderedDict(entries)

    def write(self, path):
        with open(path, 'wb') as f:
            f.write(json.dumps(self.get_json(), indent=2).encode('utf-8'))
        self.path = path

    def skip_category(self, category):
        """Skip this test for the given category.

        Category should be SKIP_ARROW or SKIP_FLIGHT.
        """
        self.skip.add(category)
        return self


def get_field(name, type_, **kwargs):
    if type_ == 'binary':
        return BinaryField(name, **kwargs)
    elif type_ == 'utf8':
        return StringField(name, **kwargs)
    elif type_ == 'largebinary':
        return LargeBinaryField(name, **kwargs)
    elif type_ == 'largeutf8':
        return LargeStringField(name, **kwargs)
    elif type_.startswith('fixedsizebinary_'):
        byte_width = int(type_.split('_')[1])
        return FixedSizeBinaryField(name, byte_width=byte_width, **kwargs)

    dtype = np.dtype(type_)

    if dtype.kind in ('i', 'u'):
        signed = dtype.kind == 'i'
        bit_width = dtype.itemsize * 8
        return IntegerField(name, signed, bit_width, **kwargs)
    elif dtype.kind == 'f':
        bit_width = dtype.itemsize * 8
        return FloatingPointField(name, bit_width, **kwargs)
    elif dtype.kind == 'b':
        return BooleanField(name, **kwargs)
    else:
        raise TypeError(dtype)


def _generate_file(name, fields, batch_sizes, dictionaries=None, skip=None,
                   metadata=None):
    schema = Schema(fields, metadata=metadata)
    batches = []
    for size in batch_sizes:
        columns = []
        for field in fields:
            col = field.generate_column(size)
            columns.append(col)

        batches.append(RecordBatch(size, columns))

    return File(name, schema, batches, dictionaries, skip=skip)


def generate_custom_metadata_case():
    def meta(items):
        # Generate a simple block of metadata where each value is '{}'.
        # Keys are delimited by whitespace in `items`.
        return [(k, '{}') for k in items.split()]

    fields = [
        get_field('sort_of_pandas', 'int8', metadata=meta('pandas')),

        get_field('lots_of_meta', 'int8', metadata=meta('a b c d .. w x y z')),

        get_field(
            'unregistered_extension', 'int8',
            metadata=[
                ('ARROW:extension:name', '!nonexistent'),
                ('ARROW:extension:metadata', ''),
                ('ARROW:integration:allow_unregistered_extension', 'true'),
            ]),

        ListField('list_with_odd_values',
                  get_field('item', 'int32', metadata=meta('odd_values'))),
    ]

    batch_sizes = [1]
    return _generate_file('custom_metadata', fields, batch_sizes,
                          metadata=meta('schema_custom_0 schema_custom_1'))


def generate_duplicate_fieldnames_case():
    fields = [
        get_field('ints', 'int8'),
        get_field('ints', 'int32'),

        StructField('struct', [get_field('', 'int32'), get_field('', 'utf8')]),
    ]

    batch_sizes = [1]
    return _generate_file('duplicate_fieldnames', fields, batch_sizes)


def generate_primitive_case(batch_sizes, name='primitive'):
    types = ['bool', 'int8', 'int16', 'int32', 'int64',
             'uint8', 'uint16', 'uint32', 'uint64',
             'float32', 'float64', 'binary', 'utf8',
             'fixedsizebinary_19', 'fixedsizebinary_120']

    fields = []

    for type_ in types:
        fields.append(get_field(type_ + "_nullable", type_, nullable=True))
        fields.append(get_field(type_ + "_nonnullable", type_, nullable=False))

    return _generate_file(name, fields, batch_sizes)


def generate_primitive_large_offsets_case(batch_sizes):
    types = ['largebinary', 'largeutf8']

    fields = []

    for type_ in types:
        fields.append(get_field(type_ + "_nullable", type_, nullable=True))
        fields.append(get_field(type_ + "_nonnullable", type_, nullable=False))

    return _generate_file('primitive_large_offsets', fields, batch_sizes)


def generate_null_case(batch_sizes):
    # Interleave null with non-null types to ensure the appropriate number of
    # buffers (0) is read and written
    fields = [
        NullField(name='f0'),
        get_field('f1', 'int32'),
        NullField(name='f2'),
        get_field('f3', 'float64'),
        NullField(name='f4')
    ]
    return _generate_file('null', fields, batch_sizes)


def generate_null_trivial_case(batch_sizes):
    # Generate a case with no buffers
    fields = [
        NullField(name='f0'),
    ]
    return _generate_file('null_trivial', fields, batch_sizes)


def generate_decimal128_case():
    fields = [
        DecimalField(name='f{}'.format(i), precision=precision, scale=2,
                     bit_width=128)
        for i, precision in enumerate(range(3, 39))
    ]

    possible_batch_sizes = 7, 10
    batch_sizes = [possible_batch_sizes[i % 2] for i in range(len(fields))]
    # 'decimal' is the original name for the test, and it must match
    # provide "gold" files that test backwards compatibility, so they
    # can be appropriately skipped.
    return _generate_file('decimal', fields, batch_sizes)


def generate_decimal256_case():
    fields = [
        DecimalField(name='f{}'.format(i), precision=precision, scale=5,
                     bit_width=256)
        for i, precision in enumerate(range(37, 70))
    ]

    possible_batch_sizes = 7, 10
    batch_sizes = [possible_batch_sizes[i % 2] for i in range(len(fields))]
    return _generate_file('decimal256', fields, batch_sizes)


def generate_datetime_case():
    fields = [
        DateField('f0', DateField.DAY),
        DateField('f1', DateField.MILLISECOND),
        TimeField('f2', 's'),
        TimeField('f3', 'ms'),
        TimeField('f4', 'us'),
        TimeField('f5', 'ns'),
        TimestampField('f6', 's'),
        TimestampField('f7', 'ms'),
        TimestampField('f8', 'us'),
        TimestampField('f9', 'ns'),
        TimestampField('f10', 'ms', tz=None),
        TimestampField('f11', 's', tz='UTC'),
        TimestampField('f12', 'ms', tz='US/Eastern'),
        TimestampField('f13', 'us', tz='Europe/Paris'),
        TimestampField('f14', 'ns', tz='US/Pacific'),
    ]

    batch_sizes = [7, 10]
    return _generate_file("datetime", fields, batch_sizes)


def generate_interval_case():
    fields = [
        DurationIntervalField('f1', 's'),
        DurationIntervalField('f2', 'ms'),
        DurationIntervalField('f3', 'us'),
        DurationIntervalField('f4', 'ns'),
        YearMonthIntervalField('f5'),
        DayTimeIntervalField('f6'),
    ]

    batch_sizes = [7, 10]
    return _generate_file("interval", fields, batch_sizes)


def generate_map_case():
    fields = [
        MapField('map_nullable', get_field('key', 'utf8', nullable=False),
                 get_field('value', 'int32')),
    ]

    batch_sizes = [7, 10]
    return _generate_file("map", fields, batch_sizes)


def generate_non_canonical_map_case():
    fields = [
        MapField('map_other_names',
                 get_field('some_key', 'utf8', nullable=False),
                 get_field('some_value', 'int32'),
                 entries_name='some_entries'),
    ]

    batch_sizes = [7]
    return _generate_file("map_non_canonical", fields, batch_sizes)


def generate_nested_case():
    fields = [
        ListField('list_nullable', get_field('item', 'int32')),
        FixedSizeListField('fixedsizelist_nullable',
                           get_field('item', 'int32'), 4),
        StructField('struct_nullable', [get_field('f1', 'int32'),
                                        get_field('f2', 'utf8')]),
        # Fails on Go (ARROW-8452)
        # ListField('list_nonnullable', get_field('item', 'int32'),
        #           nullable=False),
    ]

    batch_sizes = [7, 10]
    return _generate_file("nested", fields, batch_sizes)


def generate_recursive_nested_case():
    fields = [
        ListField('lists_list',
                  ListField('inner_list', get_field('item', 'int16'))),
        ListField('structs_list',
                  StructField('inner_struct',
                              [get_field('f1', 'int32'),
                               get_field('f2', 'utf8')])),
    ]

    batch_sizes = [7, 10]
    return _generate_file("recursive_nested", fields, batch_sizes)


def generate_nested_large_offsets_case():
    fields = [
        LargeListField('large_list_nullable', get_field('item', 'int32')),
        LargeListField('large_list_nonnullable',
                       get_field('item', 'int32'), nullable=False),
        LargeListField('large_list_nested',
                       ListField('inner_list', get_field('item', 'int16'))),
    ]

    batch_sizes = [0, 13]
    return _generate_file("nested_large_offsets", fields, batch_sizes)


def generate_unions_case():
    fields = [
        SparseUnionField('sparse', [get_field('f1', 'int32'),
                                    get_field('f2', 'utf8')],
                         type_ids=[5, 7]),
        DenseUnionField('dense', [get_field('f1', 'int16'),
                                  get_field('f2', 'binary')],
                        type_ids=[10, 20]),
        SparseUnionField('sparse', [get_field('f1', 'float32', nullable=False),
                                    get_field('f2', 'bool')],
                         type_ids=[5, 7], nullable=False),
        DenseUnionField('dense', [get_field('f1', 'uint8', nullable=False),
                                  get_field('f2', 'uint16'),
                                  NullField('f3')],
                        type_ids=[42, 43, 44], nullable=False),
    ]

    batch_sizes = [0, 11]
    return _generate_file("union", fields, batch_sizes)


def generate_dictionary_case():
    dict0 = Dictionary(0, StringField('dictionary1'), size=10, name='DICT0')
    dict1 = Dictionary(1, StringField('dictionary1'), size=5, name='DICT1')
    dict2 = Dictionary(2, get_field('dictionary2', 'int64'),
                       size=50, name='DICT2')

    fields = [
        DictionaryField('dict0', get_field('', 'int8'), dict0),
        DictionaryField('dict1', get_field('', 'int32'), dict1),
        DictionaryField('dict2', get_field('', 'int16'), dict2)
    ]
    batch_sizes = [7, 10]
    return _generate_file("dictionary", fields, batch_sizes,
                          dictionaries=[dict0, dict1, dict2])


def generate_dictionary_unsigned_case():
    dict0 = Dictionary(0, StringField('dictionary0'), size=5, name='DICT0')
    dict1 = Dictionary(1, StringField('dictionary1'), size=5, name='DICT1')
    dict2 = Dictionary(2, StringField('dictionary2'), size=5, name='DICT2')

    # TODO: JavaScript does not support uint64 dictionary indices, so disabled
    # for now

    # dict3 = Dictionary(3, StringField('dictionary3'), size=5, name='DICT3')
    fields = [
        DictionaryField('f0', get_field('', 'uint8'), dict0),
        DictionaryField('f1', get_field('', 'uint16'), dict1),
        DictionaryField('f2', get_field('', 'uint32'), dict2),
        # DictionaryField('f3', get_field('', 'uint64'), dict3)
    ]
    batch_sizes = [7, 10]
    return _generate_file("dictionary_unsigned", fields, batch_sizes,
                          dictionaries=[dict0, dict1, dict2])


def generate_nested_dictionary_case():
    dict0 = Dictionary(0, StringField('str'), size=10, name='DICT0')

    list_of_dict = ListField(
        'list',
        DictionaryField('str_dict', get_field('', 'int8'), dict0))
    dict1 = Dictionary(1, list_of_dict, size=30, name='DICT1')

    struct_of_dict = StructField('struct', [
        DictionaryField('str_dict_a', get_field('', 'int8'), dict0),
        DictionaryField('str_dict_b', get_field('', 'int8'), dict0)
    ])
    dict2 = Dictionary(2, struct_of_dict, size=30, name='DICT2')

    fields = [
        DictionaryField('list_dict', get_field('', 'int8'), dict1),
        DictionaryField('struct_dict', get_field('', 'int8'), dict2)
    ]

    batch_sizes = [10, 13]
    return _generate_file("nested_dictionary", fields, batch_sizes,
                          dictionaries=[dict0, dict1, dict2])


def generate_extension_case():
    dict0 = Dictionary(0, StringField('dictionary0'), size=5, name='DICT0')

    uuid_type = ExtensionType('uuid', 'uuid-serialized',
                              FixedSizeBinaryField('', 16))
    dict_ext_type = ExtensionType(
        'dict-extension', 'dict-extension-serialized',
        DictionaryField('str_dict', get_field('', 'int8'), dict0))

    fields = [
        ExtensionField('uuids', uuid_type),
        ExtensionField('dict_exts', dict_ext_type),
    ]

    batch_sizes = [0, 13]
    return _generate_file("extension", fields, batch_sizes,
                          dictionaries=[dict0])


def get_generated_json_files(tempdir=None):
    tempdir = tempdir or tempfile.mkdtemp(prefix='arrow-integration-')

    def _temp_path():
        return

    file_objs = [
        generate_primitive_case([], name='primitive_no_batches'),
        generate_primitive_case([17, 20], name='primitive'),
        generate_primitive_case([0, 0, 0], name='primitive_zerolength'),

        generate_primitive_large_offsets_case([17, 20])
        .skip_category('Go')
        .skip_category('JS'),

        generate_null_case([10, 0])
        .skip_category('JS')   # TODO(ARROW-7900)
        .skip_category('Go'),  # TODO(ARROW-7901)

        generate_null_trivial_case([0, 0])
        .skip_category('JS')   # TODO(ARROW-7900)
        .skip_category('Go'),  # TODO(ARROW-7901)

        generate_decimal128_case()
        .skip_category('Go')  # TODO(ARROW-7948): Decimal + Go
        .skip_category('Rust'),

        generate_decimal256_case()
        .skip_category('Go')  # TODO(ARROW-7948): Decimal + Go
        .skip_category('JS')
        .skip_category('Rust'),

        generate_datetime_case(),

        generate_interval_case()
        .skip_category('JS')  # TODO(ARROW-5239): Intervals + JS
        .skip_category('Rust'),

        generate_map_case()
        .skip_category('Go')  # TODO(ARROW-5620): Map + Go
        .skip_category('Rust'),

        generate_non_canonical_map_case()
        .skip_category('Go')     # TODO(ARROW-5620)
        .skip_category('Java')   # TODO(ARROW-8715)
        .skip_category('JS')     # TODO(ARROW-8716)
        .skip_category('Rust'),

        generate_nested_case(),

        generate_recursive_nested_case()
        .skip_category('Go'),    # TODO(ARROW-8453)

        generate_nested_large_offsets_case()
        .skip_category('Go')
        .skip_category('JS')
        .skip_category('Rust'),

        generate_unions_case()
        .skip_category('Go')
        .skip_category('JS')
        .skip_category('Rust'),

        generate_custom_metadata_case()
        .skip_category('Go')
        .skip_category('JS'),

        generate_duplicate_fieldnames_case()
        .skip_category('Go')
        .skip_category('JS'),

        # TODO(ARROW-3039, ARROW-5267): Dictionaries in GO
        generate_dictionary_case()
        .skip_category('Go'),

        generate_dictionary_unsigned_case()
        .skip_category('Go')     # TODO(ARROW-9378)
        .skip_category('Java'),  # TODO(ARROW-9377)

        generate_nested_dictionary_case()
        .skip_category('Go')
        .skip_category('Java')  # TODO(ARROW-7779)
        .skip_category('JS')
        .skip_category('Rust'),

        generate_extension_case()
        .skip_category('Go')
        .skip_category('JS')
        .skip_category('Rust'),
    ]

    generated_paths = []
    for file_obj in file_objs:
        out_path = os.path.join(tempdir, 'generated_' +
                                file_obj.name + '.json')
        file_obj.write(out_path)
        generated_paths.append(file_obj)

    return generated_paths