package com.correx.apps.server

import com.correx.apps.server.logging.LoggingEventStore
import com.correx.apps.server.registry.FileSystemWorkflowRegistry
import com.correx.apps.server.registry.ProviderRegistry
import com.correx.apps.server.workspace.WorkspaceResolver
import com.correx.core.approvals.domain.DefaultApprovalEngine
import com.correx.core.artifacts.kind.ArtifactKind
import com.correx.core.artifacts.kind.ConfigArtifactKind
import com.correx.core.artifacts.kind.JsonSchema
import com.correx.core.artifactstore.ArtifactStore
import com.correx.core.config.ConfigLoader
import com.correx.core.config.CorrexConfig
import com.correx.core.config.ProviderConfig
import com.correx.core.context.builder.DefaultContextPackBuilder
import com.correx.core.context.compression.DefaultContextCompressor
import com.correx.core.context.model.TokenBudget
import com.correx.core.events.EventDispatcher
import com.correx.core.events.stores.EventStore
import com.correx.core.inference.CapabilityScore
import com.correx.core.inference.DefaultInferenceRouter
import com.correx.core.inference.GenerationConfig
import com.correx.core.inference.InferenceProjector
import com.correx.core.inference.InferenceRepository
import com.correx.core.inference.InferenceRouter
import com.correx.core.inference.ModelCapability
import com.correx.core.kernel.orchestration.DefaultOrchestrationReducer
import com.correx.core.kernel.orchestration.DefaultSessionOrchestrator
import com.correx.core.kernel.orchestration.OrchestrationConfig
import com.correx.core.kernel.orchestration.OrchestrationProjector
import com.correx.core.kernel.orchestration.OrchestrationRepository
import com.correx.core.kernel.orchestration.OrchestratorEngines
import com.correx.core.kernel.orchestration.OrchestratorRepositories
import com.correx.core.kernel.orchestration.WorkspaceContext
import com.correx.core.kernel.orchestration.WorkspaceTools
import com.correx.core.kernel.orchestration.WorkspaceToolRegistryProvider
import com.correx.core.kernel.retry.DefaultRetryCoordinator
import com.correx.core.risk.DefaultRiskAssessor
import com.correx.core.router.l3.L3MetadataRehydrator
import com.correx.core.router.model.RouterConfig
import com.correx.core.sessions.DefaultSessionReducer
import com.correx.core.sessions.DefaultSessionRepository
import com.correx.core.sessions.SessionProjector
import com.correx.core.sessions.projections.replay.DefaultEventReplayer
import com.correx.core.tools.registry.ToolRegistry
import com.correx.core.transitions.resolution.DefaultTransitionResolver
import com.correx.core.validation.artifact.ArtifactPayloadValidator
import com.correx.core.validation.pipeline.ValidationPipeline
import com.correx.core.validation.semantic.SemanticValidator
import com.correx.core.validation.semantic.rules.CycleExitRule
import com.correx.core.validation.semantic.rules.MissingToolRule
import com.correx.core.toolintent.ToolCallAssessor
import com.correx.core.toolintent.WorkspacePolicy
import com.correx.core.toolintent.rules.ExecInterpreterRule
import com.correx.core.toolintent.rules.NetworkHostRule
import com.correx.core.toolintent.rules.PathContainmentRule
import com.correx.apps.server.freestyle.FreestyleDriver
import com.correx.core.artifacts.kind.DefaultArtifactKindRegistry
import com.correx.core.events.types.ArtifactId
import com.correx.infrastructure.InfrastructureModule
import com.correx.infrastructure.inference.DefaultProviderRegistry
import com.correx.infrastructure.workflow.ExecutionPlanCompiler
import com.correx.infrastructure.inference.FirstAvailableRoutingStrategy
import com.correx.infrastructure.inference.commons.ManagedInferenceRouter
import com.correx.infrastructure.inference.commons.AmdResourceProbe
import com.correx.infrastructure.inference.commons.NvidiaResourceProbe
import com.correx.infrastructure.inference.commons.ResourceProbe
import com.correx.infrastructure.inference.commons.UnavailableProbe
import com.correx.core.inference.InferenceProvider
import com.correx.infrastructure.inference.llama.cpp.LlamaCppInferenceProvider
import com.correx.infrastructure.tools.DispatchingToolExecutor
import com.correx.infrastructure.tools.FileEditConfig
import com.correx.infrastructure.tools.FileReadConfig
import com.correx.infrastructure.tools.FileWriteConfig
import com.correx.infrastructure.tools.ShellConfig
import com.correx.infrastructure.tools.ToolConfig
import io.ktor.server.engine.embeddedServer
import io.ktor.server.netty.Netty
import kotlinx.coroutines.runBlocking
import kotlinx.serialization.json.Json
import org.slf4j.LoggerFactory
import java.nio.file.Files
import java.nio.file.Path
import java.nio.file.Paths

private val log = LoggerFactory.getLogger("com.correx.apps.server.Main")

fun main() {
    log.info("=== correx server starting ===")
    log.info("  port         : 8080")
    val artifactStore = InfrastructureModule.createArtifactStore()
    val eventStore = LoggingEventStore(InfrastructureModule.createEventStore(artifactStore))

    logStoresInfo(eventStore, artifactStore)

    val correxConfig = ConfigLoader.load()
    val eventDispatcher = EventDispatcher(eventStore)

    // Managed path: [[models]] present → correx spawns llama-server at boot
    // Static path:  [[models]] absent  → legacy [[providers]] / env-var behavior
    val firstProvider: InferenceProvider
    val infraRegistry: DefaultProviderRegistry
    val inferenceRouter: InferenceRouter
    // Non-null only on the managed path; backs manual model swap/pin from clients.
    var modelSwapper: ManagedInferenceRouter? = null
    // Live GPU/RAM gauge; NVIDIA-backed on the managed path when nvidia-smi is present, else unavailable.
    var resourceProbe: ResourceProbe = UnavailableProbe

    if (correxConfig.models.isNotEmpty()) {
        val settings = correxConfig.modelsSettings
        val modelManager = InfrastructureModule.createModelManager(settings, eventStore, eventDispatcher)
        val descriptors = correxConfig.models.map { InfrastructureModule.modelConfigToDescriptor(it) }
        val targetId = settings.defaultModel ?: correxConfig.models.first().id
        val targetModelConfig = correxConfig.models.firstOrNull { it.id == targetId }
            ?: correxConfig.models.first()
        log.info("Managed path: loading model '{}' from '{}'", targetModelConfig.id, targetModelConfig.modelPath)
        val descriptor = InfrastructureModule.modelConfigToDescriptor(targetModelConfig)
        val managedProvider = runBlocking { modelManager.load(descriptor) }
        firstProvider = managedProvider
        infraRegistry = InfrastructureModule.createProviderRegistry(listOf(managedProvider))
        // Also register any static providers from [[providers]] (non-managed remotes)
        correxConfig.providers.mapNotNull { providerConfig ->
            when (providerConfig.type.lowercase()) {
                "llamacpp" -> buildProviderFromConfig(providerConfig)
                else -> { log.error("Unknown provider type: {}", providerConfig.type); null }
            }
        }.forEach { infraRegistry.register(it) }
        // Managed router owns per-stage model selection: it ensures the resolved model is resident
        // before routing (stage.modelId > capability match > default), evicting the prior model.
        val managedRouter = ManagedInferenceRouter(modelManager, descriptors, targetModelConfig.id)
        inferenceRouter = managedRouter
        modelSwapper = managedRouter
        resourceProbe = when {
            NvidiaResourceProbe.isAvailable() ->
                NvidiaResourceProbe(pidSupplier = { modelManager.currentPid() })
            AmdResourceProbe.isAvailable() ->
                AmdResourceProbe(pidSupplier = { modelManager.currentPid() })
            else -> UnavailableProbe
        }
        // Shutdown hook to kill the managed llama-server
        Runtime.getRuntime().addShutdownHook(Thread {
            log.info("Shutdown: unloading managed model '{}'", targetModelConfig.id)
            runBlocking {
                runCatching { modelManager.unload(targetModelConfig.id) }
                    .onFailure { log.warn("Error during managed model unload: {}", it.message) }
            }
        })
    } else {
        val staticProviders = buildProviders(correxConfig)
        require(staticProviders.isNotEmpty()) { "At least one provider must be configured" }
        firstProvider = staticProviders.first()
        infraRegistry = InfrastructureModule.createProviderRegistry(staticProviders)
        inferenceRouter = DefaultInferenceRouter(infraRegistry, FirstAvailableRoutingStrategy())
    }

    val modelId = firstProvider.id.value
    logModelInfo(modelId, infraRegistry)

    val repositories = buildRepositories(eventStore)
    val approvalEngine = DefaultApprovalEngine()
    val toolsConfig = correxConfig.tools
    val sandboxRoot = System.getenv("CORREX_SANDBOX_ROOT")
        ?.let { Path.of(it) }
        ?: toolsConfig.sandboxRoot.takeIf { it.isNotEmpty() }?.let { Path.of(it) }
        ?: Paths.get(System.getProperty("user.home"), ".config", "correx", "sandbox")
    val workingDir = System.getenv("CORREX_WORKING_DIR")
        ?.let { Path.of(it) }
        ?: toolsConfig.workingDir.takeIf { it.isNotEmpty() }?.let { Path.of(it) }
        ?: Path.of("").toAbsolutePath()
    val shellAllowedExecutables = toolsConfig.shellAllowedExecutables.toSet()
    val workspaceRoot = System.getenv("CORREX_WORKSPACE_ROOT")
        ?.let { Path.of(it) }
        ?: toolsConfig.workspaceRoot.takeIf { it.isNotEmpty() }?.let { Path.of(it) }
        ?: workingDir
    val toolRegistry = InfrastructureModule.createToolRegistry(
        buildToolConfig(
            workspaceRoot,
            workingDir,
            shellAllowedExecutables,
            toolsConfig,
        ),
    )
    val toolExecutor = InfrastructureModule.createToolExecutor(
        registry = toolRegistry,
        eventDispatcher = eventDispatcher,
        workDir = sandboxRoot,
        artifactStore = artifactStore,
    )

    logToolInfo(sandboxRoot, workingDir, shellAllowedExecutables, toolRegistry)

    val privilegedLocations = toolsConfig.privilegedLocations.map { Path.of(it) }
    val workspacePolicy = WorkspacePolicy(workspaceRoot, privilegedLocations)
    val toolCallAssessor = ToolCallAssessor(
        rules = listOf(
            PathContainmentRule(),
            ExecInterpreterRule(toolsConfig.interpreterExecutables.toSet()),
            NetworkHostRule(
                allowedHosts = toolsConfig.networkAllowedHosts.toSet(),
                deniedHosts = toolsConfig.networkDeniedHosts.toSet(),
            ),
        ),
    )

    val wsToolRegistryProvider = WorkspaceToolRegistryProvider { workspace ->
        val wsRegistry = InfrastructureModule.createToolRegistry(
            buildToolConfigForWorkspace(workspace, shellAllowedExecutables, toolsConfig),
        )
        val wsExecutor = DispatchingToolExecutor(wsRegistry)
        WorkspaceTools(registry = wsRegistry, executor = wsExecutor)
    }

    val engines = OrchestratorEngines(
        transitionResolver = DefaultTransitionResolver { condition, ctx -> condition.evaluate(ctx) },
        contextPackBuilder = DefaultContextPackBuilder(DefaultContextCompressor()),
        inferenceRouter = inferenceRouter,
        validationPipeline = ValidationPipeline(
            validators = listOf(
                ArtifactPayloadValidator(artifactStore),
                SemanticValidator(
                    rules = listOf(
                        CycleExitRule(requirePolicyForCycles = false),
                        MissingToolRule(),
                    ),
                ),
            ),
        ),
        approvalEngine = approvalEngine,
        riskAssessor = DefaultRiskAssessor(),
        promptResolver = { path -> InfrastructureModule.createPromptLoader().load(path) },
        toolRegistry = toolRegistry,
        toolExecutor = toolExecutor,
        toolCallAssessor = toolCallAssessor,
        workspacePolicy = workspacePolicy,
        workspaceToolRegistryProvider = wsToolRegistryProvider,
    )
    val decisionJournalRepository = InfrastructureModule.createDecisionJournalRepository(eventStore)
    val defaultOrchestrationConfig = OrchestrationConfig(
        sandboxRoot = sandboxRoot,
        defaultSystemPromptPath = toolsConfig.defaultSystemPromptPath,
    )
    val orchestrator = DefaultSessionOrchestrator(
        repositories = repositories,
        engines = engines,
        retryCoordinator = DefaultRetryCoordinator(eventStore),
        artifactStore = artifactStore,
        tokenizer = firstProvider.tokenizer,
        decisionJournalRepository = decisionJournalRepository,
    )
    val routerConfig = correxConfig.router
    val embedder = InfrastructureModule.createEmbedderFromConfig(routerConfig.embedder)
    val l3MemoryStore = InfrastructureModule.createL3MemoryStoreFromConfig(routerConfig.l3)
    // Rebuild the L3 metadata map from the ChatTurnEvent log before serving queries:
    // persisted vectors survive restart but the metadata map does not (no-op for
    // non-rehydratable backends like in_memory).
    runBlocking { L3MetadataRehydrator(eventStore).rehydrate(l3MemoryStore) }
    // Map the config-file [router] block onto the domain RouterConfig the facade uses.
    val domainRouterConfig = RouterConfig(
        conversationKeepLast = routerConfig.conversationKeepLast,
        retrievalK = routerConfig.retrievalK,
        tokenBudget = TokenBudget(limit = routerConfig.tokenBudget),
        generationConfig = GenerationConfig(
            temperature = routerConfig.generation.temperature,
            topP = routerConfig.generation.topP,
            maxTokens = routerConfig.generation.maxTokens,
        ),
        narrationGenerationConfig = GenerationConfig(
            temperature = routerConfig.narration.temperature,
            topP = routerConfig.narration.topP,
            maxTokens = routerConfig.narration.maxTokens,
        ),
    )
    val routerFacade = InfrastructureModule.createRouterFacade(
        eventStore = eventStore,
        inferenceRouter = inferenceRouter,
        config = domainRouterConfig,
        tokenizer = firstProvider.tokenizer,
        embedder = embedder,
        l3MemoryStore = l3MemoryStore,
    )
    val sessionUndoService = com.correx.apps.server.undo.SessionUndoService(
        eventStore = eventStore,
        artifactStore = artifactStore,
        bootRoots = setOf(workspaceRoot, workingDir),
    )
    val bootWorkspaceContext = WorkspaceContext(
        workspaceRoot = workspaceRoot,
        workingDir = workingDir,
        allowedPaths = setOf(workspaceRoot, workingDir),
        privilegedLocations = privilegedLocations,
    )
    val allowedWorkspaceRoots = toolsConfig.allowedWorkspaceRoots.map { Path.of(it) }
    val workspaceResolver = WorkspaceResolver(
        bootDefault = bootWorkspaceContext,
        privilegedLocations = privilegedLocations,
        allowedWorkspaceRoots = allowedWorkspaceRoots,
    )
    val projectMemory = if (correxConfig.project.enabled) {
        com.correx.apps.server.memory.ProjectMemoryService(
            config = correxConfig.project,
            embedder = embedder,
            l3MemoryStore = l3MemoryStore,
            journalRepository = decisionJournalRepository,
            eventStore = eventStore,
            indexer = com.correx.apps.server.memory.RepoMapIndexer(
                maxDepth = correxConfig.project.maxDepth,
                ignoreGlobs = correxConfig.project.ignoreGlobs,
            ),
        )
    } else {
        null
    }
    val configArtifactKinds = loadConfigArtifactKinds(correxConfig)
    val artifactKindRegistry = DefaultArtifactKindRegistry().also { reg ->
        configArtifactKinds.forEach { reg.register(it) }
    }
    val freestyleDriver = FreestyleDriver(
        eventStore = eventStore,
        compiler = ExecutionPlanCompiler(artifactKindRegistry),
        planContent = { sid -> orchestrator.validatedArtifactContent(sid, ArtifactId("execution_plan")) },
        config = defaultOrchestrationConfig,
        runPhase2 = { sid, graph, cfg -> orchestrator.run(sid, graph, cfg) },
    )
    val module = ServerModule(
        orchestrator = orchestrator,
        eventStore = eventStore,
        artifactStore = artifactStore,
        sessionRepository = repositories.sessionRepository,
        workflowRegistry = FileSystemWorkflowRegistry(
            InfrastructureModule.createWorkflowLoader(configArtifactKinds),
        ),
        providerRegistry = infraRegistry.asServerRegistry(),
        defaultOrchestrationConfig = defaultOrchestrationConfig,
        routerFacade = routerFacade,
        orchestrationRepository = repositories.orchestrationRepository,
        approvalRepository = repositories.approvalRepository,
        toolRegistry = toolRegistry,
        sessionUndoService = sessionUndoService,
        modelSwapper = modelSwapper,
        resourceProbe = resourceProbe,
        workspaceResolver = workspaceResolver,
        narrationMaxPerRun = routerConfig.narration.maxPerRun,
        projectMemory = projectMemory,
        freestyleDriver = freestyleDriver,
    )
    module.start()
    log.info("==============================")

    embeddedServer(Netty, port = 8080) { configureServer(module) }.start(wait = true)
}

private fun logModelInfo(
    modelId: String,
    infraRegistry: DefaultProviderRegistry,
) {
    log.info("==========MODEL INFO==========")
    log.info("  model id     : {}", modelId)
    log.info("  providers    : {} registered", infraRegistry.listAll().size)
}

private fun logToolInfo(
    sandboxRoot: Path?,
    workingDir: Path?,
    shellAllowedExecutables: Set<String>,
    toolRegistry: ToolRegistry,
) {
    log.info("==========TOOLS INFO==========")
    log.info("  sandbox root : {}", sandboxRoot)
    log.info("  working dir  : {}", workingDir)
    if (shellAllowedExecutables.isEmpty()) {
        log.warn("  shell tool   : no executable allowlist configured — all executables allowed")
    }
    log.info("  tools        : {}", toolRegistry.all().map { it.name })
}

private fun logStoresInfo(
    eventStore: LoggingEventStore,
    artifactStore: ArtifactStore,
) {
    log.info("==========STORE INFO==========")
    log.info("  event store  : {}", eventStore::class.simpleName)
    log.info("  artifact store: {}", artifactStore::class.simpleName)
}

// Build user-declared artifact kinds from config. Each [[artifacts]] entry points to a JSON
// schema file; we parse it into a JsonSchema and wrap it in a ConfigArtifactKind. A schema that
// can't be loaded is skipped with a warning — any workflow that references the kind then fails
// loudly at load time ("Unknown artifact kind"), rather than running with an unvalidated kind.
private fun loadConfigArtifactKinds(config: CorrexConfig): List<ArtifactKind> {
    val configDir = ConfigLoader.configPath().parent
    return config.artifacts.mapNotNull { decl ->
        val schemaPath = resolveConfigRelativePath(decl.schemaPath, configDir)
        val schema = runCatching {
            Json.decodeFromString(JsonSchema.serializer(), Files.readString(schemaPath))
        }.getOrElse { e ->
            System.err.println(
                "Warning: artifact kind '${decl.id}' schema '$schemaPath' failed to load: ${e.message}",
            )
            return@mapNotNull null
        }
        ConfigArtifactKind(id = decl.id, schema = schema, llmEmitted = decl.llmEmitted)
    }
}

private fun resolveConfigRelativePath(raw: String, configDir: Path?): Path {
    val expanded = if (raw.startsWith("~/")) {
        Paths.get(System.getProperty("user.home"), raw.removePrefix("~/"))
    } else {
        Path.of(raw)
    }
    return when {
        expanded.isAbsolute -> expanded
        configDir != null -> configDir.resolve(expanded)
        else -> expanded
    }
}

private fun buildProviders(config: CorrexConfig): List<LlamaCppInferenceProvider> {
    return if (config.providers.isNotEmpty()) {
        log.info("Loading {} provider(s) from config", config.providers.size)
        config.providers.mapNotNull { providerConfig ->
            when (providerConfig.type.lowercase()) {
                "llamacpp" -> buildProviderFromConfig(providerConfig)
                else -> {
                    log.error("Unknown provider type: {}", providerConfig.type)
                    null
                }
            }
        }
    } else {
        log.info("No providers in config; using env var fallback")
        listOf(
            InfrastructureModule.createLlamaCppProvider(
                modelId = System.getenv("CORREX_MODEL_ID") ?: "default",
                modelPath = System.getenv("CORREX_MODEL_PATH") ?: "",
                baseUrl = System.getenv("CORREX_LLAMA_URL") ?: "http://127.0.0.1:10000",
            ),
        )
    }
}

private fun buildProviderFromConfig(config: ProviderConfig): LlamaCppInferenceProvider {
    val capabilities = parseCapabilitiesFromConfig(config.capabilities)
    return InfrastructureModule.createLlamaCppProvider(
        modelId = config.modelId,
        modelPath = config.modelPath,
        baseUrl = config.url,
        capabilities = capabilities,
    )
}

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: {}", capName)
                return@forEach
            }
        }
        result.add(CapabilityScore(capability, score))
    }
    return result
}

private fun buildRepositories(
    eventStore: EventStore,
): OrchestratorRepositories = OrchestratorRepositories(
    eventStore = eventStore,
    inferenceRepository = InferenceRepository(DefaultEventReplayer(eventStore, InferenceProjector())),
    orchestrationRepository = OrchestrationRepository(
        DefaultEventReplayer(eventStore, OrchestrationProjector(DefaultOrchestrationReducer())),
    ),
    sessionRepository = DefaultSessionRepository(
        DefaultEventReplayer(eventStore, SessionProjector(DefaultSessionReducer())),
    ),
    artifactRepository = InfrastructureModule.createArtifactRepository(eventStore),
    approvalRepository = InfrastructureModule.createApprovalRepository(eventStore),
)

// File tools write IN PLACE into the workspace (cwd / allow-listed dirs). The old
// sandboxRoot dual-write is retired; reversibility is event-sourced (FileWrittenEvent +
// CAS pre/post images, captured by SandboxedToolExecutor). sandboxRoot survives only as
// the executor's scratch workDir, never as a write destination.
private fun buildToolConfig(
    workspaceRoot: Path,
    workingDir: Path,
    shellAllowedExecutables: Set<String>,
    toolsConfig: com.correx.core.config.ToolsConfig,
): ToolConfig {
    val allowed = setOf(workspaceRoot, workingDir)
    return ToolConfig(
        shell = ShellConfig(
            enabled = toolsConfig.shellEnabled,
            allowedExecutables = shellAllowedExecutables,
            workingDir = workingDir,
        ),
        fileRead = FileReadConfig(
            enabled = toolsConfig.fileReadEnabled,
            allowedPaths = allowed,
        ),
        fileWrite = FileWriteConfig(
            allowedPaths = allowed,
            enabled = toolsConfig.fileWriteEnabled,
            workingDir = workingDir,
        ),
        fileEdit = FileEditConfig(
            enabled = toolsConfig.fileEditEnabled,
            allowedPaths = allowed,
            workingDir = workingDir,
        ),
    )
}

private fun buildToolConfigForWorkspace(
    workspace: WorkspaceContext,
    shellAllowedExecutables: Set<String>,
    toolsConfig: com.correx.core.config.ToolsConfig,
): ToolConfig = ToolConfig(
    shell = ShellConfig(
        enabled = toolsConfig.shellEnabled,
        allowedExecutables = shellAllowedExecutables,
        workingDir = workspace.workingDir,
    ),
    fileRead = FileReadConfig(
        enabled = toolsConfig.fileReadEnabled,
        allowedPaths = workspace.allowedPaths,
    ),
    fileWrite = FileWriteConfig(
        enabled = toolsConfig.fileWriteEnabled,
        allowedPaths = workspace.allowedPaths,
        workingDir = workspace.workingDir,
    ),
    fileEdit = FileEditConfig(
        enabled = toolsConfig.fileEditEnabled,
        allowedPaths = workspace.allowedPaths,
        workingDir = workspace.workingDir,
    ),
)

private fun DefaultProviderRegistry.asServerRegistry(): ProviderRegistry = object : ProviderRegistry {
    override fun listAll() = this@asServerRegistry.listAll()
    override suspend fun healthCheckAll() = this@asServerRegistry.healthCheckAll()
}
