// 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.

import Foundation

public protocol ArrowArrayHolderBuilder {
    func toHolder() throws -> ArrowArrayHolder
    func appendAny(_ val: Any?)
}

public class ArrowArrayBuilder<T: ArrowBufferBuilder, U: ArrowArray<T.ItemType>>: ArrowArrayHolderBuilder {
    let type: ArrowType
    let bufferBuilder: T
    public var length: UInt {return self.bufferBuilder.length}
    public var capacity: UInt {return self.bufferBuilder.capacity}
    public var nullCount: UInt {return self.bufferBuilder.nullCount}
    public var offset: UInt {return self.bufferBuilder.offset}

    fileprivate init(_ type: ArrowType) throws {
        self.type = type
        self.bufferBuilder = try T()
    }

    public func append(_ vals: T.ItemType?...) {
        for val in vals {
            self.bufferBuilder.append(val)
        }
    }

    public func append(_ vals: [T.ItemType?]) {
        for val in vals {
            self.bufferBuilder.append(val)
        }
    }

    public func append(_ val: T.ItemType?) {
        self.bufferBuilder.append(val)
    }

    public func appendAny(_ val: Any?) {
        self.bufferBuilder.append(val as? T.ItemType)
    }

    public func finish() throws -> ArrowArray<T.ItemType> {
        let buffers = self.bufferBuilder.finish()
        let arrowData = try ArrowData(self.type, buffers: buffers, nullCount: self.nullCount)
        let array = try U(arrowData)
        return array
    }

    public func getStride() -> Int {
        return self.type.getStride()
    }

    public func toHolder() throws -> ArrowArrayHolder {
        return try ArrowArrayHolderImpl(self.finish())
    }
}

public class NumberArrayBuilder<T>: ArrowArrayBuilder<FixedBufferBuilder<T>, FixedArray<T>> {
    fileprivate convenience init() throws {
        try self.init(ArrowType(ArrowType.infoForNumericType(T.self)))
    }
}

public class StringArrayBuilder: ArrowArrayBuilder<VariableBufferBuilder<String>, StringArray> {
    fileprivate convenience init() throws {
        try self.init(ArrowType(ArrowType.ArrowString))
    }
}

public class BinaryArrayBuilder: ArrowArrayBuilder<VariableBufferBuilder<Data>, BinaryArray> {
    fileprivate convenience init() throws {
        try self.init(ArrowType(ArrowType.ArrowBinary))
    }
}

public class BoolArrayBuilder: ArrowArrayBuilder<BoolBufferBuilder, BoolArray> {
    fileprivate convenience init() throws {
        try self.init(ArrowType(ArrowType.ArrowBool))
    }
}

public class Date32ArrayBuilder: ArrowArrayBuilder<Date32BufferBuilder, Date32Array> {
    fileprivate convenience init() throws {
        try self.init(ArrowType(ArrowType.ArrowDate32))
    }
}

public class Date64ArrayBuilder: ArrowArrayBuilder<Date64BufferBuilder, Date64Array> {
    fileprivate convenience init() throws {
        try self.init(ArrowType(ArrowType.ArrowDate64))
    }
}

public class Time32ArrayBuilder: ArrowArrayBuilder<FixedBufferBuilder<Time32>, Time32Array> {
    fileprivate convenience init(_ unit: ArrowTime32Unit) throws {
        try self.init(ArrowTypeTime32(unit))
    }
}

public class Time64ArrayBuilder: ArrowArrayBuilder<FixedBufferBuilder<Time64>, Time64Array> {
    fileprivate convenience init(_ unit: ArrowTime64Unit) throws {
        try self.init(ArrowTypeTime64(unit))
    }
}

public class StructArrayBuilder: ArrowArrayBuilder<StructBufferBuilder, StructArray> {
    let builders: [any ArrowArrayHolderBuilder]
    let fields: [ArrowField]
    public init(_ fields: [ArrowField], builders: [any ArrowArrayHolderBuilder]) throws {
        self.fields = fields
        self.builders = builders
        try super.init(ArrowNestedType(ArrowType.ArrowStruct, fields: fields))
        self.bufferBuilder.initializeTypeInfo(fields)
    }

    public init(_ fields: [ArrowField]) throws {
        self.fields = fields
        var builders = [any ArrowArrayHolderBuilder]()
        for field in fields {
            builders.append(try ArrowArrayBuilders.loadBuilder(arrowType: field.type))
        }

        self.builders = builders
        try super.init(ArrowNestedType(ArrowType.ArrowStruct, fields: fields))
    }

    public override func append(_ values: [Any?]?) {
        self.bufferBuilder.append(values)
        if let anyValues = values {
            for index in 0..<builders.count {
                self.builders[index].appendAny(anyValues[index])
            }
        } else {
            for index in 0..<builders.count {
                self.builders[index].appendAny(nil)
            }
        }
    }

    public override func finish() throws -> StructArray {
        let buffers = self.bufferBuilder.finish()
        var childData = [ArrowData]()
        for builder in self.builders {
            childData.append(try builder.toHolder().array.arrowData)
        }

        let arrowData = try ArrowData(self.type, buffers: buffers,
                                      children: childData, nullCount: self.nullCount,
                                      length: self.length)
        let structArray = try StructArray(arrowData)
        return structArray
    }
}

public class ArrowArrayBuilders {
    public static func loadBuilder( // swiftlint:disable:this cyclomatic_complexity
        _ builderType: Any.Type) throws -> ArrowArrayHolderBuilder {
        if builderType == Int8.self || builderType == Int8?.self {
            return try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<Int8>
        } else if builderType == Int16.self || builderType == Int16?.self {
            return try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<Int16>
        } else if builderType == Int32.self || builderType == Int32?.self {
            return try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<Int32>
        } else if builderType == Int64.self || builderType == Int64?.self {
            return try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<Int64>
        } else if builderType == Float.self || builderType == Float?.self {
            return try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<Float>
        } else if builderType == UInt8.self || builderType == UInt8?.self {
            return try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<UInt8>
        } else if builderType == UInt16.self || builderType == UInt16?.self {
            return try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<UInt16>
        } else if builderType == UInt32.self || builderType == UInt32?.self {
            return try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<UInt32>
        } else if builderType == UInt64.self || builderType == UInt64?.self {
            return try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<UInt64>
        } else if builderType == Double.self || builderType == Double?.self {
            return try ArrowArrayBuilders.loadNumberArrayBuilder() as NumberArrayBuilder<Double>
        } else if builderType == String.self || builderType == String?.self {
            return try ArrowArrayBuilders.loadStringArrayBuilder()
        } else if builderType == Bool.self || builderType == Bool?.self {
            return try ArrowArrayBuilders.loadBoolArrayBuilder()
        } else if builderType == Date.self || builderType == Date?.self {
            return try ArrowArrayBuilders.loadDate64ArrayBuilder()
        } else {
            throw ArrowError.invalid("Invalid type for builder: \(builderType)")
        }
    }

    public static func isValidBuilderType<T>(_ type: T.Type) -> Bool {
        return type == Int8?.self || type == Int16?.self ||
            type == Int32?.self || type == Int64?.self ||
            type == UInt8?.self || type == UInt16?.self ||
            type == UInt32?.self || type == UInt64?.self ||
            type == String?.self || type == Double?.self ||
            type == Float?.self || type == Date?.self ||
            type == Bool?.self || type == Bool.self ||
            type == Int8.self || type == Int16.self ||
            type == Int32.self || type == Int64.self ||
            type == UInt8.self || type == UInt16.self ||
            type == UInt32.self || type == UInt64.self ||
            type == String.self || type == Double.self ||
            type == Float.self || type == Date.self
    }

    public static func loadStructArrayBuilderForType<T>(_ obj: T) throws -> StructArrayBuilder {
        let mirror = Mirror(reflecting: obj)
        var builders = [ArrowArrayHolderBuilder]()
        var fields = [ArrowField]()
        for (property, value) in mirror.children {
            guard let propertyName = property else {
                continue
            }

            let builderType = type(of: value)
            let arrowType = ArrowType(ArrowType.infoForType(builderType))
            fields.append(ArrowField(propertyName, type: arrowType, isNullable: true))
            builders.append(try loadBuilder(arrowType: arrowType))
        }

        return try StructArrayBuilder(fields, builders: builders)
    }

    public static func loadBuilder( // swiftlint:disable:this cyclomatic_complexity
        arrowType: ArrowType) throws -> ArrowArrayHolderBuilder {
        switch arrowType.id {
        case .uint8:
            return try loadNumberArrayBuilder() as NumberArrayBuilder<UInt8>
        case .uint16:
            return try loadNumberArrayBuilder() as NumberArrayBuilder<UInt16>
        case .uint32:
            return try loadNumberArrayBuilder() as NumberArrayBuilder<UInt32>
        case .uint64:
            return try loadNumberArrayBuilder() as NumberArrayBuilder<UInt64>
        case .int8:
            return try loadNumberArrayBuilder() as NumberArrayBuilder<Int8>
        case .int16:
            return try loadNumberArrayBuilder() as NumberArrayBuilder<Int16>
        case .int32:
            return try loadNumberArrayBuilder() as NumberArrayBuilder<Int32>
        case .int64:
            return try loadNumberArrayBuilder() as NumberArrayBuilder<Int64>
        case .double:
            return try loadNumberArrayBuilder() as NumberArrayBuilder<Double>
        case .float:
            return try loadNumberArrayBuilder() as NumberArrayBuilder<Float>
        case .string:
            return try StringArrayBuilder()
        case .boolean:
            return try BoolArrayBuilder()
        case .binary:
            return try BinaryArrayBuilder()
        case .date32:
            return try Date32ArrayBuilder()
        case .date64:
            return try Date64ArrayBuilder()
        case .time32:
            guard let timeType = arrowType as? ArrowTypeTime32 else {
                throw ArrowError.invalid("Expected arrow type for \(arrowType.id) not found")
            }
            return try Time32ArrayBuilder(timeType.unit)
        case .time64:
            guard let timeType = arrowType as? ArrowTypeTime64 else {
                throw ArrowError.invalid("Expected arrow type for \(arrowType.id) not found")
            }
            return try Time64ArrayBuilder(timeType.unit)
        default:
            throw ArrowError.unknownType("Builder not found for arrow type: \(arrowType.id)")
        }
    }

    public static func loadNumberArrayBuilder<T>() throws -> NumberArrayBuilder<T> {
        let type = T.self
        if type == Int8.self {
            return try NumberArrayBuilder<T>()
        } else if type == Int16.self {
            return try NumberArrayBuilder<T>()
        } else if type == Int32.self {
            return try NumberArrayBuilder<T>()
        } else if type == Int64.self {
            return try NumberArrayBuilder<T>()
        } else if type == UInt8.self {
            return try NumberArrayBuilder<T>()
        } else if type == UInt16.self {
            return try NumberArrayBuilder<T>()
        } else if type == UInt32.self {
            return try NumberArrayBuilder<T>()
        } else if type == UInt64.self {
            return try NumberArrayBuilder<T>()
        } else if type == Float.self {
            return try NumberArrayBuilder<T>()
        } else if type == Double.self {
            return try NumberArrayBuilder<T>()
        } else {
            throw ArrowError.unknownType("Type is invalid for NumberArrayBuilder")
        }
    }

    public static func loadStringArrayBuilder() throws -> StringArrayBuilder {
        return try StringArrayBuilder()
    }

    public static func loadBoolArrayBuilder() throws -> BoolArrayBuilder {
        return try BoolArrayBuilder()
    }

    public static func loadDate32ArrayBuilder() throws -> Date32ArrayBuilder {
        return try Date32ArrayBuilder()
    }

    public static func loadDate64ArrayBuilder() throws -> Date64ArrayBuilder {
        return try Date64ArrayBuilder()
    }

    public static func loadBinaryArrayBuilder() throws -> BinaryArrayBuilder {
        return try BinaryArrayBuilder()
    }

    public static func loadTime32ArrayBuilder(_ unit: ArrowTime32Unit) throws -> Time32ArrayBuilder {
        return try Time32ArrayBuilder(unit)
    }

    public static func loadTime64ArrayBuilder(_ unit: ArrowTime64Unit) throws -> Time64ArrayBuilder {
        return try Time64ArrayBuilder(unit)
    }
}
