# Concat implementation for ArrowWeaveList. Should work in all cases!
#
# Any places where we extend lists need to follow our merge_types implementation.
# Fortunately, concat is the only place that extends lists! Everything else should
# make use of concat to do so.


import dataclasses
import typing

import numpy as np
import pyarrow as pa
from pyarrow import compute as pc

from weave.legacy.weave import errors
from weave.legacy.weave import weave_types as types
from weave.legacy.weave.arrow.list_ import (
    ArrowWeaveList,
    ArrowWeaveListGeneric,
    is_list_arrowweavelist,
    is_object_arrowweavelist,
    is_ref_arrowweavelist,
    is_taggedvalue_arrowweavelist,
    is_typedict_arrowweavelist,
    make_vec_none,
    offsets_starting_at_zero,
    unsafe_awl_construction,
)
from weave.legacy.weave.language_features.tagging import tagged_value_type

DEBUG = False

from weave.legacy.weave import engine_trace

tracer = engine_trace.tracer()


@dataclasses.dataclass
class UnionMember:
    values: ArrowWeaveList
    mask: pa.Array
    offsets: pa.Array


def indent_print(indent: int = 0, *args) -> None:
    if DEBUG:
        print("    " * indent + " ".join([str(a) for a in args]))


def pa_concat_arrays(arrays: list[typing.Union[pa.Array, pa.ChunkedArray]]) -> pa.Array:
    arrays = [a if isinstance(a, pa.Array) else a.combine_chunks() for a in arrays]

    with tracer.trace("pa_concat_arrays"):
        try:
            return pa.concat_arrays(arrays)
        except pa.lib.ArrowNotImplementedError as e:
            # pyarrow doesn't support concatenating all dictionaries, so we have to do it ourselves.
            if all(isinstance(a, pa.DictionaryArray) for a in arrays):
                # We can concatenate dictionaries by concatenating the indices and the dictionary
                # values separately.

                index_offset = 0
                indices = []
                for a in arrays:
                    indices.append(pc.add(a.indices, index_offset).cast(pa.int32()))
                    index_offset += len(a.dictionary)

                indices = pa_concat_arrays(indices)
                dictionary = pa_concat_arrays([a.dictionary for a in arrays])
                return pa.DictionaryArray.from_arrays(indices, dictionary)
            else:
                raise


def _concatenate_typeddicts(
    l1: ArrowWeaveListGeneric[types.TypedDict],
    l2: ArrowWeaveListGeneric[types.TypedDict],
    depth: int = 0,
) -> ArrowWeaveList:
    # Columns are in the order they're encountered.
    all_keys = {**l1.object_type.property_types, **l2.object_type.property_types}
    if not all_keys:
        # Both are empty
        return ArrowWeaveList(
            pa_concat_arrays([l1._arrow_data, l2._arrow_data]),
            types.TypedDict({}),
            invalid_reason="Possibly nullable",
        )
    properties: dict[str, ArrowWeaveList] = {}
    for key in all_keys:
        indent_print(depth, "TypedDict key", key)
        if key not in l1.object_type.property_types:
            properties[key] = _concatenate(
                make_vec_none(len(l1._arrow_data)), l2.column(key), depth=depth + 1
            )
        elif key not in l2.object_type.property_types:
            properties[key] = _concatenate(
                l1.column(key), make_vec_none(len(l2)), depth=depth + 1
            )
        else:
            properties[key] = _concatenate(
                l1.column(key), l2.column(key), depth=depth + 1
            )

    return ArrowWeaveList(
        pa.StructArray.from_arrays(
            [properties[key]._arrow_data for key in all_keys],
            list(all_keys),
            mask=pa_concat_arrays(
                [
                    pa.compute.is_null(l1._arrow_data),
                    pa.compute.is_null(l2._arrow_data),
                ]
            ),
        ),
        types.TypedDict({key: properties[key].object_type for key in all_keys}),
        l1._artifact,
        invalid_reason="Possibly nullable",
    )


def _concatenate_objects(
    l1: ArrowWeaveListGeneric[types.ObjectType],
    l2: ArrowWeaveListGeneric[types.ObjectType],
    depth: int = 0,
) -> ArrowWeaveList:
    assert l1.object_type.name == l2.object_type.name
    attrs: dict[str, ArrowWeaveList] = {}
    merged_type = typing.cast(
        types.ObjectType, types.merge_types(l1.object_type, l2.object_type)
    )
    for key in merged_type.property_types():
        attrs[key] = _concatenate(l1.column(key), l2.column(key), depth=depth + 1)
    return ArrowWeaveList(
        pa.StructArray.from_arrays(
            [attrs[key]._arrow_data for key in attrs],
            list(attrs.keys()),
            mask=pa_concat_arrays(
                [
                    pa.compute.is_null(l1._arrow_data),
                    pa.compute.is_null(l2._arrow_data),
                ]
            ),
        ),
        merged_type,
        l1._artifact,
        invalid_reason="Possibly nullable",
    )


def _concatenate_taggedvalues(
    l1: ArrowWeaveListGeneric[tagged_value_type.TaggedValueType],
    l2: ArrowWeaveListGeneric[tagged_value_type.TaggedValueType],
    depth: int = 0,
) -> ArrowWeaveList:
    tag = _concatenate(l1.tagged_value_tag(), l2.tagged_value_tag(), depth=depth + 1)
    value = _concatenate(
        l1.tagged_value_value(), l2.tagged_value_value(), depth=depth + 1
    )
    return ArrowWeaveList(
        pa.StructArray.from_arrays(
            [tag._arrow_data, value._arrow_data],
            ["_tag", "_value"],
            mask=pa_concat_arrays(
                [
                    pa.compute.is_null(l1._arrow_data),
                    pa.compute.is_null(l2._arrow_data),
                ]
            ),
        ),
        tagged_value_type.TaggedValueType(tag.object_type, value.object_type),  # type: ignore
        l1._artifact,
        invalid_reason="Possibly nullable",
    )


def convert_null_array_to_list(
    array: typing.Union[pa.ListArray, pa.NullArray],
) -> pa.ListArray:
    if isinstance(array, pa.NullArray):
        return pc.cast(array, pa.list_(pa.null()))
    else:
        return array


def _concatenate_lists(
    l1: ArrowWeaveListGeneric[types.List],
    l2: ArrowWeaveListGeneric[types.List],
    depth: int = 0,
) -> ArrowWeaveList:
    l1_arrow_data = convert_null_array_to_list(l1._arrow_data)
    l2_arrow_data = convert_null_array_to_list(l2._arrow_data)

    self_values: ArrowWeaveList = ArrowWeaveList(
        l1_arrow_data.flatten(), l1.object_type.object_type, l2._artifact
    )
    other_values: ArrowWeaveList = ArrowWeaveList(
        l2_arrow_data.flatten(), l2.object_type.object_type, l2._artifact
    )
    concatted_values = _concatenate(self_values, other_values, depth=depth + 1)
    new_offsets = pa_concat_arrays(
        [
            offsets_starting_at_zero(l1_arrow_data)[:-1],
            pa.compute.add(
                offsets_starting_at_zero(l2_arrow_data),
                len(self_values),
            ).cast(pa.int32()),
        ]
    )
    return ArrowWeaveList(
        pa.ListArray.from_arrays(
            new_offsets,
            concatted_values._arrow_data,
            mask=pa_concat_arrays(
                [
                    pa.compute.is_null(l1_arrow_data),
                    pa.compute.is_null(l2_arrow_data),
                ]
            ),
        ),
        types.List(
            types.merge_types(l1.object_type.object_type, l2.object_type.object_type)
        ),
        l1._artifact,
        invalid_reason="Possibly nullable",
    )


def _concatenate_non_unions(
    self: ArrowWeaveList, other: ArrowWeaveList, depth: int = 0
) -> typing.Optional[ArrowWeaveList]:
    if is_typedict_arrowweavelist(self) and is_typedict_arrowweavelist(other):
        return _concatenate_typeddicts(self, other, depth=depth)
    elif is_object_arrowweavelist(self) and is_object_arrowweavelist(other):
        if self.object_type.name == other.object_type.name:
            return _concatenate_objects(self, other, depth=depth)
    elif is_taggedvalue_arrowweavelist(self) and is_taggedvalue_arrowweavelist(other):
        return _concatenate_taggedvalues(self, other, depth=depth)
    elif is_list_arrowweavelist(self) and is_list_arrowweavelist(other):
        # Lists with different object types
        return _concatenate_lists(self, other, depth=depth)

    if isinstance(self.object_type, types.Number) and isinstance(
        other.object_type, types.Number
    ):
        # Numeric special case
        indent_print(
            depth, "Extend case number types", self.object_type, other.object_type
        )

        # Weave numeric types are odd for legacy reasons. We have three numeric types
        # Number, Float, Int. Unfortunately Number is a Union of Float and Int, but
        # also assignable to Int.

        # We should fix that. But for now we rely on the arrow type for this special
        # case.

        if self._arrow_data.type == other._arrow_data.type:
            new_arrow_type = self._arrow_data.type
        else:
            # If types aren't equal, just cast to float64. This is not entirely
            # correct! Very large integers will be truncated.
            new_arrow_type = pa.float64()

        new_weave_type: types.Type
        if pa.types.is_integer(new_arrow_type):
            new_weave_type = types.Int()
        elif pa.types.is_floating(new_arrow_type):
            new_weave_type = types.Float()
        else:
            raise errors.WeaveInternalError(
                f'Unexpected type "{new_arrow_type}" for "Number" type.'
            )

        if len(self) == 0:
            data = other._arrow_data
        elif len(other) == 0:
            data = self._arrow_data
        else:
            data = pa_concat_arrays(
                [
                    self._arrow_data.cast(new_arrow_type),
                    other._arrow_data.cast(new_arrow_type),
                ]
            )
        return ArrowWeaveList(
            data, new_weave_type, self._artifact, invalid_reason="Possibly nullable"
        )
    elif (
        # Types exactly equal case
        self.object_type == other.object_type
        or (
            # Types are refs, so arrow type is string
            is_ref_arrowweavelist(self)
            and is_ref_arrowweavelist(other)
            and self.object_type.__class__ == other.object_type.__class__
        )
    ):
        indent_print(depth, "Extend case equal", self.object_type, other.object_type)
        if len(self) == 0:
            data = other._arrow_data
        elif len(other) == 0:
            data = self._arrow_data
        else:
            data = pa_concat_arrays([self._arrow_data, other._arrow_data])
        return ArrowWeaveList(
            data,
            types.merge_types(self.object_type, other.object_type),
            self._artifact,
            invalid_reason="Possibly nullable",
        )
    return None


def _concatenate(
    self: "ArrowWeaveList", other: "ArrowWeaveList", depth=0
) -> "ArrowWeaveList":
    if depth > 50:
        raise ValueError("Maximum recursion depth exceeded")
    indent_print(depth, "EXTEND TOP", depth, self.object_type, other.object_type)

    # We use UnknownType for two things:
    #   1. to indicate that we don't know the type of an object,
    #   2. as the object type of a zero length array.
    # The former case, where we don't know the type of an object,
    # should never happen here! This is runtime and we know the types of
    # things.

    # Zero length arrays of unknown type
    if self.object_type == types.UnknownType():
        if len(self) == 0:
            return other
        else:
            raise errors.WeaveInternalError(
                'Encountered non-zero length "UnknownType" array'
            )
    if other.object_type == types.UnknownType():
        if len(other) == 0:
            return self
        else:
            raise errors.WeaveInternalError(
                'Encountered non-zero length "UnknownType" array'
            )

    # Cases where one of the types is None
    if isinstance(self.object_type, types.NoneType) and isinstance(
        other.object_type, types.NoneType
    ):
        return make_vec_none(len(self) + len(other))
    elif isinstance(self.object_type, types.NoneType):
        indent_print(depth, "Extend case self None")
        if len(self) == 0:
            return other
        elif len(other) == 0:
            return self
        self_data = pa.nulls(len(self), other._arrow_data.type)
        other_data = other._arrow_data
        data = pa_concat_arrays([self_data, other_data])
        return ArrowWeaveList(data, types.optional(other.object_type), other._artifact)
    elif isinstance(other.object_type, types.NoneType):
        indent_print(depth, "Extend case other None")
        if len(self) == 0:
            return other
        elif len(other) == 0:
            return self
        self_data = self._arrow_data
        other_data = pa.nulls(len(other), self._arrow_data.type)
        data = pa_concat_arrays([self_data, other_data])
        return ArrowWeaveList(data, types.optional(self.object_type), self._artifact)

    # Separate nulls from type
    self_nullable, self_non_none_type = types.split_none(self.object_type)
    other_nullable, other_non_none_type = types.split_none(other.object_type)
    self = self._with_object_type(
        self_non_none_type, invalid_reason="Possibly nullable"
    )
    other = other._with_object_type(
        other_non_none_type, invalid_reason="Possibly nullable"
    )
    result_nullable = self_nullable or other_nullable

    if not isinstance(self.object_type, types.UnionType) and not isinstance(
        other.object_type, types.UnionType
    ):
        result = _concatenate_non_unions(self, other, depth=depth)
        if result is not None:
            if result_nullable:
                return result._with_object_type(types.optional(result.object_type))
            else:
                return result._clear_invalid_reason()

    # Otherwise, we have types that can't be merged, we're going to produce a union

    if isinstance(self.object_type, types.UnionType):
        self_field_types = self.object_type.members
        self_type_codes = self._arrow_data.type_codes
        self_offsets = self._arrow_data.offsets
        self_fields = [
            self._arrow_data.field(i) for i in range(len(self._arrow_data.type))
        ]
    else:
        self_field_types = [self.object_type]
        self_type_codes = pa.array(np.zeros(len(self), dtype=np.int8))
        self_offsets = pa.array(np.arange(len(self), dtype=np.int32))
        self_fields = [self._arrow_data]

    if isinstance(other.object_type, types.UnionType):
        other_field_types = other.object_type.members
        other_type_codes = other._arrow_data.type_codes
        other_offsets = other._arrow_data.offsets
        other_fields = [
            other._arrow_data.field(i) for i in range(len(other._arrow_data.type))
        ]
    else:
        other_field_types = [other.object_type]
        other_type_codes = pa.array(np.zeros(len(other), dtype=np.int8))
        other_offsets = pa.array(np.arange(len(other), dtype=np.int32))
        other_fields = [other._arrow_data]

    assert len(self_field_types) == len(self_fields)
    assert len(other_field_types) == len(other_fields)

    # produce merged union
    new_member_awls: list[UnionMember] = []

    indent_print(depth, "MERGING UNIONS")
    indent_print(depth, "SELF NULLABLE", self_nullable)
    # indent_print(depth, "SELF MEMBERS")
    # for member in self_fields:
    #     indent_print(depth + 1, member)
    indent_print(depth, "OTHER NULLABLE", other_nullable)
    # indent_print(depth, "OTHER MEMBERS")
    # for member in other_fields:
    #     indent_print(depth + 1, member)

    remaining_other_indexes = {i: True for i in range(len(other_fields))}
    for self_i, (self_member_type, self_field) in enumerate(
        zip(self_field_types, self_fields)
    ):
        for other_i in remaining_other_indexes:
            other_member_type = other_field_types[other_i]
            other_field = other_fields[other_i]
            if types.types_are_mergeable(self_member_type, other_member_type):
                indent_print(
                    depth,
                    "SELF OTHER FIELD merge",
                    self_i,
                    other_i,
                    "self type:",
                    self_member_type,
                    "other type:",
                    other_member_type,
                )
                concatted = _concatenate_non_unions(
                    ArrowWeaveList(
                        self_field, self_member_type, invalid_reason="Possibly nullable"
                    ),
                    ArrowWeaveList(
                        other_field,
                        other_member_type,
                        invalid_reason="Possibly nullable",
                    ),
                )
                if concatted is None:
                    raise errors.WeaveInternalError(
                        "Could not concatenate mergable types"
                    )
                new_member_awls.append(
                    UnionMember(
                        values=concatted,
                        mask=pa_concat_arrays(
                            [
                                pa.compute.equal(self_type_codes, self_i).cast(
                                    pa.int8()
                                ),
                                pa.compute.equal(other_type_codes, other_i).cast(
                                    pa.int8()
                                ),
                            ]
                        ),
                        offsets=pa_concat_arrays(
                            [
                                self_offsets,
                                pa.compute.add(other_offsets, len(self_field)).cast(
                                    pa.int32()
                                ),
                            ]
                        ),
                    )
                )
                indent_print(
                    depth,
                    "SELF OTHER FIELD merge RESULT",
                    new_member_awls[-1].values.object_type,
                )
                remaining_other_indexes.pop(other_i)
                break
        else:
            # no matching other member found
            indent_print(depth, "SELF FIELD solo", self_i, self_member_type)
            new_member_awls.append(
                UnionMember(
                    values=ArrowWeaveList(
                        self_fields[self_i],
                        self_member_type,
                        self._artifact,
                        invalid_reason="Possibly nullable",
                    ),
                    mask=pa_concat_arrays(
                        [
                            pa.compute.equal(self_type_codes, self_i).cast(pa.int8()),
                            pa.array(np.zeros(len(other), dtype=np.int8)),
                        ]
                    ),
                    # offsets=pa.array(np.arange(len(self), dtype=np.int32)),
                    offsets=pa_concat_arrays(
                        [
                            self_offsets,
                            pa.array(np.zeros(len(other), dtype=np.int32)),
                        ]
                    ),
                )
            )
    for other_i in remaining_other_indexes:
        other_member_type = other_field_types[other_i]
        indent_print(depth, "OTHER FIELD solo", other_i, other_member_type)
        other_field = other_fields[other_i]
        new_member_awls.append(
            UnionMember(
                values=ArrowWeaveList(
                    other_field,
                    other_member_type,
                    other._artifact,
                    invalid_reason="Possibly nullable",
                ),
                mask=pa_concat_arrays(
                    [
                        pa.array(np.zeros(len(self), dtype=np.int8)),
                        pa.compute.equal(other_type_codes, other_i).cast(pa.int8()),
                    ]
                ),
                # offsets=pa.array(np.arange(len(other), dtype=np.int32)),
                offsets=pa_concat_arrays(
                    [
                        pa.array(np.zeros(len(self), dtype=np.int32)),
                        other_offsets,
                    ]
                ),
            )
        )

    # I think we will always have at least two union members here. Because we
    # must have had a union on one side coming in.
    if len(new_member_awls) < 2:
        raise errors.WeaveInternalError("Not enough union members after merge")

    member_types = [m.values.object_type for m in new_member_awls]

    result_type_codes = pa.array(np.zeros(len(self) + len(other), dtype=np.int8))
    for i, new_member_awl in enumerate(new_member_awls):
        result_type_codes = pa.compute.add(
            result_type_codes, pa.compute.multiply(new_member_awl.mask, i)
        )
    result_offsets = pa.array(np.zeros(len(self) + len(other), dtype=np.int32))
    for i, new_member_awl in enumerate(new_member_awls):
        result_offsets = pa.compute.add(
            result_offsets,
            pa.compute.multiply(new_member_awl.mask, new_member_awl.offsets),
        )

    assert len(member_types) == len(new_member_awls)
    expected_object_type_members = len(new_member_awls)
    new_object_type = types.UnionType(*member_types)
    if len(new_object_type.members) != expected_object_type_members:
        indent_print(
            depth,
            "ERROR PRODUCING UNION",
            len(new_object_type.members),
            len(new_member_awls),
        )
        indent_print(depth, "new_object.type.members")
        for m in new_object_type.members:
            indent_print(depth + 1, m)
        indent_print(depth, "new_member_awls")
        for m2 in new_member_awls:
            indent_print(depth + 1, m2.values.object_type)

        assert len(new_object_type.members) == expected_object_type_members

    indent_print(depth, "NEW MEMBER TYPE", new_object_type)
    indent_print(depth, "RESULT NULLABLE", result_nullable)
    return ArrowWeaveList(
        pa.UnionArray.from_dense(
            result_type_codes.cast(pa.int8()),
            result_offsets.cast(pa.int32()),
            [m.values._arrow_data for m in new_member_awls],
        ),
        types.optional(new_object_type) if result_nullable else new_object_type,
        self._artifact,
    )


def concatenate(
    self: "ArrowWeaveList", other: "ArrowWeaveList", depth=0
) -> "ArrowWeaveList":
    with unsafe_awl_construction("Possibly nullable"):
        result = _concatenate(self, other, depth)
    result.validate()
    return result


def _merge_concat_split(
    arr: list[ArrowWeaveList],
) -> tuple[list[ArrowWeaveList], list[ArrowWeaveList]]:
    if len(arr) < 2:
        raise ValueError("arr must have length of at least 2")
    middle_index = len(arr) // 2
    return arr[:middle_index], arr[middle_index:]


def concatenate_all(arr: list[ArrowWeaveList]) -> ArrowWeaveList:
    if len(arr) == 0:
        raise ValueError("arr must not be empty")
    if len(arr) == 1:
        return arr[0]
    left, right = _merge_concat_split(arr)
    return concatenate_all(left).concat(concatenate_all(right))
