package com.omixlab.lckcontrol.cortex

import android.content.Context
import android.graphics.Bitmap
import android.hardware.HardwareBuffer
import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.media.MediaMetadataRetriever
import android.media.MediaMuxer
import android.util.Log
import com.omixlab.lckcontrol.streaming.NativeStreamingEngine
import dagger.hilt.android.qualifiers.ApplicationContext
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.launch
import java.io.File
import java.io.FileOutputStream
import java.util.UUID
import javax.inject.Inject
import javax.inject.Singleton

@Singleton
class CortexManager @Inject constructor(
    @ApplicationContext private val context: Context,
    private val cortexPreferences: CortexPreferences,
) {
    companion object {
        private const val TAG = "CortexManager"
        private const val CORTEX_DIR = "cortex"
        private const val VIDEO_BITRATE = 8_000_000
        private const val AUDIO_BITRATE = 128_000
        private const val SAMPLE_RATE = 48_000
        private const val CHANNELS = 2
        private const val KEYFRAME_INTERVAL = 2
    }

    private val scope = CoroutineScope(SupervisorJob() + Dispatchers.IO)

    private var cortexEngine: NativeStreamingEngine? = null
    private var currentSessionId: String? = null
    private var poolWidth = 0
    private var poolHeight = 0

    private val _sessions = MutableStateFlow<List<CortexSession>>(emptyList())
    val sessions: StateFlow<List<CortexSession>> = _sessions.asStateFlow()

    private val _isRecording = MutableStateFlow(false)
    val isRecording: StateFlow<Boolean> = _isRecording.asStateFlow()

    private val _storageUsedBytes = MutableStateFlow(0L)
    val storageUsedBytes: StateFlow<Long> = _storageUsedBytes.asStateFlow()

    var onBufferReleased: ((Int) -> Unit)? = null

    val hasCortexEngine: Boolean get() = cortexEngine != null

    private fun cortexBaseDir(): File = File(context.filesDir, CORTEX_DIR)

    fun onTexturePoolRegistered(width: Int, height: Int) {
        poolWidth = width
        poolHeight = height
        if (cortexPreferences.isEnabled()) {
            startCortexOnlyEngine()
        }
    }

    fun onTexturePoolUnregistered() {
        stopCortexEngine()
        poolWidth = 0
        poolHeight = 0
    }

    fun onStreamingStarting(engine: NativeStreamingEngine) {
        // Stop cortex-only engine — streaming engine takes over
        stopCortexEngine()

        // Enable cortex recording on the streaming engine
        if (cortexPreferences.isEnabled()) {
            val sessionDir = getOrCreateSessionDir()
            engine.enableCortexRecording(sessionDir.absolutePath, cortexPreferences.getMaxMinutes())
            engine.onCortexSegment = { segPath, keyframeData ->
                scope.launch { handleSegmentReady(segPath, keyframeData) }
            }
            _isRecording.value = true
            Log.i(TAG, "Cortex enabled on streaming engine: ${sessionDir.absolutePath}")
        }
    }

    fun onStreamingStopped() {
        // Restart cortex-only engine if we have a texture pool and cortex is enabled
        if (poolWidth > 0 && poolHeight > 0 && cortexPreferences.isEnabled()) {
            startCortexOnlyEngine()
        }
    }

    fun submitVideoFrame(buffer: HardwareBuffer, timestampNs: Long, fenceFd: Int, bufferIndex: Int) {
        cortexEngine?.submitVideoFrame(buffer, timestampNs, fenceFd, bufferIndex)
    }

    fun submitAudioFrame(pcmData: ByteArray, timestampNs: Long) {
        cortexEngine?.submitAudioFrame(pcmData, timestampNs)
    }

    fun deleteSession(sessionId: String) {
        val dir = File(cortexBaseDir(), sessionId)
        if (dir.exists()) {
            dir.deleteRecursively()
        }
        refreshSessions()
    }

    fun deleteAllSessions() {
        // Don't delete the current active session
        val activeId = currentSessionId
        cortexBaseDir().listFiles()?.forEach { dir ->
            if (dir.isDirectory && dir.name != activeId) {
                dir.deleteRecursively()
            }
        }
        refreshSessions()
    }

    fun refreshSessions() {
        scope.launch {
            val baseDir = cortexBaseDir()
            if (!baseDir.exists()) {
                _sessions.value = emptyList()
                _storageUsedBytes.value = 0L
                return@launch
            }

            var totalSize = 0L
            val sessionList = baseDir.listFiles()
                ?.filter { it.isDirectory }
                ?.mapNotNull { dir ->
                    val segFiles = dir.listFiles { f -> f.name.endsWith(".seg") } ?: emptyArray()
                    if (segFiles.isEmpty()) return@mapNotNull null

                    val thumbnails = dir.listFiles { f -> f.name.endsWith(".jpg") }
                        ?.sortedBy { it.name }
                        ?.map { it.absolutePath }
                        ?: emptyList()

                    val size = dir.listFiles()?.sumOf { it.length() } ?: 0L
                    totalSize += size

                    CortexSession(
                        sessionId = dir.name,
                        directory = dir,
                        segmentCount = segFiles.size,
                        thumbnailPaths = thumbnails,
                        totalSizeBytes = size,
                        startTime = segFiles.minOf { it.lastModified() },
                    )
                }
                ?.sortedByDescending { it.startTime }
                ?: emptyList()

            _sessions.value = sessionList
            _storageUsedBytes.value = totalSize
        }
    }

    private fun startCortexOnlyEngine() {
        if (cortexEngine != null) return
        if (poolWidth <= 0 || poolHeight <= 0) return

        Log.i(TAG, "Starting cortex-only engine: ${poolWidth}x${poolHeight}")

        val eng = NativeStreamingEngine()
        eng.create(
            width = poolWidth,
            height = poolHeight,
            videoBitrate = VIDEO_BITRATE,
            audioBitrate = AUDIO_BITRATE,
            sampleRate = SAMPLE_RATE,
            channels = CHANNELS,
            keyframeInterval = KEYFRAME_INTERVAL,
        )

        // No addDestination calls — zero sinks

        eng.onBufferReleased = { index ->
            onBufferReleased?.invoke(index)
        }

        val sessionDir = getOrCreateSessionDir()
        eng.enableCortexRecording(sessionDir.absolutePath, cortexPreferences.getMaxMinutes())

        eng.onCortexSegment = { segPath, keyframeData ->
            scope.launch { handleSegmentReady(segPath, keyframeData) }
        }

        if (eng.start()) {
            cortexEngine = eng
            _isRecording.value = true
            Log.i(TAG, "Cortex-only engine started")
        } else {
            eng.destroy()
            Log.e(TAG, "Failed to start cortex-only engine")
        }
    }

    private fun stopCortexEngine() {
        val eng = cortexEngine ?: return
        Log.i(TAG, "Stopping cortex-only engine")
        eng.disableCortexRecording()
        eng.stop()
        eng.destroy()
        cortexEngine = null
        _isRecording.value = false
        refreshSessions()
    }

    private fun getOrCreateSessionDir(): File {
        val id = currentSessionId ?: UUID.randomUUID().toString().also { currentSessionId = it }
        val dir = File(cortexBaseDir(), id)
        dir.mkdirs()
        return dir
    }

    private suspend fun handleSegmentReady(segPath: String, keyframeData: ByteArray) {
        Log.i(TAG, "Segment ready: $segPath")
        generateThumbnail(segPath, keyframeData)
        refreshSessions()
    }

    private fun generateThumbnail(segPath: String, keyframeData: ByteArray) {
        try {
            val segFile = File(segPath)
            val segName = segFile.nameWithoutExtension // e.g. seg_000000
            val segIndex = segName.removePrefix("seg_")
            val thumbFile = File(segFile.parentFile, "thumb_$segIndex.jpg")

            // Mux the single keyframe into a temp MP4 so MediaMetadataRetriever can decode it
            val tempMp4 = File(segFile.parentFile, "thumb_temp_$segIndex.mp4")
            try {
                val muxer = MediaMuxer(tempMp4.absolutePath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4)

                val format = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, poolWidth, poolHeight)

                // Parse SPS/PPS from the .seg file header
                // For now, use the keyframe data directly — the codec config should be set on the format
                // We need to read the SPS/PPS from the segment file
                val spsPps = readSpsPpsFromSegment(segPath)
                if (spsPps != null) {
                    // Split SPS/PPS (Annex-B: [00 00 00 01 SPS] [00 00 00 01 PPS])
                    val startCode = byteArrayOf(0, 0, 0, 1)
                    var ppsOffset = -1
                    for (i in 4 until spsPps.size - 3) {
                        if (spsPps[i] == 0.toByte() && spsPps[i + 1] == 0.toByte() &&
                            spsPps[i + 2] == 0.toByte() && spsPps[i + 3] == 1.toByte()
                        ) {
                            ppsOffset = i
                            break
                        }
                    }
                    if (ppsOffset > 0) {
                        format.setByteBuffer("csd-0", java.nio.ByteBuffer.wrap(spsPps, 0, ppsOffset))
                        format.setByteBuffer("csd-1", java.nio.ByteBuffer.wrap(spsPps, ppsOffset, spsPps.size - ppsOffset))
                    } else {
                        format.setByteBuffer("csd-0", java.nio.ByteBuffer.wrap(spsPps))
                    }
                }

                val trackIndex = muxer.addTrack(format)
                muxer.start()

                val bufferInfo = MediaCodec.BufferInfo().apply {
                    offset = 0
                    size = keyframeData.size
                    presentationTimeUs = 0
                    flags = MediaCodec.BUFFER_FLAG_KEY_FRAME
                }
                muxer.writeSampleData(trackIndex, java.nio.ByteBuffer.wrap(keyframeData), bufferInfo)

                muxer.stop()
                muxer.release()

                // Extract frame with MediaMetadataRetriever
                val retriever = MediaMetadataRetriever()
                retriever.setDataSource(tempMp4.absolutePath)
                val bitmap = retriever.getFrameAtTime(0, MediaMetadataRetriever.OPTION_CLOSEST_SYNC)
                retriever.release()

                if (bitmap != null) {
                    FileOutputStream(thumbFile).use { fos ->
                        bitmap.compress(Bitmap.CompressFormat.JPEG, 80, fos)
                    }
                    bitmap.recycle()
                    Log.i(TAG, "Thumbnail generated: ${thumbFile.absolutePath}")
                }
            } finally {
                tempMp4.delete()
            }
        } catch (e: Exception) {
            Log.e(TAG, "Failed to generate thumbnail for $segPath", e)
        }
    }

    private fun readSpsPpsFromSegment(segPath: String): ByteArray? {
        try {
            val file = File(segPath)
            val raf = java.io.RandomAccessFile(file, "r")
            // Magic(4) + version(4) + width(4) + height(4) = 16 bytes
            raf.seek(16)
            val spsPpsSize = Integer.reverseBytes(raf.readInt()).let {
                // The file is written in native byte order (little-endian on ARM)
                // RandomAccessFile.readInt() reads big-endian, so we need to handle this
                // Actually, our writeVal writes raw bytes, so we should read raw bytes too
                raf.seek(16)
                val buf = ByteArray(4)
                raf.read(buf)
                (buf[0].toInt() and 0xFF) or
                    ((buf[1].toInt() and 0xFF) shl 8) or
                    ((buf[2].toInt() and 0xFF) shl 16) or
                    ((buf[3].toInt() and 0xFF) shl 24)
            }
            if (spsPpsSize <= 0 || spsPpsSize > 1024) {
                raf.close()
                return null
            }
            val data = ByteArray(spsPpsSize)
            raf.read(data)
            raf.close()
            return data
        } catch (e: Exception) {
            Log.e(TAG, "Failed to read SPS/PPS from segment", e)
            return null
        }
    }
}
