package com.correx.apps.cli.commands

import com.correx.apps.cli.CorrexCli
import com.correx.apps.cli.DEFAULT_PORT
import com.github.ajalt.clikt.core.CliktCommand
import com.github.ajalt.clikt.parameters.arguments.argument
import com.github.ajalt.clikt.parameters.options.default
import com.github.ajalt.clikt.parameters.options.option
import io.ktor.client.HttpClient
import io.ktor.client.call.body
import io.ktor.client.engine.cio.CIO
import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
import io.ktor.client.request.get
import io.ktor.client.statement.bodyAsText
import io.ktor.serialization.kotlinx.json.json
import java.util.Locale
import kotlinx.coroutines.runBlocking
import kotlinx.serialization.Serializable
import kotlinx.serialization.json.Json

@Serializable
data class FailureMetricsDto(
    val inferenceFailures: Long = 0,
    val inferenceTimeouts: Long = 0,
    val toolFailures: Long = 0,
    val toolRejections: Long = 0,
    val stageFailures: Long = 0,
    val workflowFailures: Long = 0,
)

@Serializable
data class ProviderStatsDto(
    val provider: String,
    val completedCount: Long,
    val totalLatencyMs: Long,
    val promptTokens: Long,
    val completionTokens: Long,
    val tokensPerSecond: Double,
)

@Serializable
data class ToolStatsDto(
    val toolName: String,
    val completedCount: Long,
    val totalDurationMs: Long,
)

@Serializable
data class TierApprovalStatsDto(
    val tier: String,
    val requestedCount: Long,
    val resolvedCount: Long,
    val totalWaitMs: Long,
    val avgWaitMs: Long,
)

@Serializable
data class QualityStatsDto(
    val sourceFetches: Long = 0,
    val lowQualityExtractions: Long = 0,
    val extractionQualityRate: Double = 1.0,
    val retrievedHits: Long = 0,
    val droppedHits: Long = 0,
    val retrievalPrecision: Double = 1.0,
    val briefDriftCount: Long = 0,
    val capabilityGapCount: Long = 0,
)

@Serializable
data class StatsReportDto(
    val sessionId: String,
    val eventCount: Long = 0,
    val sessionDurationMs: Long = 0,
    val inferenceCount: Long = 0,
    val inferenceMs: Long = 0,
    val promptTokens: Long = 0,
    val completionTokens: Long = 0,
    val tokensPerSecond: Double = 0.0,
    val perProvider: List<ProviderStatsDto> = emptyList(),
    val toolCount: Long = 0,
    val toolMs: Long = 0,
    val perTool: List<ToolStatsDto> = emptyList(),
    val approvalsRequested: Long = 0,
    val approvalsResolved: Long = 0,
    val approvalsPending: Long = 0,
    val approvalWaitMs: Long = 0,
    val avgApprovalWaitMs: Long = 0,
    val perTier: List<TierApprovalStatsDto> = emptyList(),
    val failures: FailureMetricsDto = FailureMetricsDto(),
    val quality: QualityStatsDto = QualityStatsDto(),
    val inferencePct: Double = 0.0,
    val toolPct: Double = 0.0,
    val approvalWaitPct: Double = 0.0,
)

private const val MS_PER_SECOND = 1000L
private const val SECONDS_PER_MINUTE = 60L
private const val MINUTES_PER_HOUR = 60L
private const val PERCENT_MULTIPLIER = 100.0

private fun humanDuration(ms: Long): String {
    if (ms <= 0L) return "0s"
    val totalSeconds = ms / MS_PER_SECOND
    val hours = totalSeconds / (SECONDS_PER_MINUTE * MINUTES_PER_HOUR)
    val minutes = (totalSeconds / SECONDS_PER_MINUTE) % MINUTES_PER_HOUR
    val seconds = totalSeconds % SECONDS_PER_MINUTE
    return buildString {
        if (hours > 0) append("${hours}h ")
        if (hours > 0 || minutes > 0) append("${minutes}m ")
        append("${seconds}s")
    }.trim()
}

private fun otherPct(report: StatsReportDto): Double =
    (PERCENT_MULTIPLIER - report.inferencePct - report.toolPct - report.approvalWaitPct).coerceAtLeast(0.0)

fun renderStats(report: StatsReportDto): String {
    val lines = mutableListOf<String>()
    lines += "Session:    ${report.sessionId}"
    lines += "Events:     ${report.eventCount}"
    lines += "Duration:   ${humanDuration(report.sessionDurationMs)}"
    lines += ""

    lines += "Inference"
    lines += "  calls:      ${report.inferenceCount}"
    lines += "  tokens:     ${report.promptTokens} prompt + ${report.completionTokens} completion = " +
        "${report.promptTokens + report.completionTokens}"
    lines += "  time:       ${report.inferenceMs} ms  (%.1f tok/s)".format(Locale.ROOT, report.tokensPerSecond)
    for (p in report.perProvider) {
        lines += "    %-20s %4d calls  %6d tok  %8d ms  (%.1f tok/s)".format(
            Locale.ROOT,
            p.provider, p.completedCount, p.promptTokens + p.completionTokens, p.totalLatencyMs, p.tokensPerSecond,
        )
    }
    lines += ""

    lines += "Tools"
    lines += "  calls:      ${report.toolCount}   time: ${report.toolMs} ms"
    for (t in report.perTool) {
        lines += "    %-24s %4d  %8d ms".format(Locale.ROOT, t.toolName, t.completedCount, t.totalDurationMs)
    }
    lines += ""

    lines += "Approvals"
    lines += "  requested:  ${report.approvalsRequested}   resolved: ${report.approvalsResolved}   " +
        "pending: ${report.approvalsPending}"
    lines += "  total wait: ${report.approvalWaitMs} ms  (avg ${report.avgApprovalWaitMs} ms)"
    for (tier in report.perTier) {
        lines += "    %-4s req %d  res %d  wait %d ms (avg %d ms)".format(
            Locale.ROOT,
            tier.tier, tier.requestedCount, tier.resolvedCount, tier.totalWaitMs, tier.avgWaitMs,
        )
    }
    lines += ""

    val f = report.failures
    lines += "Failures"
    lines += "  inference: ${f.inferenceFailures}  timeouts: ${f.inferenceTimeouts}  " +
        "tool: ${f.toolFailures}  rejected: ${f.toolRejections}  " +
        "stage: ${f.stageFailures}  workflow: ${f.workflowFailures}"
    lines += ""

    val q = report.quality
    lines += "Signal quality"
    lines += "  extraction:  %d fetched, %d low-quality  (%.0f%% clean)".format(
        Locale.ROOT, q.sourceFetches, q.lowQualityExtractions, q.extractionQualityRate * PERCENT_MULTIPLIER,
    )
    lines += "  retrieval:   %d kept, %d filtered  (%.0f%% precision)".format(
        Locale.ROOT, q.retrievedHits, q.droppedHits, q.retrievalPrecision * PERCENT_MULTIPLIER,
    )
    lines += "  brief drift: ${q.briefDriftCount}   capability gaps: ${q.capabilityGapCount}"
    lines += ""

    lines += "Time accounting (% of session wall time)"
    lines += "  inference: %.1f%%   tools: %.1f%%   approval-wait: %.1f%%   other: %.1f%%".format(
        Locale.ROOT,
        report.inferencePct, report.toolPct, report.approvalWaitPct, otherPct(report),
    )

    return lines.joinToString("\n")
}

class StatsCommand : CliktCommand(name = "stats") {
    private val sessionId by argument("sessionId")
    private val host by option("--host").default("localhost")
    private val port by option("--port").default("$DEFAULT_PORT")

    private val statsJson = Json { ignoreUnknownKeys = true }

    override fun run(): Unit = runBlocking {
        val portInt = port.toIntOrNull() ?: DEFAULT_PORT
        val outputJson = (currentContext.findRoot().command as? CorrexCli)?.json ?: false
        val url = "http://$host:$portInt/sessions/$sessionId/metrics"

        val client = HttpClient(CIO) {
            install(ContentNegotiation) { json(statsJson) }
        }

        runCatching {
            if (outputJson) {
                println(client.get(url).bodyAsText())
            } else {
                println(renderStats(client.get(url).body()))
            }
        }.onFailure { e ->
            System.err.println("Error fetching stats: ${e.message}")
        }

        client.close()
    }
}
