package com.correx.apps.cli.commands

import com.correx.apps.cli.CorrexCli
import com.correx.apps.cli.DEFAULT_PORT
import com.correx.apps.cli.ws.CliWsClient
import com.github.ajalt.clikt.core.CliktCommand
import com.github.ajalt.clikt.parameters.options.default
import com.github.ajalt.clikt.parameters.options.flag
import com.github.ajalt.clikt.parameters.options.option
import com.github.ajalt.clikt.parameters.options.required
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.plugins.websocket.WebSockets
import io.ktor.client.request.post
import io.ktor.client.request.setBody
import io.ktor.http.ContentType
import io.ktor.http.contentType
import io.ktor.serialization.kotlinx.json.json
import io.ktor.websocket.WebSocketSession
import io.ktor.websocket.send
import kotlinx.coroutines.runBlocking
import kotlinx.serialization.Serializable
import kotlinx.serialization.encodeToString
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.JsonObject
import kotlinx.serialization.json.jsonObject
import kotlinx.serialization.json.jsonPrimitive

@Serializable
private data class StartSessionRequest(val workflowId: String, val sessionId: String?, val intent: String? = null)

@Serializable
private data class StartSessionResponse(val sessionId: String)

@Serializable
private data class ApprovalResponsePayload(
    val type: String,
    val requestId: String,
    val decision: String,
    val steeringNote: String?,
)

private data class RunContext(
    val sid: String,
    val outputJson: Boolean,
    val autoApprove: Boolean,
    val isTty: Boolean,
)

private data class ApprovalContext(
    val session: WebSocketSession,
    val requestId: String,
    val toolName: String,
    val preview: String?,
)

class RunCommand : CliktCommand(name = "run") {
    private val workflow by option("--workflow", help = "Path to workflow definition").required()
    private val sessionId by option("--session", help = "Existing session ID to resume")
    private val autoApprove by option("--auto-approve", help = "Auto-approve all approval requests").flag()
    private val intent by option("--intent", help = "Freestyle intent to seed a new session")
    private val host by option("--host", help = "Server host").default("localhost")
    private val port by option("--port", help = "Server port").default("$DEFAULT_PORT")

    private val serializer = Json { ignoreUnknownKeys = true; classDiscriminator = "type" }

    override fun run(): Unit = runBlocking {
        val portInt = port.toIntOrNull() ?: DEFAULT_PORT
        val outputJson = (currentContext.findRoot().command as? CorrexCli)?.json ?: false

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

        val sid = startSession(client, host, portInt, outputJson) ?: run {
            client.close()
            return@runBlocking
        }

        printLine("session_started", mapOf("sessionId" to sid, "workflow" to workflow), outputJson)

        val ctx = RunContext(sid, outputJson, autoApprove, System.console() != null)
        var exitCode = 0
        val wsClient = CliWsClient(client)

        wsClient.sessionStream(host, portInt, sid) { session, raw ->
            val msgJson = runCatching { serializer.parseToJsonElement(raw).jsonObject }.getOrNull()
            val keepGoing = handleMessage(session, msgJson, ctx) { exitCode = it }
            keepGoing
        }

        client.close()
        if (exitCode != 0) {
            throw SystemExitException(exitCode)
        }
    }

    private fun resolveWorkflowId(path: String): String =
        java.nio.file.Path.of(path).fileName.toString().removeSuffix(".toml")

    private suspend fun startSession(
        client: HttpClient,
        host: String,
        portInt: Int,
        outputJson: Boolean,
    ): String? = runCatching {
        val resp = client.post("http://$host:$portInt/sessions") {
            contentType(ContentType.Application.Json)
            setBody(StartSessionRequest(resolveWorkflowId(workflow), sessionId, intent))
        }
        resp.body<StartSessionResponse>().sessionId
    }.getOrElse { e ->
        printLine("error", mapOf("message" to "Failed to start session: ${e.message}"), outputJson)
        null
    }

    @Suppress("CyclomaticComplexMethod", "LongMethod")
    private suspend fun handleMessage(
        session: WebSocketSession,
        msgJson: JsonObject?,
        ctx: RunContext,
        setExitCode: (Int) -> Unit,
    ): Boolean {
        val type = msgJson?.get("type")?.jsonPrimitive?.content ?: "unknown"
        fun field(key: String): String? = msgJson?.get(key)?.jsonPrimitive?.content

        return when (type) {
            "SessionCompleted" -> {
                printLine("session_completed", mapOf("sessionId" to ctx.sid), ctx.outputJson)
                setExitCode(0)
                false
            }

            "SessionFailed" -> {
                val reason = field("reason") ?: "unknown"
                printLine("session_failed", mapOf("sessionId" to ctx.sid, "reason" to reason), ctx.outputJson)
                setExitCode(1)
                false
            }

            "StageStarted" -> {
                printLine("stage_started", mapOf("stageId" to (field("stageId") ?: "")), ctx.outputJson)
                true
            }

            "StageCompleted" -> {
                printLine("stage_completed", mapOf("stageId" to (field("stageId") ?: "")), ctx.outputJson)
                true
            }

            "StageFailed" -> {
                printLine(
                    "stage_failed",
                    mapOf("stageId" to (field("stageId") ?: ""), "reason" to (field("reason") ?: "")),
                    ctx.outputJson,
                )
                true
            }

            "ToolStarted" -> {
                printLine("tool_started", mapOf("tool" to (field("toolName") ?: "")), ctx.outputJson)
                true
            }

            "ToolCompleted" -> {
                printLine("tool_completed", mapOf("tool" to (field("toolName") ?: "")), ctx.outputJson)
                true
            }

            "ToolFailed" -> {
                printLine(
                    "tool_failed",
                    mapOf("tool" to (field("toolName") ?: ""), "reason" to (field("reason") ?: "")),
                    ctx.outputJson,
                )
                true
            }

            "ToolRejected" -> {
                printLine(
                    "tool_rejected",
                    mapOf("tool" to (field("toolName") ?: ""), "reason" to (field("reason") ?: "")),
                    ctx.outputJson,
                )
                true
            }

            "ApprovalRequired" -> {
                val approvalCtx = ApprovalContext(
                    session = session,
                    requestId = field("requestId") ?: "",
                    toolName = field("toolName") ?: "",
                    preview = field("preview"),
                )
                handleApproval(approvalCtx, ctx, setExitCode)
            }

            else -> true
        }
    }

    private suspend fun handleApproval(
        approvalCtx: ApprovalContext,
        ctx: RunContext,
        setExitCode: (Int) -> Unit,
    ): Boolean {
        val decision = resolveDecision(ctx, approvalCtx.toolName, approvalCtx.preview)

        val steeringNote = if (decision == "STEER" && ctx.isTty) {
            print("Steering note: ")
            readlnOrNull()
        } else {
            null
        }

        val response = ApprovalResponsePayload(
            type = "ApprovalResponse",
            requestId = approvalCtx.requestId,
            decision = decision,
            steeringNote = steeringNote,
        )
        approvalCtx.session.send(serializer.encodeToString(response))

        return if (decision == "REJECT" && !ctx.autoApprove && !ctx.isTty) {
            setExitCode(2)
            false
        } else {
            true
        }
    }

    private fun resolveDecision(ctx: RunContext, toolName: String, preview: String?): String = when {
        ctx.autoApprove -> "APPROVE"
        !ctx.isTty -> {
            System.err.println(
                "WARNING: approval required but no TTY and --auto-approve not set — denying",
            )
            "REJECT"
        }

        else -> promptApproval(toolName, preview)
    }

    private fun promptApproval(toolName: String, preview: String?): String {
        println("\nApproval required for tool: $toolName")
        preview?.let { println("Preview: $it") }
        print("[a]pprove / [r]eject / [s]teer: ")
        return when (readlnOrNull()?.trim()?.lowercase()) {
            "a", "approve" -> "APPROVE"
            "s", "steer" -> "STEER"
            else -> "REJECT"
        }
    }

    private fun printLine(event: String, data: Map<String, String?>, outputJson: Boolean) {
        if (outputJson) {
            val fields = data.entries.joinToString(", ") { (k, v) ->
                "\"$k\": ${if (v == null) "null" else "\"$v\""}"
            }
            println("{\"event\": \"$event\", $fields}")
        } else {
            val details = data.entries.joinToString(" ") { (k, v) -> "$k=${v ?: "null"}" }
            println("[$event] $details")
        }
    }
}

class SystemExitException(code: Int) : Exception("exit $code")
