package com.optix.android.ai

import com.google.gson.Gson
import com.google.gson.JsonArray
import com.google.gson.JsonObject
import com.optix.android.domain.model.ToolChain
import com.optix.android.domain.model.ToolChainStep
import com.optix.android.tools.registry.ToolRegistry
import timber.log.Timber
import javax.inject.Inject
import javax.inject.Singleton

/**
 * Parses AI assistant output into a structured [ToolChain].
 *
 * The AI can emit a chain in one of two forms:
 *  1. A fenced JSON block tagged ```chain ... ``` (or a bare JSON array) whose shape is
 *     `{ "name": "...", "steps": [ { "tool": "slug", "params": { ... } }, ... ] }`.
 *  2. A structured text block of lines like:
 *       STEP 1: lan_scanner  subnet=192.168.1.0/24
 *       STEP 2: port_scanner target=$TARGET start_port=1 end_port=1024
 *
 * Every step's [ToolChainStep.toolSlug] is validated against [ToolRegistry]; steps that
 * reference an unknown tool are skipped and reported.
 */
@Singleton
class ChainParser @Inject constructor(
    private val toolRegistry: ToolRegistry,
    private val gson: Gson,
) {

    data class ParseResult(
        val chain: ToolChain?,
        val skippedSteps: List<String>,
    )

    fun parse(text: String): ParseResult {
        val jsonChain = extractJsonChain(text)
        val chain = jsonChain ?: parseStructuredSteps(text)
        if (chain == null) return ParseResult(null, emptyList())

        val validSlugs = toolRegistry.listNames().toSet()
        val steps = mutableListOf<ToolChainStep>()
        val skipped = mutableListOf<String>()

        chain.steps.forEach { step ->
            if (step.toolSlug in validSlugs) {
                steps.add(step)
            } else {
                skipped.add(step.toolSlug)
                Timber.w("ChainParser: skipping step with unknown tool slug '%s'", step.toolSlug)
            }
        }

        if (steps.isEmpty()) return ParseResult(null, skipped)
        return ParseResult(chain.copy(steps = steps), skipped)
    }

    // ------------------------------------------------------------------
    // JSON extraction
    // ------------------------------------------------------------------

    private fun extractJsonChain(text: String): ToolChain? {
        // Prefer a fenced ```chain block, fall back to a bare JSON object/array.
        val fenced = Regex("""(?s)```(?:json|chain)?\s*(\{.*?\}|\[.*?\])\s*```""")
            .find(text)?.groupValues?.get(1)
        val candidate = fenced
            ?: Regex("""(\{\s*"(?:name|steps)"|\{\s*"tool")""").find(text)?.let {
                text.substring(it.range.first)
            }

        if (candidate == null) return null
        return try {
            val trimmed = candidate.trim().removeSuffix("```").trim()
            if (trimmed.startsWith("[")) {
                val steps = parseJsonArray(gson.fromJson(trimmed, JsonArray::class.java)) ?: return null
                ToolChain(steps = steps, name = "AI Chain")
            } else {
                parseJsonObject(gson.fromJson(trimmed, JsonObject::class.java))
            }
        } catch (e: Exception) {
            Timber.w(e, "ChainParser: failed to parse JSON chain")
            null
        }
    }

    private fun parseJsonObject(obj: JsonObject): ToolChain? {
        val stepsArray = obj.getAsJsonArray("steps") ?: return null
        val name = obj.get("name")?.takeIf { !it.isJsonNull }?.asString ?: "AI Chain"
        val steps = parseJsonArray(stepsArray) ?: return null
        return ToolChain(steps = steps, name = name)
    }

    private fun parseJsonArray(arr: JsonArray): List<ToolChainStep>? {
        val steps = arr.mapNotNull { elem ->
            if (!elem.isJsonObject) return@mapNotNull null
            val o = elem.asJsonObject
            val tool = o.get("tool")?.takeIf { !it.isJsonNull }?.asString
                ?: o.get("toolSlug")?.takeIf { !it.isJsonNull }?.asString
                ?: return@mapNotNull null
            val params = mutableMapOf<String, String>()
            o.getAsJsonObject("params")?.entrySet()?.forEach { (k, v) ->
                params[k] = when {
                    v.isJsonNull -> ""
                    v.isJsonPrimitive -> v.asString
                    else -> v.toString()
                }
            }
            ToolChainStep(toolSlug = tool, params = params)
        }
        return steps.ifEmpty { null }
    }

    // ------------------------------------------------------------------
    // Structured text extraction:  STEP n: slug  key=value key2=value2
    // ------------------------------------------------------------------

    private fun parseStructuredSteps(text: String): ToolChain? {
        val regex = Regex(
            """(?im)^\s*(?:step\s*\d+\s*[:.]?\s*)?([a-z0-9_]+)\s*(?::\s*)?([a-z0-9_]+)\b(?:\s+(.*))?$"""
        )
        val matches = regex.findAll(text).toList()
        val steps = matches.mapNotNull { m ->
            val maybeTool = m.groupValues[2]
            // Heuristic: a real step has a verb prefix marker ("STEP 1:") or the
            // first capture is a known ordering word. We accept lines that begin
            // with "step" naturally above; here require the leading "step n".
            val leadingStep = Regex("(?im)^\\s*step\\s*\\d+", RegexOption.IGNORE_CASE)
                .containsMatchIn(m.value)
            if (!leadingStep) return@mapNotNull null
            if (maybeTool !in toolRegistry.listNames()) return@mapNotNull null
            val rest = m.groupValues[3]
            ToolChainStep(toolSlug = maybeTool, params = parseKeyValue(rest))
        }
        return steps.ifEmpty { null }?.let { ToolChain(steps = it, name = "AI Chain") }
    }

    private fun parseKeyValue(rest: String?): Map<String, String> {
        if (rest.isNullOrBlank()) return emptyMap()
        return rest.split(Regex("\\s+"))
            .mapNotNull { token ->
                val idx = token.indexOf('=')
                if (idx <= 0) null else token.substring(0, idx) to token.substring(idx + 1)
            }
            .toMap()
    }
}
