package com.correx.core.kernel.orchestration

import com.correx.core.approvals.DefaultApprovalRepository
import com.correx.core.approvals.domain.ApprovalEngine
import com.correx.core.approvals.model.ApprovalDecision
import com.correx.core.artifacts.kind.ArtifactKindRegistry
import com.correx.core.artifacts.kind.KindInference
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.context.builder.ContextPackBuilder
import com.correx.core.context.model.ContextPack
import com.correx.core.events.events.ClarificationAnswer
import com.correx.core.events.events.ClarificationAnsweredEvent
import com.correx.core.events.events.ClarificationRequestedEvent
import com.correx.core.events.events.InferenceCompletedEvent
import com.correx.core.events.events.InferenceFailedEvent
import com.correx.core.events.events.InferenceStartedEvent
import com.correx.core.events.events.InferenceTimeoutEvent
import com.correx.core.events.events.ToolRequest
import com.correx.core.toolintent.ToolCallAssessor
import com.correx.core.toolintent.WorkspacePolicy
import com.correx.core.toolintent.WorldProbe
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ApprovalRequestId
import com.correx.core.events.types.ArtifactId
import com.correx.core.events.types.ClarificationRequestId
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.InferenceRepository
import com.correx.core.inference.InferenceRequest
import com.correx.core.inference.InferenceResponse
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.PromptRenderer
import com.correx.core.inference.ResponseFormat
import com.correx.core.inference.Tokenizer
import com.correx.core.inference.ToolDefinition
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.ToolFunction
import com.correx.core.kernel.execution.WorkflowResult
import com.correx.core.risk.RiskAssessor
import com.correx.core.tools.contract.ToolCapability
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.registry.ToolRegistry
import com.correx.core.transitions.evaluation.PromptResolver
import com.correx.core.transitions.execution.StageExecutionResult
import com.correx.core.transitions.graph.StageConfig
import com.correx.core.transitions.graph.WorkflowGraph
import com.correx.core.transitions.resolution.TransitionResolver
import com.correx.core.validation.artifact.ArtifactExtractionPipeline
import com.correx.core.validation.model.ValidationContext
import com.correx.core.validation.pipeline.ValidationOutcome
import com.correx.core.validation.pipeline.ValidationPipeline
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.withTimeout
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonArray
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.JsonPrimitive
import com.correx.core.journal.DecisionJournalRenderer
import com.correx.core.journal.DefaultDecisionJournalRepository
import com.correx.core.kernel.orchestration.subagent.SubagentRunner
import org.slf4j.LoggerFactory
import java.util.*
import java.util.concurrent.*
import java.util.concurrent.atomic.*
import kotlin.coroutines.cancellation.CancellationException

// Ceiling on agentic tool rounds per stage entry. Originally 10 as a tight leash on early local
// models that would thrash tools; the model now drives tools reliably, and the tight cap started
// starving legitimate multi-step work — notably the recovery stage, which spends most rounds
// re-investigating (list_dir/read/rebuild) before it reaches the write, and got bounced out before
// acting on a correct diagnosis. The real runaway guards are the read-loop and rejection-loop
// nudges (both fire after 3 unproductive rounds), so this ceiling only needs to be generous enough
// that productive stages aren't cut off mid-work. Value now lives in OrchestrationTuning.maxToolRounds.

// Consecutive read-only tool rounds (no file_write/file_edit) that still owe a file_written
// artifact before we force the write nudge. A model that keeps calling read tools every round
// trips neither the prose nudge nor the stage_complete nudge, so without this it silently burns
// every round reading and never writes (2026-07-05: define_types read-looped 40 turns → failed).
// Value now lives in OrchestrationTuning.readLoopNudgeThreshold.

// Consecutive rounds in which EVERY tool call was rejected/denied (plane-2 BLOCKED or a stage/policy
// ERROR) with nothing succeeding. Distinct from the read-loop breaker, which only fires for stages
// that owe a file_written artifact — a JSON-emitting stage (discovery, analyst) that fixates on a
// rejected path (e.g. list/read of a not-yet-created dir) would otherwise thrash to MAX_TOOL_ROUNDS
// with no progress. After this many all-rejected rounds we force the model to stop retrying and
// produce its output; a single successful tool call resets the counter.
// Value now lives in OrchestrationTuning.rejectionLoopNudgeThreshold.
internal val WRITE_TOOL_NAMES = setOf("file_write", "file_edit")
internal const val STAGE_COMPLETE_TOOL = "stage_complete"
internal const val EMIT_ARTIFACT_TOOL = "emit_artifact"
internal const val SCOPE_PROPOSAL_TOOL = "propose_scope"
internal const val READ_BEFORE_WRITE_CODE = "READ_BEFORE_WRITE"
internal const val REFERENCE_EXISTS_CODE = "REFERENCE_EXISTS"
internal const val OUTPUT_SUMMARY_LIMIT = 500

// Global floor on tool-result text entering model context, applied beneath each tool's own
// outputCompressor. A tool without a compressor (or one whose output survives compression) can still
// flood a small model's window, so any Success body over this many chars is shown head+tail with a
// marker and the FULL raw output spilled to the artifact store for retrieval via tool_output.
// ponytail: fixed constants; lift to OrchestrationTuning if operators need to tune per-deployment.
internal const val TOOL_RESULT_MAX_CHARS = 8_000
internal const val TOOL_RESULT_HEAD_LINES = 60
internal const val TOOL_RESULT_TAIL_LINES = 60
internal const val TOOL_OUTPUT_TOOL = "tool_output"

/**
 * Frame an over-cap tool output as `header` + head lines + a truncation marker (naming the
 * [tool_output] ref that retrieves the full text) + tail lines. Head and tail are each char-capped
 * so a single pathological long line can't defeat the line-count bound. Pure — the caller spills the
 * full output and supplies its [ref].
 */
internal fun frameTruncatedToolResult(header: String, compressed: String, ref: String): String {
    val lines = compressed.split("\n")
    val head = lines.take(TOOL_RESULT_HEAD_LINES).joinToString("\n").take(TOOL_RESULT_MAX_CHARS / 2)
    val tail = lines.takeLast(TOOL_RESULT_TAIL_LINES).joinToString("\n").takeLast(TOOL_RESULT_MAX_CHARS / 2)
    val marker = "… output truncated (${lines.size} lines); " +
        "call $TOOL_OUTPUT_TOOL(ref=\"$ref\") for the full output …"
    return "$header\n$head\n$marker\n$tail"
}
// Read-on-demand doc catalog: the top-N docs (by repo-map recency score) surfaced as
// `path — descriptor` so a stage learns which docs exist and can file_read one when relevant,
// instead of docs being force-fed (or excluded outright). One line each, hard-capped, so it can
// stay always-on without re-poisoning context (2026-07-07 doc-injection rework).
// Value now lives in OrchestrationTuning.docsCatalogMax.

// A stage prompt must explicitly ask about documentation for .md/docs paths to enter the repo
// layout listing; otherwise docs only reach context through a real retrieval hit.
internal val DOCS_REQUEST_REGEX =
    Regex("""\b(docs?|readme|adr|changelog|documentation|markdown)\b|\.md\b""", RegexOption.IGNORE_CASE)

internal fun isDocPath(path: String): Boolean =
    path.endsWith(".md", ignoreCase = true) ||
        path.split('/').any { it.equals("docs", ignoreCase = true) || it.equals("doc", ignoreCase = true) }

// Source types whose entries land in the REQUIRED context bucket (see stampBuckets).
internal val REQUIRED_SOURCE_TYPES = setOf(
    "systemPrompt",
    "agentPrompt",
    "agentInstructions",
    "claimedTask",
    "schemaInstruction",
    "steeringNote",
    "clarificationAnswer",
    "rejectionFeedback",
    "retryFeedback",
    "neededArtifact",
    "criticFeedback",
    // #290: original intent stays unprunable via the REQUIRED bucket now that it renders as
    // L1/USER instead of relying on the old L0/SYSTEM never-drop placement.
    "initialIntent",
)

// HTTP statuses that are transient despite being 4xx (F-002 retry classification).
internal const val HTTP_TIMEOUT = 408
internal const val HTTP_TOO_MANY_REQUESTS = 429

// Cap on clarification rounds per stage — OrchestrationTuning.maxClarificationRounds.

// Static-analysis output caps: the tail retained in the recorded event vs. the (larger) slice fed
// back to the model so it can fix the failure. Tails, because the error summary sits at the end.
internal const val CONTRACT_EVIDENCE_CAP = 400
// Gate 3 semantic review: a FAIL blocks only on a correctness finding at/above this confidence.
// Per-gate retry budgeting + hybrid salvage (design 2026-07-06-per-gate-retry-budgets.md) is now the
// authority on when review retries stop — the gate simply reports the block and the RetryCoordinator
// charges/exhausts the "review" budget, then the salvage judge decides continue-vs-fail. This cap
// remains only as an absolute backstop against a pathological run whose review findings *change*
// every attempt (so the budget never charges, i.e. the run looks like perpetual "progress"): after
// this many blocks, findings surface without blocking so a stuck reviewer can never trap a stage
// forever. Set well above any legitimate budget+salvage sequence. Objective text is capped too.
// REVIEW_BLOCK_MIN_CONFIDENCE / REVIEW_BLOCK_RETRY_CAP now in OrchestrationTuning.
internal const val REVIEW_OBJECTIVE_CAP = 4_000
internal const val STATIC_ANALYSIS_SUMMARY_CAP = 2_000
internal const val STATIC_ANALYSIS_FEEDBACK_CAP = 6_000
/** Profile command alias run before a build gate to make it runnable (e.g. `npm ci`). */
internal const val SETUP_COMMAND_ALIAS = "setup"

// KindInference kinds for build manifests — a written file of one of these declares a buildable
// project toolchain even before any source lands, so the auto build gate fires (see
// SessionOrchestrator.sessionProducedBuildTarget). Kept in sync with KindInference's manifest kinds.
internal val BUILD_MANIFEST_KINDS = setOf("package_json", "gradle_module")

@SuppressWarnings(
    "ForbiddenComment",
    "UnusedParameter",
    "ReturnCount",
    "TooGenericExceptionCaught",
    "TooManyFunctions",
    "NestedBlockDepth",
    "LongMethod",
)
abstract class SessionOrchestrator(
    repositories: OrchestratorRepositories,
    engines: OrchestratorEngines,
    internal val artifactStore: ArtifactStore,
    internal val decisionJournalRepository: DefaultDecisionJournalRepository,
    internal val decisionJournalRenderer: DecisionJournalRenderer = DecisionJournalRenderer(),
    internal val artifactKindRegistry: ArtifactKindRegistry? = null,
    internal val repoKnowledgeRetriever: RepoKnowledgeRetriever? = null,
    internal val readyTaskCounter: ReadyTaskCounter? = null,
    internal val taskClaimCoordinator: TaskClaimCoordinator? = null,
    internal val tuning: OrchestrationTuning = OrchestrationTuning(),
) {
    internal val log = LoggerFactory.getLogger(this::class.java)
    internal val eventStore: EventStore = repositories.eventStore
    internal val transitionResolver: TransitionResolver = engines.transitionResolver
    internal val contextPackBuilder: ContextPackBuilder = engines.contextPackBuilder
    internal val inferenceRouter: InferenceRouter = engines.inferenceRouter
    internal val validationPipeline: ValidationPipeline = engines.validationPipeline
    internal val riskAssessor: RiskAssessor = engines.riskAssessor
    internal val promptResolver: PromptResolver = engines.promptResolver
    internal val toolExecutor: ToolExecutor? = engines.toolExecutor
    internal val toolRegistry: ToolRegistry? = engines.toolRegistry
    internal val toolCallAssessor: ToolCallAssessor? = engines.toolCallAssessor
    internal val workspacePolicy: WorkspacePolicy? = engines.workspacePolicy
    internal val worldProbe: WorldProbe = engines.worldProbe
    internal val staticAnalysisRunner: StaticAnalysisRunner? = engines.staticAnalysisRunner
    internal val lspDiagnosticsRunner: LspDiagnosticsRunner? = engines.lspDiagnosticsRunner
    internal val contractAssertionEvaluator: ContractAssertionEvaluator? = engines.contractAssertionEvaluator
    internal val planCompilationCheck: PlanCompilationCheck? = engines.planCompilationCheck
    internal val semanticReviewer: SemanticReviewer? = engines.semanticReviewer
    internal val workspaceToolRegistryProvider: WorkspaceToolRegistryProvider? = engines.workspaceToolRegistryProvider
    internal val inferenceRepository: InferenceRepository = repositories.inferenceRepository
    internal val orchestrationRepository: OrchestrationRepository = repositories.orchestrationRepository
    internal open val tokenizer: Tokenizer? = null
    internal val approvalEngine: ApprovalEngine = engines.approvalEngine
    internal val approvalRepository: DefaultApprovalRepository = repositories.approvalRepository
    internal abstract val cancellations: ConcurrentHashMap<SessionId, AtomicBoolean>

    internal abstract val subagentRunner: SubagentRunner
    internal val pendingApprovals: ConcurrentHashMap<ApprovalRequestId, CompletableDeferred<ApprovalDecision>> =
        ConcurrentHashMap()

    // A stage whose produced artifact carries open questions parks its run on one of these until the
    // operator answers (submitClarification completes it). Separate from pendingApprovals so a
    // clarification round never consumes the failure-retry budget.
    internal val pendingClarifications:
        ConcurrentHashMap<ClarificationRequestId, CompletableDeferred<List<ClarificationAnswer>>> =
        ConcurrentHashMap()

    /**
     * Request-ids of clarifications this process is still parked on (the [CompletableDeferred] is
     * live and awaitable). A reconnecting client uses this to re-push only clarifications that can
     * actually be answered — a request whose deferred died with a server restart is excluded, so the
     * client never shows a modal that would resume nothing.
     */
    internal val artifactContentCache: ConcurrentHashMap<String, String> = ConcurrentHashMap()

    // ACR Store 1: in-process memo of describe().render(), keyed on (repoRoot, path, contentHash).
    // Disposable — the source of truth is FileWrittenEvent + CAS; this only skips recomputing a pure
    // function of content-addressed bytes. Empty string = a computed "no descriptor" (negative cache).
    internal val descriptorMemo: ConcurrentHashMap<String, String> = ConcurrentHashMap()

    /** Drops a terminated session's cached artifact contents (the heaviest per-session state — full
     *  file/JSON payloads). Safe: rehydrateArtifactContentCache rebuilds it from durable events if the
     *  session is ever resumed. Called on WorkflowCompleted/WorkflowFailed. */
    internal val artifactExtractionPipeline = ArtifactExtractionPipeline()

    abstract suspend fun run(
        sessionId: SessionId,
        graph: WorkflowGraph,
        config: OrchestrationConfig,
    ): WorkflowResult

    abstract suspend fun cancel(sessionId: SessionId)

    /** Live clarification request ids across sessions (member: consumed cross-module by the WS layer). */
    fun liveClarificationRequestIds(): Set<String> =
        pendingClarifications.keys.mapTo(mutableSetOf()) { it.value }

    /**
     * The clarification a headless caller should answer for this session: the newest still-live,
     * unanswered request. Null if the session is not parked on a clarification. Lets a REST answer
     * route resolve the (stageId, requestId) server-side so a script need only supply answer values —
     * the WS path knows them because it received the [ClarificationRequestedEvent]; a curl caller does not.
     */
    suspend fun pendingClarificationFor(sessionId: SessionId): ClarificationRequestedEvent? {
        val live = liveClarificationRequestIds()
        if (live.isEmpty()) return null
        val events = eventStore.read(sessionId)
        val answered = events.mapNotNull { it.payload as? ClarificationAnsweredEvent }
            .mapTo(mutableSetOf()) { it.requestId }
        return events.mapNotNull { it.payload as? ClarificationRequestedEvent }
            .lastOrNull { it.requestId.value in live && it.requestId !in answered }
    }

    /** Public seam for out-of-band gates (e.g. freestyle plan review) that reuse the stage-approval flow. */
    suspend fun requestPlanApproval(sessionId: SessionId, preview: String): Boolean =
        requestStageApproval(sessionId, StageId("plan_review"), preview)

    internal open suspend fun mapValidationOutcome(
        sessionId: SessionId,
        stageId: StageId,
        context: ValidationContext,
    ): StageExecutionResult = when (val outcome = validationPipeline.validate(context)) {
        is ValidationOutcome.Passed -> StageExecutionResult.Success(emptyList())
        is ValidationOutcome.Rejected ->
            StageExecutionResult.Failure(summarizeRejection(outcome.report), retryable = outcome.retryable)

        is ValidationOutcome.NeedsApproval -> handleApproval(sessionId, stageId, outcome)
    }

    @Suppress("LongParameterList")
    internal open suspend fun runInference(
        sessionId: SessionId,
        stageId: StageId,
        contextPack: ContextPack,
        stageConfig: StageConfig,
        timeoutMs: Long,
        responseFormat: ResponseFormat = ResponseFormat.Text,
        effectives: RunEffectives = RunEffectives(toolRegistry, toolExecutor, workspacePolicy),
        // When false, send NO tools. Some local chat templates (e.g. llama.cpp's Gemma) switch to a
        // tool-calling/reasoning-channel format when tools are present and leak `<|channel>` markers
        // into the artifact text; a tools-less call yields clean JSON deterministically (and lets the
        // JSON grammar constraint re-apply, since grammar+tools is rejected by llama.cpp).
        withTools: Boolean = true,
        // When true, offer ONLY write tools (+ stage_complete / emit_artifact) — read tools are
        // stripped. Used to break a read-loop: a model that owes a file_written artifact but keeps
        // calling read tools ignores a textual nudge (its output all goes to reasoning_content), so
        // we structurally remove the option to read and leave writing as the only move.
        forceWriteOnly: Boolean = false,
    ): InferenceResult {
        // A stage that grants tools needs a tool-calling model, so request ToolCalling on top of any
        // declared capabilities — the capability-aware strategy then ranks eligible providers by their
        // ToolCalling score and routes the stage to the best tool-caller.
        val requiredCapabilities = stageConfig.requiredCapabilities +
            if (withTools && stageConfig.allowedTools.isNotEmpty()) setOf(ModelCapability.ToolCalling) else emptySet()
        // Routing itself can fail transiently (provider mid-crash-recovery) — it must be retryable
        // like any other inference failure, not escape and kill the whole session (see #299).
        val provider = try {
            inferenceRouter.route(stageId, requiredCapabilities, stageConfig.modelId)
        } catch (e: CancellationException) {
            throw e
        } catch (e: Exception) {
            return InferenceResult.Failed(e.message ?: "routing failed")
        }
        log.debug(
            "[Orchestrator] inference session={} stage={} provider={} timeoutMs={}",
            sessionId.value, stageId.value, provider.id.value, timeoutMs,
        )
        val requestId = InferenceRequestId(UUID.randomUUID().toString())
        val request = InferenceRequest(
            requestId = requestId,
            sessionId = sessionId,
            stageId = stageId,
            contextPack = contextPack,
            generationConfig = stageConfig.generationConfig,
            responseFormat = responseFormat,
            tools = if (!withTools) {
                emptyList()
            } else {
                // Read the log ONCE for the read-only check instead of once per tool inside the filter
                // (the flag is tool-independent) — this filter runs per tool per inference round.
                val readOnlyMode = isReadOnlyMode(sessionId)
                // tool_output is withheld until the session has actually spilled an over-cap output to
                // CAS — only then can it retrieve anything, and only then is its hash-ref marker in play.
                val toolNames = stageConfig.effectiveAllowedTools.let { declared ->
                    if (declared.isNotEmpty() && sessionHasSpilledOutput(sessionId)) {
                        declared + TOOL_OUTPUT_TOOL
                    } else {
                        declared
                    }
                }
                toolNames
                    .mapNotNull { effectives.registry?.resolve(it) }
                    .filter { tool ->
                        // ponytail: filter write tools while read-before-write block is active;
                        // restored once a read completes
                        !readOnlyMode || ToolCapability.FILE_WRITE !in tool.requiredCapabilities
                    }
                    .filter { tool ->
                        // Read-loop break: keep only write tools (drop file_read/list_dir/shell) so
                        // the model cannot keep reading and must write or complete.
                        !forceWriteOnly || ToolCapability.FILE_WRITE in tool.requiredCapabilities
                    }
                    .map { tool ->
                        ToolDefinition(
                            function = ToolFunction(
                                name = tool.name,
                                description = tool.description,
                                parameters = tool.parametersSchema,
                            ),
                        )
                    } + ToolDefinition(
                        function = ToolFunction(
                            name = STAGE_COMPLETE_TOOL,
                            description = "Call this tool when the stage's goal is fully met and no " +
                                    "further tool calls are needed. The orchestrator will proceed to the next stage.",
                            parameters = JsonObject(emptyMap()),
                        ),
                    ) + emitArtifactTool(stageConfig)
            },
        )

        val renderedMessages = PromptRenderer.render(contextPack)
        val promptArtifactId = artifactStore.put(Json.encodeToString(renderedMessages).toByteArray())
        emit(sessionId, InferenceStartedEvent(requestId, sessionId, stageId, provider.id, promptArtifactId))

        if (isCancelled(sessionId)) return InferenceResult.Cancelled

        return try {
            val rawResponse = withTimeout(timeoutMs) { provider.infer(request) }
            // Models frequently wrap a JSON artifact in a ```json … ``` markdown fence despite
            // being told not to. The fence makes the artifact unparseable → validation fails.
            // Strip an outer fence wrapper so the stored artifact, its CAS hash, the content
            // cache, and validation all see the same clean payload.
            val response = rawResponse.copy(text = stripCodeFence(rawResponse.text))
            log.debug(
                "[Orchestrator] inference done session={} stage={} tokens={} latencyMs={}",
                sessionId.value, stageId.value, response.tokensUsed, response.latencyMs,
            )

            if (isCancelled(sessionId)) return InferenceResult.Cancelled

            val responseArtifactId = artifactStore.put(response.text.toByteArray())
            val reasoningArtifactId = response.reasoning?.let { artifactStore.put(it.toByteArray()) }
            emit(
                sessionId,
                InferenceCompletedEvent(
                    requestId = requestId,
                    sessionId = sessionId,
                    stageId = stageId,
                    providerId = provider.id,
                    tokensUsed = response.tokensUsed,
                    latencyMs = response.latencyMs,
                    responseArtifactId = responseArtifactId,
                    reasoningArtifactId = reasoningArtifactId,
                ),
            )

            InferenceResult.Success(response, responseArtifactId)
        } catch (e: TimeoutCancellationException) {
            emit(
                sessionId,
                InferenceTimeoutEvent(
                    requestId = requestId,
                    sessionId = sessionId,
                    stageId = stageId,
                    providerId = provider.id,
                    timeoutMs = timeoutMs,
                ),
            )
            InferenceResult.Failed("TIMEOUT: ${e.message}")
        } catch (e: CancellationException) {
            throw e // never swallow
        } catch (e: Exception) {
            if (isConnectionLevelFailure(e)) {
                // Mark it down NOW instead of waiting for the next periodic health poll (~18s lag,
                // see #300) — the retry's route() call must see this provider as unavailable
                // immediately so it gates/waits (#299) rather than instantly re-selecting the dead
                // provider again.
                inferenceRouter.reportFailure(provider.id, e.message ?: "connection failure")
            }
            emit(
                sessionId,
                InferenceFailedEvent(
                    requestId = requestId,
                    sessionId = sessionId,
                    stageId = stageId,
                    providerId = provider.id,
                    reason = e.message ?: "unknown",
                ),
            )
            InferenceResult.Failed(e.message ?: "inference failed")
        }
    }

    // Connection-level failures (provider crashed/restarting mid-request) should gate routing
    // immediately; other failures (bad request, model error, HTTP 4xx) should not mark the
    // provider down since the provider itself is still reachable.
    private fun isConnectionLevelFailure(e: Exception): Boolean {
        val message = e.message.orEmpty()
        return e is java.net.ConnectException ||
            e is java.net.SocketException ||
            e is java.io.IOException || // covers ktor/CIO's IOException, which extends java.io.IOException on the JVM
            message.contains("prematurely closed", ignoreCase = true) ||
            message.contains("connection refused", ignoreCase = true) ||
            message.contains("connection reset", ignoreCase = true)
    }

    // --- token estimation ---

    internal open suspend fun estimateTokens(content: String): Int {
        val t = tokenizer
        if (t != null) {
            // A zero count for non-blank content means the tokenizer endpoint is lying
            // (e.g. a malformed request silently answered with an empty token list) —
            // trust the heuristic instead, or every budget check goes blind.
            return runCatching { t.countTokens(content) }
                .getOrElse { fallbackTokenEstimate(content) }
                .takeIf { it > 0 || content.isBlank() }
                ?: fallbackTokenEstimate(content)
        }
        return fallbackTokenEstimate(content)
    }

}


internal sealed interface InferenceResult {
    data class Success(val response: InferenceResponse, val responseArtifactId: ArtifactId? = null) : InferenceResult
    data class Failed(val reason: String) : InferenceResult
    data object Cancelled : InferenceResult
}

/**
 * Flattens an LLM tool call's JSON arguments into a [ToolRequest.parameters] map. A JSON array of
 * primitives becomes a `List<String>` so list-param accessors read it — otherwise it arrived as a
 * `.toString()` string and silently read empty, dropping e.g. a task's `affected_paths` (which in
 * turn left the write-scope and cite-before-claim gates with nothing to enforce). A primitive
 * becomes its content string; objects and arrays-of-objects keep their JSON text for tools that
 * re-parse them (e.g. task_decompose). Malformed JSON yields an empty map.
 */
// Lenient fallback for small local models (gemma) that drift out of strict JSON on large tool
// args — most commonly bare/unquoted keys (`path:` instead of `"path":`) and unquoted string
// values. Strict parse is tried first; only genuinely-off syntax hits this. Triple-quoted
// Python-style values (`"""..."""`) still fail here — that ambiguity isn't worth a fragile
// hand-rolled repair; the loud "unparseable" path below tells the model to fix its format.
internal val lenientJson = Json { isLenient = true }

internal fun parseToolArguments(arguments: String): Map<String, Any> = runCatching {
    val element = runCatching { Json.parseToJsonElement(arguments) }
        .getOrElse { lenientJson.parseToJsonElement(arguments) }
    val obj = element as? JsonObject ?: return@runCatching emptyMap()
    obj.entries.associate { (k, v) ->
        k to when {
            v is JsonPrimitive -> v.content
            v is JsonArray && v.all { it is JsonPrimitive } -> v.map { (it as JsonPrimitive).content }
            else -> v.toString()
        } as Any
    }
}.getOrElse { emptyMap() }
