package com.correx.infrastructure.inference.llama.cpp

import com.correx.core.events.EventDispatcher
import com.correx.core.events.events.ModelLoadedEvent
import com.correx.core.events.events.ModelUnloadedEvent
import com.correx.core.events.stores.EventStore
import com.correx.core.events.types.ProviderId
import com.correx.core.events.types.SessionId
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.ProviderHealth
import com.correx.core.inference.Tokenizer
import com.correx.infrastructure.inference.commons.ManagedInferenceProvider
import com.correx.infrastructure.inference.commons.ModelDescriptor
import com.correx.infrastructure.inference.commons.ModelLoadException
import com.correx.infrastructure.inference.commons.ModelManager
import io.ktor.client.HttpClient
import io.ktor.client.call.body
import io.ktor.client.request.get
import kotlinx.coroutines.delay
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.datetime.Clock
import java.io.File

/**
 * Implementation of ModelManager that manages llama.cpp server process lifecycle.
 * 
 * Enforces single-model-at-a-time policy via Mutex.
 * Model swaps are performed via process restart (kill existing, spawn new with --model flag).
 * 
 * Thread-safe: all public methods use Mutex for synchronization.
 */
@Suppress("LongParameterList", "TooGenericExceptionCaught")
class DefaultModelManager(
    private val llamaServerBin: String = "llama-server",
    private val host: String = "127.0.0.1",
    private val port: Int = 10000,
    private val healthTimeoutMs: Long = 30_000L,
    @Suppress("UnusedPrivateProperty")
    private val eventStore: EventStore,
    private val httpClient: HttpClient,
    private val eventDispatcher: EventDispatcher,
) : ModelManager {

    private var currentProcess: LlamaProcess? = null
    private var currentDescriptor: ModelDescriptor? = null
    private val mutex = Mutex()

    private val logDirPath = System.getProperty("user.home") + "/.local/share/correx/logs"

    override suspend fun load(descriptor: ModelDescriptor): InferenceProvider = mutex.withLock {
        // If same model already loaded, return existing provider
        if (currentDescriptor?.modelId == descriptor.modelId) {
            return@withLock createProvider(descriptor)
        }

        // Kill existing process if running
        stopCurrentModelInternal()

        // Spawn new process
        val process = try {
            val logFile = File(logDirPath, "llama-server.log")
            val command = listOf(
                llamaServerBin,
                "--model", descriptor.modelPath,
                "--ctx-size", descriptor.contextSize.toString(),
                "--host", host,
                "--port", port.toString(),
            ) + descriptor.extraArgs
            val lp = LlamaProcess(command, logFile)
            lp.start()
            lp
        } catch (e: Exception) {
            throw ModelLoadException("Failed to start llama-server: ${e.message}", descriptor.modelId, e)
        }

        // Health check polling
        if (!waitForHealthy(process)) {
            process.stop()
            throw ModelLoadException(
                "Health check timeout after ${healthTimeoutMs}ms",
                descriptor.modelId,
            )
        }

        // Set state
        currentProcess = process
        currentDescriptor = descriptor

        // Emit event
        val providerId = ProviderId("llama-cpp:${descriptor.modelId}")
        eventDispatcher.emit(
            payload = ModelLoadedEvent(
                modelId = descriptor.modelId,
                providerId = providerId,
                sessionId = SessionId("correx:model-manager"),
            ),
            sessionId = SessionId("correx:model-manager")
        )

        // Return provider
        createProvider(descriptor)
    }

    override suspend fun unload(modelId: String) = mutex.withLock {
        val process = currentProcess ?: throw ModelLoadException("No model currently loaded", modelId)
        val descriptor = currentDescriptor ?: throw ModelLoadException("No model currently loaded", modelId)

        if (descriptor.modelId != modelId) {
            throw ModelLoadException("Model ID mismatch: expected ${descriptor.modelId}, got $modelId", modelId)
        }

        process.stop()
        currentProcess = null
        currentDescriptor = null

        // Emit event
        val providerId = ProviderId("llama-cpp")
        eventDispatcher.emit(
            payload = ModelUnloadedEvent(
                modelId = modelId,
                providerId = providerId,
                sessionId = SessionId("correx:model-manager"),
                cancellationReason = null,
            ),
            sessionId = SessionId("correx:model-manager")
        )
    }

    override fun currentModel(): ModelDescriptor? = currentDescriptor

    /**
     * The managed `llama-server` OS pid, or null when no model is resident. Read without the mutex
     * for the live resource gauge — a stale/null read for one probe tick during a swap is acceptable
     * (telemetry only, feeds no core decision).
     */
    fun currentPid(): Long? = currentProcess?.pid

    override fun healthCheck(): ProviderHealth {
        val process = currentProcess ?: return ProviderHealth.Unavailable("No model loaded")
        return if (process.isAlive) {
            ProviderHealth.Healthy
        } else {
            ProviderHealth.Unavailable("Process not running")
        }
    }

    private suspend fun stopCurrentModelInternal() {
        currentProcess?.stop()
        currentProcess = null
        currentDescriptor = null
    }

    private suspend fun waitForHealthy(process: LlamaProcess): Boolean {
        val endTime = Clock.System.now().toEpochMilliseconds() + healthTimeoutMs
        while (Clock.System.now().toEpochMilliseconds() < endTime) {
            try {
                val response = httpClient.get("http://$host:$port/health").body<String>()
                if (response.contains("\"status\":\"healthy\"") || response.contains("ok")) {
                    return true
                }
            } catch (_: Exception) {
                // Continue polling
            }
            // A crashed process (bad --model, port clash) never gets healthy; polling the full
            // timeout is pure dead time. Bail as soon as it's gone — but only after probing health,
            // so a server that's already answering isn't missed on a lost startup/exit race.
            if (!process.isAlive) return false
            delay(POLL_INTERVAL)
        }
        return false
    }

    private fun createProvider(descriptor: ModelDescriptor): ManagedInferenceProvider {
        val providerId = ProviderId("llama-cpp:${descriptor.modelId}")
        val baseUrl = "http://$host:$port"

        val delegate = LlamaCppInferenceProvider(
            descriptor = descriptor,
            baseUrl = baseUrl,
            httpClient = httpClient,
        )

        return DefaultManagedInferenceProvider(delegate, this, descriptor, providerId)
    }

    private companion object {
        private const val POLL_INTERVAL = 1000L
    }
}

/**
 * Implementation of [ManagedInferenceProvider] that tracks model lifecycle.
 * Ensures the provider is tied to a specific [ModelManager] instance and [ModelDescriptor].
 */
class DefaultManagedInferenceProvider(
    private val delegate: InferenceProvider,
    private val manager: ModelManager,
    private val descriptor: ModelDescriptor,
    override val id: ProviderId,
) : ManagedInferenceProvider, InferenceProvider by delegate {

    override val name: String = "Managed(${delegate.name})"
    override val tokenizer: Tokenizer = delegate.tokenizer

    override suspend fun unload() {
        manager.unload(descriptor.modelId)
    }

    override fun isLoaded(): Boolean {
        return manager.currentModel()?.modelId == descriptor.modelId
    }

    override suspend fun load(newDescriptor: ModelDescriptor): ManagedInferenceProvider {
        val provider = manager.load(newDescriptor) as? ManagedInferenceProvider
            ?: throw ModelLoadException("Failed to load new model", newDescriptor.modelId)
        return provider
    }
}
