package com.correx.apps.server.inference

import com.correx.core.context.model.CompressionMetadata
import com.correx.core.context.model.ContextEntry
import com.correx.core.context.model.ContextLayer
import com.correx.core.context.model.ContextPack
import com.correx.core.context.model.EntryRole
import com.correx.core.events.events.CapabilityGapVerdict
import com.correx.core.events.types.ContextEntryId
import com.correx.core.events.types.ContextPackId
import com.correx.core.events.types.InferenceRequestId
import com.correx.core.events.types.SessionId
import com.correx.core.events.types.StageId
import com.correx.core.inference.GenerationConfig
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.ResponseFormat
import com.correx.core.kernel.orchestration.CapabilityGapInput
import com.correx.core.kernel.orchestration.CapabilityGapReflection
import com.correx.core.kernel.orchestration.CapabilityGapReflector
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json
import org.slf4j.LoggerFactory
import java.util.UUID

private const val REFLECTION_MAX_TOKENS = 4_096
private val REFLECTION_STAGE_ID = StageId("capability_gap_reflection")

/**
 * Vikunja #30 part 2: the "are you sure?" reflection judge. Given the plan-compile gate's
 * capability-gap findings for a session, asks a single bounded question — revise onto an existing
 * capability, or justify the need for a new one — and maps the answer to a [CapabilityGapReflection]
 * per gap. Mirrors [SemanticReviewerImpl]: talks to the provider directly with a self-contained
 * prompt, one inference call for the whole batch (no per-gap loop, no retry), and degrades to a
 * safe RESOLVED/no-escalation default (see [defaultReflection]) on any failure so a broken
 * reflection never wedges the plan-compile gate it sits behind.
 */
class CapabilityGapReflectorImpl(
    private val inferenceRouter: InferenceRouter,
    private val modelId: String? = null,
) : CapabilityGapReflector {

    override suspend fun reflect(sessionId: SessionId, gaps: List<CapabilityGapInput>): List<CapabilityGapReflection> {
        if (gaps.isEmpty()) return emptyList()
        return runCatching {
            val provider = inferenceRouter.route(REFLECTION_STAGE_ID, setOf(ModelCapability.General), modelId)
            val prompt = buildPrompt(gaps)
            val request = InferenceRequest(
                requestId = InferenceRequestId(UUID.randomUUID().toString()),
                sessionId = sessionId,
                stageId = REFLECTION_STAGE_ID,
                contextPack = pack(sessionId, prompt),
                generationConfig = GenerationConfig(temperature = 0.2, topP = 0.9, maxTokens = REFLECTION_MAX_TOKENS),
                responseFormat = ResponseFormat.Text,
            )
            parse(provider.infer(request).text, gaps)
        }.getOrElse {
            log.warn(
                "[CapabilityGapReflector] reflection failed session={}: {}",
                sessionId.value, it.message,
            )
            gaps.map(::defaultReflection)
        }
    }

    /** Malformed/missing verdict for a gap defaults to advisory RESOLVED-nothing — never NEEDS_TOOL
     * by default, so a broken reflection pass cannot force an escalation on its own. */
    private fun defaultReflection(gap: CapabilityGapInput) = CapabilityGapReflection(
        stageId = gap.stageId,
        impliedCapability = gap.impliedCapability,
        verdict = CapabilityGapVerdict.RESOLVED,
        detail = "reflection unavailable — treated as advisory, no escalation",
    )

    private fun buildPrompt(gaps: List<CapabilityGapInput>): String {
        val findings = gaps.joinToString("\n\n") { gap ->
            val granted = gap.grantedCapabilities.joinToString(", ").ifBlank { "(none)" }
            """
            - stage: ${gap.stageId}
              implied capability: ${gap.impliedCapability}
              triggering phrase: "${gap.phrase}"
              tools already granted to this stage give it: $granted
            """.trimIndent()
        }
        return """
        A deterministic scan of a compiled plan flagged stages whose brief implies a tool capability
        they were not granted. Before this is escalated to a human operator, decide for EACH gap
        whether it is a real gap or a false alarm.

        IMPORTANT domain fact: the file_write tool auto-creates parent directories — there is no
        separate mkdir/create-directory tool. A stage that says "create a directory" or "create the
        file" and already has file_write does NOT need a new capability.

        GAPS:
        $findings

        For EACH gap, respond with either:
        - RESOLVED: the stage's existing tools are actually sufficient; state the revised intent.
        - NEEDS_TOOL: the stage genuinely needs the missing capability; justify why no existing tool
          covers it.

        Respond with ONLY a JSON object, no prose, no markdown fence:
        {"reflections":[
          {"stage":"stage-id","verdict":"RESOLVED|NEEDS_TOOL","detail":"revised intent or justification"}
        ]}
        """.trimIndent()
    }

    private fun pack(sessionId: SessionId, prompt: String) = ContextPack(
        id = ContextPackId(UUID.randomUUID().toString()),
        sessionId = sessionId,
        stageId = REFLECTION_STAGE_ID,
        layers = mapOf(
            ContextLayer.L0 to listOf(
                ContextEntry(
                    id = ContextEntryId(UUID.randomUUID().toString()),
                    layer = ContextLayer.L0,
                    content = prompt,
                    sourceType = "capability_gap_reflection",
                    sourceId = "reflector",
                    tokenEstimate = prompt.length / 4,
                    role = EntryRole.USER,
                ),
            ),
        ),
        budgetUsed = prompt.length / 4,
        budgetLimit = prompt.length / 4,
        compressionMetadata = CompressionMetadata(),
    )

    private fun parse(raw: String, gaps: List<CapabilityGapInput>): List<CapabilityGapReflection> {
        val byStage = gaps.associateBy { it.stageId }
        val dto = extractJson(raw)
            ?.let { json ->
                runCatching { lenient.decodeFromString(ReflectionBatchDto.serializer(), json) }.getOrNull()
            }
            ?: return gaps.map(::defaultReflection)
        val byStageResult = dto.reflections.mapNotNull { r ->
            val gap = byStage[r.stage] ?: return@mapNotNull null
            val verdict = runCatching { CapabilityGapVerdict.valueOf(r.verdict.uppercase()) }
                .getOrDefault(CapabilityGapVerdict.RESOLVED)
            gap.stageId to CapabilityGapReflection(
                stageId = gap.stageId,
                impliedCapability = gap.impliedCapability,
                verdict = verdict,
                detail = r.detail.ifBlank { "no detail provided" },
            )
        }.toMap()
        return gaps.map { gap -> byStageResult[gap.stageId] ?: defaultReflection(gap) }
    }

    private fun extractJson(raw: String): String? {
        val start = raw.indexOf('{')
        val end = raw.lastIndexOf('}')
        return if (start in 0 until end) raw.substring(start, end + 1) else null
    }

    @Serializable
    private data class ReflectionBatchDto(val reflections: List<ReflectionDto> = emptyList())

    @Serializable
    private data class ReflectionDto(
        val stage: String = "",
        val verdict: String = "RESOLVED",
        val detail: String = "",
    )

    companion object {
        private val log = LoggerFactory.getLogger(CapabilityGapReflectorImpl::class.java)
        private val lenient = Json { ignoreUnknownKeys = true; isLenient = true }
    }
}
