package com.guiagent.ocr

import android.graphics.BitmapFactory
import com.google.gson.Gson
import fi.iki.elonen.NanoHTTPD
import java.io.ByteArrayOutputStream

class OcrHttpServer(port: Int = 18900) : NanoHTTPD(port) {

    private val gson = Gson()
    private val defaultPath = "/sdcard/ocr_screen.png"

    override fun serve(session: IHTTPSession): Response {
        return when (session.uri) {
            "/ocr" -> handleOcr(session)
            "/snap" -> handleSnap(session)
            "/health" -> jsonResponse(mapOf("status" to "ok", "engine" to "mlkit-chinese"))
            else -> newFixedLengthResponse(Response.Status.NOT_FOUND, MIME_PLAINTEXT, "404")
        }
    }

    /** 读文件方式 OCR */
    private fun handleOcr(session: IHTTPSession): Response {
        val params = session.parms ?: emptyMap()
        val imagePath = params["path"] ?: defaultPath
        return doOcr(params["text"]) { OcrEngine.recognize(imagePath) }
    }

    /** POST 图片数据直接 OCR，不存文件 */
    private fun handleSnap(session: IHTTPSession): Response {
        val params = session.parms ?: emptyMap()

        if (session.method == Method.POST) {
            // NanoHTTPD parseBody 将 binary data 存到临时文件
            val bodyFiles = HashMap<String, String>()
            session.parseBody(bodyFiles)

            // postData 键对应临时文件路径
            val tmpPath = bodyFiles["postData"]
            if (tmpPath != null) {
                val imageBytes = java.io.File(tmpPath).readBytes()
                val bitmap = BitmapFactory.decodeByteArray(imageBytes, 0, imageBytes.size)
                if (bitmap != null) {
                    return doOcr(params["text"]) { OcrEngine.recognizeBitmap(bitmap) }
                }
                return jsonResponse(mapOf("error" to "decode failed", "size" to imageBytes.size, "count" to 0))
            }
            return jsonResponse(mapOf("error" to "no body received", "count" to 0))
        }

        // GET: 读文件方式 fallback
        return handleOcr(session)
    }

    private fun doOcr(query: String?, recognize: () -> List<TextBox>): Response {
        val startTime = System.currentTimeMillis()
        var results = recognize()

        if (!query.isNullOrBlank()) {
            results = results.filter { it.text.contains(query) }
        }

        val elapsed = System.currentTimeMillis() - startTime

        val response = mapOf(
            "results" to results.map { box ->
                mapOf(
                    "text" to box.text,
                    "x" to box.x,
                    "y" to box.y,
                    "w" to box.w,
                    "h" to box.h,
                    "cx" to box.cx,
                    "cy" to box.cy,
                    "confidence" to box.confidence
                )
            },
            "count" to results.size,
            "elapsed_ms" to elapsed
        )
        return jsonResponse(response)
    }

    private fun jsonResponse(data: Any): Response {
        val json = gson.toJson(data)
        return newFixedLengthResponse(Response.Status.OK, "application/json", json)
    }
}