package com.akexorcist.materialdesign

import android.content.Context
import androidx.core.view.MarginLayoutParamsCompat
import androidx.core.view.ViewCompat
import android.util.AttributeSet
import android.util.Log
import android.view.View
import android.widget.LinearLayout
import com.google.android.material.shape.AbsoluteCornerSize
import com.google.android.material.shape.CornerSize
import com.google.android.material.shape.ShapeAppearanceModel
import com.google.android.material.button.MaterialButton
import androidx.core.view.accessibility.AccessibilityNodeInfoCompat.CollectionItemInfoCompat
import androidx.core.view.accessibility.AccessibilityNodeInfoCompat
import androidx.core.view.AccessibilityDelegateCompat
import android.view.ViewGroup
import java.util.*

class MaterialButtonGroup : LinearLayout {
    constructor(context: Context?) : super(context, null)

    constructor(context: Context?, attrs: AttributeSet?) : super(
        context,
        attrs,
        R.attr.materialButtonToggleGroupStyle
    )

    constructor(context: Context?, attrs: AttributeSet?, defStyleAttr: Int) : super(
        context,
        attrs,
        defStyleAttr
    )

    private val originalCornerData: ArrayList<CornerData> = ArrayList()

    override fun addView(child: View?, index: Int, params: ViewGroup.LayoutParams?) {
        if (child !is MaterialButton) {
            Log.e("MaterialButtonGroup", "Child views must be of type MaterialButton.")
            return
        }
        super.addView(child, index, params)
        setGeneratedIdIfNeeded(child)
        setChildButton(child)

        // Saves original corner data
        val shapeAppearanceModel = child.shapeAppearanceModel
        originalCornerData.add(
            CornerData(
                shapeAppearanceModel.topLeftCornerSize,
                shapeAppearanceModel.bottomLeftCornerSize,
                shapeAppearanceModel.topRightCornerSize,
                shapeAppearanceModel.bottomRightCornerSize
            )
        )
        ViewCompat.setAccessibilityDelegate(child, object : AccessibilityDelegateCompat() {
            override fun onInitializeAccessibilityNodeInfo(
                host: View,
                info: AccessibilityNodeInfoCompat
            ) {
                super.onInitializeAccessibilityNodeInfo(host, info)
                info.setCollectionItemInfo(
                    CollectionItemInfoCompat.obtain(
                        0,
                        1,
                        getIndexWithinVisibleButtons(host),
                        1,
                        false,
                        (host as MaterialButton).isSelected
                    )
                )
            }
        })
    }

    override fun onViewRemoved(child: View?) {
        super.onViewRemoved(child)
        val indexOfChild = indexOfChild(child)
        if (indexOfChild >= 0) {
            originalCornerData.removeAt(indexOfChild)
        }
        updateChildShapes()
        adjustChildMarginsAndUpdateLayout()
    }

    override fun onMeasure(widthMeasureSpec: Int, heightMeasureSpec: Int) {
        updateChildShapes()
        adjustChildMarginsAndUpdateLayout()
        super.onMeasure(widthMeasureSpec, heightMeasureSpec)
    }

    private fun adjustChildMarginsAndUpdateLayout() {
        val firstVisibleChildIndex = getFirstVisibleChildIndex()
        if (firstVisibleChildIndex == -1) {
            return
        }
        for (i in firstVisibleChildIndex + 1 until childCount) {
            val currentButton = getChildButton(i)
            val previousButton = getChildButton(i - 1)
            val smallestStrokeWidth = currentButton.strokeWidth.coerceAtMost(previousButton.strokeWidth)
            val params: LayoutParams = buildLayoutParams(currentButton)
            if (orientation == HORIZONTAL) {
                MarginLayoutParamsCompat.setMarginEnd(params, 0)
                MarginLayoutParamsCompat.setMarginStart(params, -smallestStrokeWidth)
                params.topMargin = 0
            } else {
                params.bottomMargin = 0
                params.topMargin = -smallestStrokeWidth
                MarginLayoutParamsCompat.setMarginStart(params, 0)
            }
            currentButton.layoutParams = params
        }
        resetChildMargins(firstVisibleChildIndex)
    }

    private fun getChildButton(index: Int): MaterialButton {
        return getChildAt(index) as MaterialButton
    }

    private fun resetChildMargins(childIndex: Int) {
        if (childCount == 0 || childIndex == -1) {
            return
        }
        val currentButton = getChildButton(childIndex)
        val params = currentButton.layoutParams as LayoutParams
        if (orientation == VERTICAL) {
            params.topMargin = 0
            params.bottomMargin = 0
            return
        }
        MarginLayoutParamsCompat.setMarginEnd(params, 0)
        MarginLayoutParamsCompat.setMarginStart(params, 0)
        params.leftMargin = 0
        params.rightMargin = 0
    }

    private fun updateChildShapes() {
        val childCount = childCount
        val firstVisibleChildIndex: Int = getFirstVisibleChildIndex()
        val lastVisibleChildIndex: Int = getLastVisibleChildIndex()
        for (i in 0 until childCount) {
            val button: MaterialButton = getChildButton(i)
            if (button.visibility == GONE) {
                continue
            }
            val builder = button.shapeAppearanceModel.toBuilder()
            val newCornerData = getNewCornerData(i, firstVisibleChildIndex, lastVisibleChildIndex)
            updateBuilderWithCornerData(builder, newCornerData)
            button.shapeAppearanceModel = builder.build()
        }
    }

    private fun getFirstVisibleChildIndex(): Int {
        val childCount = childCount
        for (i in 0 until childCount) {
            if (isChildVisible(i)) {
                return i
            }
        }
        return -1
    }

    private fun getLastVisibleChildIndex(): Int {
        val childCount = childCount
        for (i in childCount - 1 downTo 0) {
            if (isChildVisible(i)) {
                return i
            }
        }
        return -1
    }

    private fun isChildVisible(i: Int): Boolean {
        val child = getChildAt(i)
        return child.visibility != GONE
    }

    private fun getVisibleButtonCount(): Int {
        var count = 0
        for (i in 0 until childCount) {
            if (getChildAt(i) is MaterialButton && isChildVisible(i)) {
                count++
            }
        }
        return count
    }

    private fun getIndexWithinVisibleButtons(child: View?): Int {
        if (child !is MaterialButton) {
            return -1
        }
        var index = 0
        for (i in 0 until childCount) {
            if (getChildAt(i) === child) {
                return index
            }
            if (getChildAt(i) is MaterialButton && isChildVisible(i)) {
                index++
            }
        }
        return -1
    }

    private fun setGeneratedIdIfNeeded(materialButton: MaterialButton) {
        if (materialButton.id == NO_ID) {
            materialButton.id = ViewCompat.generateViewId()
        }
    }

    private fun setChildButton(materialButton: MaterialButton) {
        materialButton.insetTop = 0
        materialButton.insetBottom = 0
    }

    private fun buildLayoutParams(child: View): LayoutParams {
        val layoutParams = child.layoutParams
        return if (layoutParams is LayoutParams) layoutParams
        else LayoutParams(layoutParams.width, layoutParams.height)
    }

    private fun getNewCornerData(
        index: Int,
        firstVisibleChildIndex: Int,
        lastVisibleChildIndex: Int
    ): CornerData? {
        val cornerData: CornerData = originalCornerData[index]
        if (firstVisibleChildIndex == lastVisibleChildIndex) {
            return cornerData
        }
        val isHorizontal = orientation == HORIZONTAL
        if (index == firstVisibleChildIndex) {
            return if (isHorizontal) CornerData.start(cornerData, this)
            else CornerData.top(cornerData)
        }
        return if (index == lastVisibleChildIndex) {
            if (isHorizontal) CornerData.end(cornerData, this)
            else CornerData.bottom(cornerData)
        } else null
    }

    private fun updateBuilderWithCornerData(
        shapeAppearanceModelBuilder: ShapeAppearanceModel.Builder,
        cornerData: CornerData?
    ) {
        if (cornerData == null) {
            shapeAppearanceModelBuilder.setAllCornerSizes(0f)
            return
        }
        shapeAppearanceModelBuilder
            .setTopLeftCornerSize(cornerData.topLeft)
            .setBottomLeftCornerSize(cornerData.bottomLeft)
            .setTopRightCornerSize(cornerData.topRight)
            .setBottomRightCornerSize(cornerData.bottomRight)
    }
}

class CornerData(
    var topLeft: CornerSize,
    var bottomLeft: CornerSize,
    var topRight: CornerSize,
    var bottomRight: CornerSize
) {
    companion object {
        private val noCorner: CornerSize = AbsoluteCornerSize(0f)

        fun start(orig: CornerData, view: View): CornerData {
            return if (view.isLayoutRtl()) right(orig) else left(orig)
        }

        fun end(orig: CornerData, view: View): CornerData {
            return if (view.isLayoutRtl()) left(orig) else right(orig)
        }

        fun left(orig: CornerData): CornerData {
            return CornerData(orig.topLeft, orig.bottomLeft, noCorner, noCorner)
        }

        fun right(orig: CornerData): CornerData {
            return CornerData(noCorner, noCorner, orig.topRight, orig.bottomRight)
        }

        fun top(orig: CornerData): CornerData {
            return CornerData(orig.topLeft, noCorner, orig.topRight, noCorner)
        }

        fun bottom(orig: CornerData): CornerData {
            return CornerData(noCorner, orig.bottomLeft, noCorner, orig.bottomRight)
        }
    }
}

fun View.isLayoutRtl(): Boolean {
    return ViewCompat.getLayoutDirection(this) == ViewCompat.LAYOUT_DIRECTION_RTL
}