@file:Suppress("WrongConstant")

package io.visus.orbis.feature.orb.service

import android.content.Context
import android.graphics.BitmapFactory
import android.os.Handler
import android.os.Looper
import android.view.Choreographer
import android.view.Surface
import androidx.webgpu.GPU
import androidx.webgpu.GPUAdapter
import androidx.webgpu.GPUBindGroup
import androidx.webgpu.GPUBindGroupDescriptor
import androidx.webgpu.GPUBindGroupEntry
import androidx.webgpu.GPUBindGroupLayout
import androidx.webgpu.GPUBindGroupLayoutDescriptor
import androidx.webgpu.GPUBindGroupLayoutEntry
import androidx.webgpu.GPUBuffer
import androidx.webgpu.GPUBufferBindingLayout
import androidx.webgpu.GPUBufferDescriptor
import androidx.webgpu.GPUColor
import androidx.webgpu.GPUColorTargetState
import androidx.webgpu.GPUDevice
import androidx.webgpu.GPUExtent3D
import androidx.webgpu.GPUFragmentState
import androidx.webgpu.GPUInstance
import androidx.webgpu.GPUOrigin3D
import androidx.webgpu.GPUPipelineLayout
import androidx.webgpu.GPUPipelineLayoutDescriptor
import androidx.webgpu.GPURenderPassColorAttachment
import androidx.webgpu.GPURenderPassDescriptor
import androidx.webgpu.GPURenderPipeline
import androidx.webgpu.GPURenderPipelineDescriptor
import androidx.webgpu.GPUSampler
import androidx.webgpu.GPUSamplerBindingLayout
import androidx.webgpu.GPUSamplerDescriptor
import androidx.webgpu.GPUShaderModule
import androidx.webgpu.GPUShaderModuleDescriptor
import androidx.webgpu.GPUShaderSourceWGSL
import androidx.webgpu.GPUSurface
import androidx.webgpu.GPUSurfaceConfiguration
import androidx.webgpu.GPUSurfaceDescriptor
import androidx.webgpu.GPUSurfaceSourceAndroidNativeWindow
import androidx.webgpu.GPUTexelCopyBufferLayout
import androidx.webgpu.GPUTexelCopyTextureInfo
import androidx.webgpu.GPUTexture
import androidx.webgpu.GPUTextureBindingLayout
import androidx.webgpu.GPUTextureDescriptor
import androidx.webgpu.GPUTextureView
import androidx.webgpu.GPUTextureViewDescriptor
import androidx.webgpu.GPUVertexState
import androidx.webgpu.helper.Util
import androidx.webgpu.helper.initLibrary
import io.visus.orbis.feature.orb.data.OrbRenderState
import io.visus.orbis.feature.orb.util.CubemapGenerator
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import java.nio.ByteBuffer
import java.nio.ByteOrder

/**
 * WebGPU-based implementation of [OrbRenderService].
 *
 * This implementation renders a procedural orb effect using WGSL shaders with support for:
 * - Noise texture sampling for surface distortion
 * - Cubemap environment reflections
 * - Dynamic color and amplitude parameters via uniform buffers
 *
 * The renderer uses a reduced resolution (60% by default) for performance optimization
 * while maintaining the correct aspect ratio. Frame timing is synchronized with the
 * display's VSync via [android.view.Choreographer].
 *
 * GPU resources are allocated during [initialize] and [attachSurface], and must be
 * explicitly released via [release] to avoid memory leaks.
 */
class OrbRenderServiceImpl : OrbRenderService {

    private val _isInitialized = MutableStateFlow(false)
    override val isInitialized: StateFlow<Boolean> = _isInitialized.asStateFlow()

    private val _error = MutableStateFlow<String?>(null)
    override val error: StateFlow<String?> = _error.asStateFlow()

    private val _fps = MutableStateFlow(0)
    override val fps: StateFlow<Int> = _fps.asStateFlow()

    private var instance: GPUInstance? = null
    private var adapter: GPUAdapter? = null
    private var device: GPUDevice? = null
    private var gpuSurface: GPUSurface? = null
    private var shaderModule: GPUShaderModule? = null
    private var uniformBuffer: GPUBuffer? = null
    private var noiseTexture: GPUTexture? = null
    private var noiseTextureView: GPUTextureView? = null
    private var cubemapTexture: GPUTexture? = null
    private var cubemapTextureView: GPUTextureView? = null
    private var sampler: GPUSampler? = null
    private var bindGroupLayout: GPUBindGroupLayout? = null
    private var pipelineLayout: GPUPipelineLayout? = null
    private var renderPipeline: GPURenderPipeline? = null
    private var bindGroup: GPUBindGroup? = null

    // State
    @Volatile private var currentState = OrbRenderState()
    private var surfaceFormat: Int = TEXTURE_FORMAT_BGRA8_UNORM
    private var surfaceWidth: Int = 0
    private var surfaceHeight: Int = 0
    private var renderWidth: Int = 0
    private var renderHeight: Int = 0
    private var startTime: Long = 0L
    @Volatile private var renderLoopRunning = false

    private var choreographer: Choreographer? = null
    private var choreographerHandler: Handler? = null
    private val frameCallback = object : Choreographer.FrameCallback {
        override fun doFrame(frameTimeNanos: Long) {
            if (renderLoopRunning) {
                render()
                updateFpsCounter()
                choreographer?.postFrameCallback(this)
            }
        }
    }

    private var frameCount = 0
    private var lastFpsUpdateTime = 0L

    // Reusable uniform buffer to avoid allocations per frame
    private val uniformData = ByteBuffer.allocateDirect(48).order(ByteOrder.nativeOrder())

    // Pre-allocated render objects to avoid per-frame allocations
    private val clearColor = GPUColor(0.0, 0.0, 0.0, 0.0)

    companion object {
        private const val UNIFORM_BUFFER_SIZE = 48L

        // Resolution scale for performance (1.0 = full, 0.5 = half resolution)
        private const val RESOLUTION_SCALE = 0.6f

        // WebGPU constants (raw values to avoid RestrictedApi warnings)
        private const val TEXTURE_FORMAT_RGBA8_UNORM = 0x00000016
        private const val TEXTURE_FORMAT_BGRA8_UNORM = 0x0000001b
        private const val TEXTURE_USAGE_COPY_DST = 0x00000002
        private const val TEXTURE_USAGE_TEXTURE_BINDING = 0x00000004
        private const val TEXTURE_USAGE_RENDER_ATTACHMENT = 0x00000010
        private const val BUFFER_USAGE_COPY_DST = 0x00000008
        private const val BUFFER_USAGE_UNIFORM = 0x00000040
        private const val SHADER_STAGE_VERTEX = 0x00000001
        private const val SHADER_STAGE_FRAGMENT = 0x00000002
        private const val BUFFER_BINDING_TYPE_UNIFORM = 0x00000002
        private const val TEXTURE_SAMPLE_TYPE_FLOAT = 0x00000002
        private const val TEXTURE_VIEW_DIMENSION_2D = 0x00000002
        private const val TEXTURE_VIEW_DIMENSION_CUBE = 0x00000004
        private const val TEXTURE_DIMENSION_2D = 0x00000002
        private const val SAMPLER_BINDING_TYPE_FILTERING = 0x00000002
        private const val ADDRESS_MODE_REPEAT = 0x00000002
        private const val FILTER_MODE_LINEAR = 0x00000002
        private const val LOAD_OP_CLEAR = 0x00000002
        private const val STORE_OP_STORE = 0x00000001

        // Align to 256 bytes (WebGPU requirement for bytesPerRow)
        private fun alignTo256(value: Int): Int = (value + 255) and 0xFF.inv()
    }

    override suspend fun initialize(context: Context) {
        try {
            initLibrary()

            instance = GPU.createInstance()
            adapter = instance?.requestAdapter()
            device = adapter?.requestDevice()

            if (device == null) {
                _error.value = "Failed to create GPU device"
                return
            }

            loadShader(context)
            createUniformBuffer()
            loadNoiseTexture(context)
            createCubemapTexture()
            createSampler()
            createBindGroupLayout()
            createPipelineLayout()

            startTime = System.nanoTime()
            _isInitialized.value = true
        } catch (e: Exception) {
            _error.value = "Initialization failed: ${e.message}"
        }
    }

    override suspend fun attachSurface(surface: Surface, width: Int, height: Int) {
        val dev = device ?: return
        val inst = instance ?: return

        surfaceWidth = width
        surfaceHeight = height
        renderWidth = (width * RESOLUTION_SCALE).toInt().coerceAtLeast(1)
        renderHeight = (height * RESOLUTION_SCALE).toInt().coerceAtLeast(1)

        gpuSurface = inst.createSurface(
            GPUSurfaceDescriptor(
                surfaceSourceAndroidNativeWindow = GPUSurfaceSourceAndroidNativeWindow(
                    window = Util.windowFromSurface(surface)
                )
            )
        )

        val caps = gpuSurface?.getCapabilities(adapter!!)
        surfaceFormat = caps?.formats?.firstOrNull() ?: TEXTURE_FORMAT_BGRA8_UNORM

        gpuSurface?.configure(
            GPUSurfaceConfiguration(
                device = dev,
                width = renderWidth,
                height = renderHeight,
                format = surfaceFormat,
                usage = TEXTURE_USAGE_RENDER_ATTACHMENT
            )
        )

        createRenderPipeline()
        createBindGroup()
    }

    override fun detachSurface() {
        stopRenderLoop()
        gpuSurface?.unconfigure()
        gpuSurface?.close()
        gpuSurface = null
    }

    override fun updateRenderState(state: OrbRenderState) {
        currentState = state
    }

    override fun render() {
        val dev = device ?: return
        val surf = gpuSurface ?: return
        val pipeline = renderPipeline ?: return
        val bg = bindGroup ?: return

        try {
            updateUniforms()

            val surfaceTexture = surf.getCurrentTexture()
            val textureView = surfaceTexture.texture.createView()
            val commandEncoder = dev.createCommandEncoder()

            val renderPass = commandEncoder.beginRenderPass(
                GPURenderPassDescriptor(
                    colorAttachments = arrayOf(
                        GPURenderPassColorAttachment(
                            view = textureView,
                            loadOp = LOAD_OP_CLEAR,
                            storeOp = STORE_OP_STORE,
                            clearValue = clearColor
                        )
                    )
                )
            )

            renderPass.setPipeline(pipeline)
            renderPass.setBindGroup(0, bg)
            renderPass.draw(6)
            renderPass.end()

            val commandBuffer = commandEncoder.finish()
            dev.queue.submit(arrayOf(commandBuffer))

            // Close command buffer immediately after submit - GPU has its own copy
            commandBuffer.close()
            commandEncoder.close()
            renderPass.close()

            surf.present()
            textureView.close()
        } catch (_: Exception) {
            // Continue on error
        }
    }

    override fun startRenderLoop() {
        if (renderLoopRunning) return

        renderLoopRunning = true
        frameCount = 0
        lastFpsUpdateTime = System.nanoTime()

        choreographerHandler = Handler(Looper.getMainLooper())
        choreographerHandler?.post {
            choreographer = Choreographer.getInstance()
            choreographer?.postFrameCallback(frameCallback)
        }
    }

    override fun stopRenderLoop() {
        renderLoopRunning = false
        choreographerHandler?.post {
            choreographer?.removeFrameCallback(frameCallback)
        }
        choreographerHandler = null
        choreographer = null
    }

    private fun updateFpsCounter() {
        frameCount++

        val now = System.nanoTime()
        val elapsed = now - lastFpsUpdateTime

        if (elapsed >= 1_000_000_000L) {
            _fps.value = (frameCount * 1_000_000_000L / elapsed).toInt()
            frameCount = 0
            lastFpsUpdateTime = now
        }
    }

    override fun release() {
        stopRenderLoop()

        bindGroup?.close()
        renderPipeline?.close()
        pipelineLayout?.close()
        bindGroupLayout?.close()
        sampler?.close()
        cubemapTextureView?.close()
        cubemapTexture?.close()
        noiseTextureView?.close()
        noiseTexture?.close()
        uniformBuffer?.close()
        shaderModule?.close()
        gpuSurface?.close()
        device?.close()
        adapter?.close()
        instance?.close()

        _isInitialized.value = false
    }

    private fun loadShader(context: Context) {
        val shaderCode = context.assets.open("shaders/Orb.wgsl").bufferedReader().use {
            it.readText()
        }
        shaderModule = device?.createShaderModule(
            GPUShaderModuleDescriptor(
                shaderSourceWGSL = GPUShaderSourceWGSL(code = shaderCode)
            )
        )
    }

    private fun createUniformBuffer() {
        uniformBuffer = device?.createBuffer(
            GPUBufferDescriptor(
                usage = BUFFER_USAGE_UNIFORM or BUFFER_USAGE_COPY_DST,
                size = UNIFORM_BUFFER_SIZE
            )
        )
    }

    private fun loadNoiseTexture(context: Context) {
        val bitmap = context.assets.open("textures/noise_map.png").use { inputStream ->
            BitmapFactory.decodeStream(inputStream)
        }

        val width = bitmap.width
        val height = bitmap.height
        val bytesPerPixel = 4
        val bytesPerRow = alignTo256(width * bytesPerPixel)

        noiseTexture = device?.createTexture(
            GPUTextureDescriptor(
                usage = TEXTURE_USAGE_TEXTURE_BINDING or TEXTURE_USAGE_COPY_DST,
                size = GPUExtent3D(width, height, 1),
                format = TEXTURE_FORMAT_RGBA8_UNORM
            )
        )

        val pixels = IntArray(width * height)
        bitmap.getPixels(pixels, 0, width, 0, 0, width, height)

        // Create buffer with aligned row size
        val buffer = ByteBuffer.allocateDirect(bytesPerRow * height).order(ByteOrder.nativeOrder())

        for (y in 0 until height) {
            for (x in 0 until width) {
                val pixel = pixels[y * width + x]

                buffer.put((pixel shr 16 and 0xFF).toByte()) // R
                buffer.put((pixel shr 8 and 0xFF).toByte())  // G
                buffer.put((pixel and 0xFF).toByte())        // B
                buffer.put((pixel shr 24 and 0xFF).toByte()) // A
            }

            // Pad to aligned row size
            val padding = bytesPerRow - (width * bytesPerPixel)

            for (i in 0 until padding) {
                buffer.put(0)
            }
        }
        buffer.rewind()

        device?.queue?.writeTexture(
            destination = GPUTexelCopyTextureInfo(texture = noiseTexture!!),
            data = buffer,
            writeSize = GPUExtent3D(width, height, 1),
            dataLayout = GPUTexelCopyBufferLayout(
                offset = 0,
                bytesPerRow = bytesPerRow,
                rowsPerImage = height
            )
        )

        noiseTextureView = noiseTexture?.createView()

        bitmap.recycle()
    }

    private fun createCubemapTexture() {
        val faceSize = CubemapGenerator.getFaceSize()
        val bytesPerPixel = 4
        val bytesPerRow = alignTo256(faceSize * bytesPerPixel)

        cubemapTexture = device?.createTexture(
            GPUTextureDescriptor(
                usage = TEXTURE_USAGE_TEXTURE_BINDING or TEXTURE_USAGE_COPY_DST,
                size = GPUExtent3D(faceSize, faceSize, 6),
                format = TEXTURE_FORMAT_RGBA8_UNORM,
                dimension = TEXTURE_DIMENSION_2D
            )
        )

        val cubemapData = CubemapGenerator.generate()
        val sourceBytesPerRow = faceSize * bytesPerPixel

        for (face in 0 until 6) {
            // Create buffer with aligned row size
            val faceBuffer = ByteBuffer.allocateDirect(bytesPerRow * faceSize).order(ByteOrder.nativeOrder())

            for (y in 0 until faceSize) {
                // Copy one row from source
                cubemapData.position(face * faceSize * sourceBytesPerRow + y * sourceBytesPerRow)
                for (x in 0 until sourceBytesPerRow) {
                    faceBuffer.put(cubemapData.get())
                }
                // Pad to aligned row size
                val padding = bytesPerRow - sourceBytesPerRow
                for (i in 0 until padding) {
                    faceBuffer.put(0)
                }
            }
            faceBuffer.rewind()

            device?.queue?.writeTexture(
                destination = GPUTexelCopyTextureInfo(
                    texture = cubemapTexture!!,
                    origin = GPUOrigin3D(0, 0, face)
                ),
                data = faceBuffer,
                writeSize = GPUExtent3D(faceSize, faceSize, 1),
                dataLayout = GPUTexelCopyBufferLayout(
                    offset = 0,
                    bytesPerRow = bytesPerRow,
                    rowsPerImage = faceSize
                )
            )
        }

        cubemapTextureView = cubemapTexture?.createView(
            GPUTextureViewDescriptor(
                usage = TEXTURE_USAGE_TEXTURE_BINDING,
                dimension = TEXTURE_VIEW_DIMENSION_CUBE,
                arrayLayerCount = 6
            )
        )
    }

    private fun createSampler() {
        sampler = device?.createSampler(
            GPUSamplerDescriptor(
                addressModeU = ADDRESS_MODE_REPEAT,
                addressModeV = ADDRESS_MODE_REPEAT,
                addressModeW = ADDRESS_MODE_REPEAT,
                magFilter = FILTER_MODE_LINEAR,
                minFilter = FILTER_MODE_LINEAR
            )
        )
    }

    private fun createBindGroupLayout() {
        bindGroupLayout = device?.createBindGroupLayout(
            GPUBindGroupLayoutDescriptor(
                entries = arrayOf(
                    GPUBindGroupLayoutEntry(
                        binding = 0,
                        visibility = SHADER_STAGE_VERTEX or SHADER_STAGE_FRAGMENT,
                        buffer = GPUBufferBindingLayout(type = BUFFER_BINDING_TYPE_UNIFORM)
                    ),
                    GPUBindGroupLayoutEntry(
                        binding = 1,
                        visibility = SHADER_STAGE_FRAGMENT,
                        texture = GPUTextureBindingLayout(
                            sampleType = TEXTURE_SAMPLE_TYPE_FLOAT,
                            viewDimension = TEXTURE_VIEW_DIMENSION_2D
                        )
                    ),
                    GPUBindGroupLayoutEntry(
                        binding = 2,
                        visibility = SHADER_STAGE_FRAGMENT,
                        texture = GPUTextureBindingLayout(
                            sampleType = TEXTURE_SAMPLE_TYPE_FLOAT,
                            viewDimension = TEXTURE_VIEW_DIMENSION_CUBE
                        )
                    ),
                    GPUBindGroupLayoutEntry(
                        binding = 3,
                        visibility = SHADER_STAGE_FRAGMENT,
                        sampler = GPUSamplerBindingLayout(type = SAMPLER_BINDING_TYPE_FILTERING)
                    )
                )
            )
        )
    }

    private fun createPipelineLayout() {
        pipelineLayout = device?.createPipelineLayout(
            GPUPipelineLayoutDescriptor(bindGroupLayouts = arrayOf(bindGroupLayout!!))
        )
    }

    private fun createRenderPipeline() {
        val shader = shaderModule ?: return
        val layout = pipelineLayout ?: return

        renderPipeline = device?.createRenderPipeline(
            GPURenderPipelineDescriptor(
                layout = layout,
                vertex = GPUVertexState(module = shader, entryPoint = "vs_main"),
                fragment = GPUFragmentState(
                    module = shader,
                    entryPoint = "fs_main",
                    targets = arrayOf(GPUColorTargetState(format = surfaceFormat))
                )
            )
        )
    }

    private fun createBindGroup() {
        val ub = uniformBuffer ?: return
        val noiseView = noiseTextureView ?: return
        val cubeView = cubemapTextureView ?: return
        val samp = sampler ?: return
        val layout = bindGroupLayout ?: return

        bindGroup = device?.createBindGroup(
            GPUBindGroupDescriptor(
                layout = layout,
                entries = arrayOf(
                    GPUBindGroupEntry(binding = 0, buffer = ub, size = UNIFORM_BUFFER_SIZE),
                    GPUBindGroupEntry(binding = 1, textureView = noiseView),
                    GPUBindGroupEntry(binding = 2, textureView = cubeView),
                    GPUBindGroupEntry(binding = 3, sampler = samp)
                )
            )
        )
    }

    private fun updateUniforms() {
        val ub = uniformBuffer ?: return
        val dev = device ?: return

        val elapsedTime = (System.nanoTime() - startTime) / 1_000_000_000f
        // Use original aspect ratio for correct projection
        val aspectRatio = if (surfaceHeight > 0) surfaceWidth.toFloat() / surfaceHeight else 1f

        // Reuse the pre-allocated buffer
        uniformData.clear()

        // resolution (vec2<f32>) - offset 0 (use render resolution)
        uniformData.putFloat(renderWidth.toFloat())
        uniformData.putFloat(renderHeight.toFloat())

        // time (f32) - offset 8
        uniformData.putFloat(elapsedTime)

        // amplitude (f32) - offset 12
        uniformData.putFloat(currentState.amplitude)

        // color (vec4<f32>) - offset 16 (16-byte aligned)
        uniformData.putFloat(currentState.color.red)
        uniformData.putFloat(currentState.color.green)
        uniformData.putFloat(currentState.color.blue)
        uniformData.putFloat(currentState.color.alpha)

        // aspectRatio (f32) - offset 32
        uniformData.putFloat(aspectRatio)

        // Padding to 48 bytes
        uniformData.putFloat(0f)
        uniformData.putFloat(0f)
        uniformData.putFloat(0f)

        uniformData.rewind()

        dev.queue.writeBuffer(ub, 0, uniformData)
    }
}
