package mods.betterfoliage.render.block.vanilla

import mods.betterfoliage.BetterFoliage
import mods.betterfoliage.BetterFoliageMod
import mods.betterfoliage.chunk.BlockCtx
import mods.betterfoliage.config.Config
import mods.betterfoliage.config.DIRT_BLOCKS
import mods.betterfoliage.config.SALTWATER_BIOMES
import mods.betterfoliage.config.isSnow
import mods.betterfoliage.integration.ShadersModIntegration
import mods.betterfoliage.model.HalfBakedSpecialWrapper
import mods.betterfoliage.model.HalfBakedWrapperKey
import mods.betterfoliage.model.SpecialRenderData
import mods.betterfoliage.model.SpecialRenderModel
import mods.betterfoliage.model.SpriteSetDelegate
import mods.betterfoliage.model.buildTufts
import mods.betterfoliage.model.tuftModelSet
import mods.betterfoliage.model.tuftShapeSet
import mods.betterfoliage.render.lighting.LightingPreferredFace
import mods.betterfoliage.render.pipeline.Layers
import mods.betterfoliage.render.pipeline.RenderCtxBase
import mods.betterfoliage.render.pipeline.extendLayers
import mods.betterfoliage.resource.discovery.AbstractModelDiscovery
import mods.betterfoliage.resource.discovery.BakeWrapperManager
import mods.betterfoliage.resource.discovery.ModelBakingContext
import mods.betterfoliage.resource.discovery.ModelDiscoveryContext
import mods.betterfoliage.resource.generated.CenteredSprite
import mods.betterfoliage.util.Atlas
import mods.betterfoliage.util.Int3
import mods.betterfoliage.util.LazyInvalidatable
import mods.betterfoliage.util.get
import mods.betterfoliage.util.getBlockModel
import mods.betterfoliage.util.idxOrNull
import mods.betterfoliage.util.offset
import mods.betterfoliage.util.randomI
import net.minecraft.block.material.Material
import net.minecraft.client.Minecraft
import net.minecraft.client.renderer.RenderType
import net.minecraft.client.renderer.RenderTypeLookup
import net.minecraft.client.renderer.model.BlockModel
import net.minecraft.util.Direction.UP
import net.minecraft.util.ResourceLocation
import java.util.Random

object StandardDirtDiscovery : AbstractModelDiscovery() {
    override fun processModel(ctx: ModelDiscoveryContext) {
        if (ctx.getUnbaked() is BlockModel && ctx.blockState.block in DIRT_BLOCKS) {
            BetterFoliage.blockTypes.dirt.add(ctx.blockState)
            ctx.addReplacement(StandardDirtKey)
            ctx.blockState.block.extendLayers()
        }
        super.processModel(ctx)
    }
}

object StandardDirtKey : HalfBakedWrapperKey() {
    override fun bake(ctx: ModelBakingContext, wrapped: SpecialRenderModel) = StandardDirtModel(wrapped)
}

class DirtRenderData(
    val connectedGrassModel: SpecialRenderModel?,
    val algaeIdx: Int?,
    val reedIdx: Int?
) : SpecialRenderData {
    override fun canRenderInLayer(layer: RenderType) = when {
        connectedGrassModel != null && layer == Layers.connectedDirt -> true
        (algaeIdx != null || reedIdx != null) && layer == Layers.tufts -> true
        else -> false
    }
}

class StandardDirtModel(
    wrapped: SpecialRenderModel
) : HalfBakedSpecialWrapper(wrapped) {
    val vanillaTuftLighting = LightingPreferredFace(UP)

    override fun prepare(ctx: BlockCtx, random: Random): Any {
        if (!Config.enabled) return Unit
        val stateUp = ctx.state(UP)
        val state2Up = ctx.state(Int3(0, 2, 0))
        val isConnectedGrass = Config.connectedGrass.enabled &&
                stateUp in BetterFoliage.blockTypes.grass &&
                (Config.connectedGrass.snowEnabled || !state2Up.isSnow)

        val isWater = stateUp.material == Material.WATER
        val isDeepWater = isWater && state2Up.material == Material.WATER
        val isShallowWater = isWater && state2Up.isAir
        val isSaltWater = isWater && ctx.biome?.biomeCategory in SALTWATER_BIOMES

        return DirtRenderData(
            connectedGrassModel = if (isConnectedGrass) (getBlockModel(stateUp) as? SpecialRenderModel)?.resolve(random) else null,
            algaeIdx = random.idxOrNull(algaeModels) { Config.algae.enabled(random) && isDeepWater && isSaltWater },
            reedIdx = random.idxOrNull(reedModels) { Config.reed.enabled(random) && isShallowWater && !isSaltWater }
        )
    }

    override fun renderLayer(ctx: RenderCtxBase, data: Any, layer: RenderType) {
        if (data is DirtRenderData) {
            if (data.connectedGrassModel != null) {
                if (layer == Layers.connectedDirt) {
                    ctx.renderMasquerade(UP.offset) {
                        data.connectedGrassModel.renderLayer(ctx, ctx.state(UP), layer)
                    }
                }
            } else {
                // render non-connected grass
                super.renderLayer(ctx, data, layer)
            }

            if (layer == Layers.tufts) {
                data.algaeIdx?.let {
                    ctx.vertexLighter = vanillaTuftLighting
                    ShadersModIntegration.grass(ctx, Config.algae.shaderWind) {
                        ctx.renderQuads(algaeModels[it])
                    }
                }
                data.reedIdx?.let {
                    ctx.vertexLighter = vanillaTuftLighting
                    ShadersModIntegration.grass(ctx, Config.reed.shaderWind) {
                        ctx.renderQuads(reedModels[it])
                    }
                }
            }
        } else super.renderLayer(ctx, data, layer)
    }

    companion object {
        val algaeSprites by SpriteSetDelegate(Atlas.BLOCKS) { idx ->
            ResourceLocation(BetterFoliageMod.MOD_ID, "blocks/better_algae_$idx")
        }
        val reedSprites by SpriteSetDelegate(
            Atlas.BLOCKS,
            idFunc = { idx -> ResourceLocation(BetterFoliageMod.MOD_ID, "blocks/better_reed_$idx") },
            idRegister = { id -> CenteredSprite(id, aspectHeight = 2).register(BetterFoliage.generatedPack) }
        )
        val algaeModels by LazyInvalidatable(BakeWrapperManager) {
            val shapes = Config.algae.let { tuftShapeSet(it.size, it.heightMin, it.heightMax, it.hOffset) }
            tuftModelSet(shapes, -1) { algaeSprites[randomI()] }.buildTufts()
        }
        val reedModels by LazyInvalidatable(BakeWrapperManager) {
            val shapes = Config.reed.let { tuftShapeSet(2.0, it.heightMin, it.heightMax, it.hOffset) }
            tuftModelSet(shapes, -1) { reedSprites[randomI()] }.buildTufts()
        }
    }
}