package com.nexlab.openclaw.node.service

import android.annotation.SuppressLint
import android.app.Activity
import android.content.Context
import android.content.Intent
import android.graphics.Bitmap
import android.graphics.PixelFormat
import android.hardware.display.DisplayManager
import android.hardware.display.VirtualDisplay
import android.media.ImageReader
import android.media.projection.MediaProjection
import android.media.projection.MediaProjectionManager
import android.os.Handler
import android.os.Looper
import android.util.Base64
import android.util.DisplayMetrics
import android.view.WindowManager
import java.io.ByteArrayOutputStream

/**
 * Service for capturing screenshots using MediaProjection API
 */
class ScreenshotService(private val context: Context) {

    private var mediaProjection: MediaProjection? = null
    private var virtualDisplay: VirtualDisplay? = null
    private var imageReader: ImageReader? = null
    private var isCapturing = false

    private val windowManager = context.getSystemService(Context.WINDOW_SERVICE) as WindowManager
    private val mainHandler = Handler(Looper.getMainLooper())

    private var callback: ((Result<String>) -> Unit)? = null

    /**
     * Initialize the screenshot service with a media projection
     */
    fun initialize(projection: MediaProjection) {
        this.mediaProjection = projection
        setupVirtualDisplay()
    }

    /**
     * Handle the activity result from requesting screen capture permission
     */
    fun onActivityResult(resultCode: Int, data: Intent?, onReady: () -> Unit) {
        if (resultCode != Activity.RESULT_OK || data == null) {
            callback?.invoke(Result.failure(Exception("Screen capture permission denied")))
            return
        }

        val projectionManager = context.getSystemService(Context.MEDIA_PROJECTION_SERVICE) as MediaProjectionManager
        mediaProjection = projectionManager.getMediaProjection(resultCode, data)
        
        mediaProjection?.registerCallback(object : MediaProjection.Callback() {
            override fun onStop() {
                cleanup()
            }
        }, mainHandler)

        setupVirtualDisplay()
        onReady()
    }

    @SuppressLint("WrongConstant")
    private fun setupVirtualDisplay() {
        val metrics = DisplayMetrics()
        windowManager.defaultDisplay.getRealMetrics(metrics)
        
        val density = metrics.densityDpi
        val width = metrics.widthPixels
        val height = metrics.heightPixels

        imageReader = ImageReader.newInstance(
            width, height,
            PixelFormat.RGBA_8888, 2
        )

        virtualDisplay = mediaProjection?.createVirtualDisplay(
            "OpenClawScreenshot",
            width, height, density,
            DisplayManager.VIRTUAL_DISPLAY_FLAG_AUTO_MIRROR,
            imageReader?.surface,
            null, mainHandler
        )

        isCapturing = true
    }

    /**
     * Capture a single screenshot
     */
    fun captureScreenshot(quality: Int = 80, onResult: (Result<String>) -> Unit) {
        callback = onResult
        
        if (!isCapturing || imageReader == null) {
            onResult(Result.failure(Exception("Screenshot service not initialized")))
            return
        }

        try {
            val image = imageReader?.acquireLatestImage()
            if (image == null) {
                onResult(Result.failure(Exception("Failed to acquire image")))
                return
            }

            val planes = image.planes
            val buffer = planes[0].buffer
            val pixelStride = planes[0].pixelStride
            val rowStride = planes[0].rowStride
            val rowPadding = rowStride - pixelStride * image.width

            val bitmap = Bitmap.createBitmap(
                image.width + rowPadding / pixelStride,
                image.height,
                Bitmap.Config.ARGB_8888
            )
            bitmap.copyPixelsFromBuffer(buffer)

            // Crop to actual screen size
            val croppedBitmap = Bitmap.createBitmap(
                bitmap, 0, 0,
                image.width, image.height
            )

            val outputStream = ByteArrayOutputStream()
            croppedBitmap.compress(Bitmap.CompressFormat.JPEG, quality, outputStream)
            val base64 = Base64.encodeToString(outputStream.toByteArray(), Base64.NO_WRAP)

            bitmap.recycle()
            croppedBitmap.recycle()
            image.close()

            onResult(Result.success(base64))

        } catch (e: Exception) {
            onResult(Result.failure(e))
        }
    }

    /**
     * Start continuous screenshot streaming
     */
    fun startStream(fps: Int = 5, onFrame: (String) -> Unit) {
        // Implementation for streaming mode
        // Would use a coroutine to capture at regular intervals
    }

    /**
     * Stop screenshot capture
     */
    fun stop() {
        isCapturing = false
        cleanup()
    }

    private fun cleanup() {
        virtualDisplay?.release()
        virtualDisplay = null
        imageReader?.close()
        imageReader = null
        mediaProjection?.stop()
        mediaProjection = null
    }

    companion object {
        const val REQUEST_CODE = 1001
    }
}
