package com.correx.infrastructure

import com.correx.core.approvals.ApprovalProjector
import com.correx.core.approvals.DefaultApprovalReducer
import com.correx.core.approvals.DefaultApprovalRepository
import com.correx.core.journal.DecisionJournalProjector
import com.correx.core.journal.DefaultDecisionJournalReducer
import com.correx.core.journal.DefaultDecisionJournalRepository
import com.correx.core.artifacts.DefaultArtifactReducer
import com.correx.core.artifacts.repository.ArtifactRepository
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.config.EmbedderConfig
import com.correx.core.config.L3Config
import com.correx.core.config.ModelConfig
import com.correx.core.config.ModelsSettings
import com.correx.core.events.EventDispatcher
import com.correx.core.events.stores.EventStore
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.Embedder
import com.correx.core.inference.GenerationConfig
import com.correx.core.inference.InferenceProvider
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.ModelCapability
import com.correx.core.inference.NoopEmbedder
import com.correx.core.inference.Tokenizer
import com.correx.core.talkie.DefaultTalkieContextBuilder
import com.correx.core.talkie.DefaultTalkieFacade
import com.correx.core.talkie.IdeaReader
import com.correx.core.talkie.DefaultTalkieReducer
import com.correx.core.talkie.DefaultTalkieRepository
import com.correx.core.talkie.TalkieFacade
import com.correx.core.talkie.TalkieProjector
import com.correx.core.talkie.l3.InMemoryL3MemoryStore
import com.correx.core.talkie.l3.L3MemoryStore
import com.correx.core.talkie.model.TalkieConfig
import com.correx.core.talkie.model.WorkflowSummary
import com.correx.core.events.types.SessionId
import com.correx.core.sessions.projections.replay.DefaultEventReplayer
import com.correx.core.tools.contract.Tool
import com.correx.core.tools.contract.ToolExecutor
import com.correx.core.tools.registry.ToolRegistry
import com.correx.infrastructure.artifactscas.CasArtifactStore
import com.correx.infrastructure.artifactscas.config.CasConfig
import com.correx.infrastructure.artifactscas.index.SqliteArtifactIndex
import com.correx.infrastructure.inference.DefaultProviderRegistry
import com.correx.infrastructure.inference.commons.ModelDescriptor
import com.correx.infrastructure.inference.commons.ResidencyMode
import com.correx.infrastructure.inference.llama.cpp.DefaultModelManager
import com.correx.infrastructure.inference.llama.cpp.LlamaCppEmbedder
import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider
import com.correx.infrastructure.inference.openai.OpenAiCompatInferenceProvider
import io.ktor.client.HttpClient
import io.ktor.client.engine.cio.CIO
import com.correx.infrastructure.router.turbovec.TurboVecL3MemoryStore
import com.correx.infrastructure.router.turbovec.TurboVecSidecarConfig
import com.correx.infrastructure.persistence.SqliteEventStore
import com.correx.infrastructure.persistence.artifact.LiveArtifactRepository
import com.correx.infrastructure.tools.DefaultToolRegistry
import com.correx.infrastructure.tools.DispatchingToolExecutor
import com.correx.infrastructure.tools.SandboxedToolExecutor
import com.correx.infrastructure.tools.ToolConfig
import com.correx.infrastructure.tools.buildTools
import com.correx.core.artifacts.kind.ArtifactKind
import com.correx.core.artifacts.kind.DefaultArtifactKindRegistry
import com.correx.infrastructure.workflow.FileSystemPromptLoader
import com.correx.infrastructure.workflow.PromptLoader
import com.correx.infrastructure.workflow.TomlWorkflowLoader
import com.correx.infrastructure.workflow.WorkflowLoader
import kotlinx.coroutines.runBlocking
import org.slf4j.LoggerFactory
import java.nio.file.Files
import java.nio.file.Path
import java.nio.file.Paths
import java.nio.file.StandardCopyOption
import java.sql.DriverManager

private val log = LoggerFactory.getLogger(InfrastructureModule::class.java)

@Suppress("TooManyFunctions")
object InfrastructureModule {
    private val defaultDbPath: String =
        "${System.getProperty("user.home")}/.config/correx/correx.db"

    private val defaultArtifactsRoot: String =
        "${System.getProperty("user.home")}/.config/correx/artifacts"

    private const val LLAMA_CODING_SCORE = 0.7
    private const val LLAMA_REASONING_SCORE = 0.6
    private const val LLAMA_SUMMARIZATION_SCORE = 0.8
    private const val LLAMA_TOOL_CALLING_SCORE = 0.5

    private val DEFAULT_LLAMA_CAPABILITIES: Set<CapabilityScore> = setOf(
        CapabilityScore(ModelCapability.General, 1.0),
        CapabilityScore(ModelCapability.Coding, LLAMA_CODING_SCORE),
        CapabilityScore(ModelCapability.Reasoning, LLAMA_REASONING_SCORE),
        CapabilityScore(ModelCapability.Summarization, LLAMA_SUMMARIZATION_SCORE),
        CapabilityScore(ModelCapability.ToolCalling, LLAMA_TOOL_CALLING_SCORE),
    )

    fun createEventStore(artifactStore: ArtifactStore, dbPath: String = defaultDbPath): EventStore {
        val dir = java.io.File(dbPath).parentFile
        if (!dir.exists()) dir.mkdirs()
        return SqliteEventStore(
            connection = DriverManager.getConnection("jdbc:sqlite:$dbPath"),
            artifactStore = artifactStore,
        )
    }

    fun createArtifactStore(rootDir: Path = Paths.get(defaultArtifactsRoot)): ArtifactStore {
        Files.createDirectories(rootDir)
        val index = SqliteArtifactIndex(rootDir.resolve("index.sqlite"))
        return runBlocking { CasArtifactStore.open(CasConfig(rootDir), index) }
    }

    fun createToolRegistry(
        config: ToolConfig,
        extraTools: List<Tool> = emptyList(),
    ): ToolRegistry =
        DefaultToolRegistry.build(config.buildTools() + extraTools)

    fun createLlamaCppProvider(
        modelId: String,
        modelPath: String,
        baseUrl: String = "http://127.0.0.1:10000",
        residencyMode: ResidencyMode = ResidencyMode.PERSISTENT,
        capabilities: Set<CapabilityScore> = DEFAULT_LLAMA_CAPABILITIES,
    ): LlamaCppInferenceProvider = LlamaCppInferenceProvider(
        descriptor = ModelDescriptor(
            modelId = modelId,
            modelPath = modelPath,
            residencyMode = residencyMode,
            capabilities = capabilities,
        ),
        baseUrl = baseUrl,
    )

    fun createOpenAiCompatProvider(
        modelId: String,
        baseUrl: String,
        apiKey: String,
        idPrefix: String = "openai",
        capabilities: Set<CapabilityScore> = DEFAULT_LLAMA_CAPABILITIES,
    ): OpenAiCompatInferenceProvider = OpenAiCompatInferenceProvider(
        modelId = modelId,
        baseUrl = baseUrl,
        apiKey = apiKey,
        idPrefix = idPrefix,
        capabilities = capabilities,
    )

    fun createModelManager(
        settings: ModelsSettings,
        eventStore: EventStore,
        eventDispatcher: EventDispatcher,
        httpClient: HttpClient = HttpClient(CIO),
        healthTimeoutMs: Long = 30_000L,
    ): DefaultModelManager = DefaultModelManager(
        llamaServerBin = settings.llamaServerBin,
        host = settings.host,
        port = settings.port,
        healthTimeoutMs = healthTimeoutMs,
        eventStore = eventStore,
        httpClient = httpClient,
        eventDispatcher = eventDispatcher,
    )

    fun modelConfigToDescriptor(config: ModelConfig): ModelDescriptor = ModelDescriptor(
        modelId = config.id,
        modelPath = config.modelPath,
        residencyMode = ResidencyMode.PERSISTENT,
        contextSize = config.contextSize,
        capabilities = parseCapabilitiesFromConfig(config.capabilities),
        // [[models]] `params` (a flag->value map) passes straight through to llama-server, e.g.
        // draft-MTP speculative decoding. Flatten to token order; each stays a distinct argv entry
        // (no shell), so paths with spaces are safe.
        extraArgs = config.params.flatMap { (flag, value) -> listOf(flag, value) },
    )

    private fun parseCapabilitiesFromConfig(capsMap: Map<String, Double>): Set<CapabilityScore> {
        val result = mutableSetOf<CapabilityScore>()
        capsMap.forEach { (capName, score) ->
            val capability = when (capName.lowercase()) {
                "general" -> ModelCapability.General
                "coding" -> ModelCapability.Coding
                "reasoning" -> ModelCapability.Reasoning
                "summarization" -> ModelCapability.Summarization
                "toolcalling" -> ModelCapability.ToolCalling
                else -> {
                    log.warn("Unknown capability name '{}' in [[models]] config", capName)
                    return@forEach
                }
            }
            result.add(CapabilityScore(capability, score))
        }
        return result
    }

    fun createProviderRegistry(
        providers: List<InferenceProvider> = emptyList(),
    ): DefaultProviderRegistry = DefaultProviderRegistry(providers)

    fun createArtifactRepository(eventStore: EventStore): ArtifactRepository =
        LiveArtifactRepository(eventStore, DefaultArtifactReducer())

    fun createApprovalRepository(eventStore: EventStore): DefaultApprovalRepository =
        DefaultApprovalRepository(
            DefaultEventReplayer(
                eventStore,
                ApprovalProjector(
                    DefaultApprovalReducer(),
                ),
            ),
        )

    fun createDecisionJournalRepository(eventStore: EventStore): DefaultDecisionJournalRepository =
        DefaultDecisionJournalRepository(
            DefaultEventReplayer(
                eventStore,
                DecisionJournalProjector(DefaultDecisionJournalReducer()),
            ),
        )

    fun createWorkflowLoader(
        extraKinds: List<ArtifactKind> = emptyList(),
        samplingDefaults: GenerationConfig = GenerationConfig(temperature = 0.7, topP = 1.0, maxTokens = 0),
    ): WorkflowLoader {
        val registry = DefaultArtifactKindRegistry()
        extraKinds.forEach { registry.register(it) }
        return TomlWorkflowLoader(registry, samplingDefaults)
    }

    fun createPromptLoader(): PromptLoader = FileSystemPromptLoader()

    fun createToolExecutor(
        registry: ToolRegistry,
        eventDispatcher: EventDispatcher,
        workDir: Path,
        artifactStore: ArtifactStore? = null,
    ): ToolExecutor = SandboxedToolExecutor(
        delegate = DispatchingToolExecutor(registry),
        registry = registry,
        eventDispatcher = eventDispatcher,
        workDir = workDir,
        artifactStore = artifactStore,
    )

    fun createNoopEmbedder(dimension: Int = 1536): Embedder = NoopEmbedder(dimension)

    fun createInMemoryL3MemoryStore(): L3MemoryStore = InMemoryL3MemoryStore()

    fun createLlamaCppEmbedder(config: EmbedderConfig): Embedder {
        val url = config.url ?: throw IllegalArgumentException(
            "EmbedderConfig with backend 'llamacpp' requires non-null 'url'"
        )
        return LlamaCppEmbedder(
            url = url,
            modelId = config.modelId,
            dimension = config.dimension,
        )
    }

    fun createTurboVecL3MemoryStore(config: L3Config): L3MemoryStore {
        val scriptPath = config.scriptPath?.let { Paths.get(it) }
            ?: extractBundledTurboVecScript()

        val persistPath = config.persistPath?.let { Paths.get(it) }

        val sidecarConfig = TurboVecSidecarConfig(
            pythonExecutable = config.pythonExecutable,
            scriptPath = scriptPath,
            dim = config.dim,
            bitWidth = config.bitWidth,
            persistPath = persistPath,
        )

        return TurboVecL3MemoryStore(sidecarConfig)
    }

    private fun extractBundledTurboVecScript(): Path {
        val cacheDir = Paths.get(System.getProperty("user.home"), ".cache", "correx")
        Files.createDirectories(cacheDir)

        val targetPath = cacheDir.resolve("turbovec_sidecar.py")

        val inputStream = InfrastructureModule::class.java.getResourceAsStream("/python/turbovec_sidecar.py")
            ?: throw IllegalStateException("Could not find bundled turbovec_sidecar.py in classpath")

        inputStream.use { input ->
            Files.copy(input, targetPath, StandardCopyOption.REPLACE_EXISTING)
        }

        return targetPath
    }

    fun createEmbedderFromConfig(config: EmbedderConfig): Embedder {
        return when (config.backend) {
            "noop" -> createNoopEmbedder(config.dimension)
            "llamacpp" -> createLlamaCppEmbedder(config)
            else -> throw IllegalArgumentException(
                "Unknown embedder backend: '${config.backend}'. Supported: noop, llamacpp"
            )
        }
    }

    fun createL3MemoryStoreFromConfig(config: L3Config): L3MemoryStore {
        return when (config.backend) {
            "in_memory" -> {
                log.warn(
                    "L3 backend is 'in_memory': non-durable, intended for tests/dev only — " +
                        "all router memory is lost on restart. " +
                        "Set [l3] backend = \"turbovec\" for persistence."
                )
                createInMemoryL3MemoryStore()
            }
            "turbovec" -> createTurboVecL3MemoryStore(config)
            else -> throw IllegalArgumentException(
                "Unknown L3 backend: '${config.backend}'. Supported: in_memory, turbovec"
            )
        }
    }

    fun createTalkieFacade(
        eventStore: EventStore,
        inferenceRouter: InferenceRouter,
        config: TalkieConfig = TalkieConfig(),
        tokenizer: Tokenizer? = null,
        embedder: Embedder = createNoopEmbedder(),
        l3MemoryStore: L3MemoryStore = createInMemoryL3MemoryStore(),
        workflowSummaryProvider: () -> List<WorkflowSummary> = { emptyList() },
        sessionProfileProvider: suspend (SessionId) -> String? = { null },
    ): TalkieFacade {
        val reducer = DefaultTalkieReducer()
        val projector = TalkieProjector(reducer)
        val replayer = DefaultEventReplayer(eventStore, projector)
        val repository = DefaultTalkieRepository(replayer)
        val ideaReader = IdeaReader(eventStore)
        val contextBuilder = DefaultTalkieContextBuilder(config, tokenizer, ideaProvider = { ideaReader.activeIdeas() })
        return DefaultTalkieFacade(
            repository,
            contextBuilder,
            inferenceRouter,
            eventStore,
            config,
            null,
            embedder,
            l3MemoryStore,
            workflowSummaryProvider,
            sessionProfileProvider,
        )
    }
}
