package com.optix.android.ui.screens.ai

import androidx.lifecycle.ViewModel
import androidx.lifecycle.viewModelScope
import com.optix.android.ai.AiStreamEvent
import com.optix.android.domain.model.AiMessage
import com.optix.android.domain.model.AiProviderConfig
import com.optix.android.domain.model.DEFAULT_SYSTEM_PROMPT
import com.optix.android.domain.model.HostResult
import com.optix.android.domain.model.MessageRole
import com.optix.android.domain.model.ToolChain
import com.optix.android.domain.repository.ChainRepository
import com.optix.android.domain.repository.NodeRepository
import com.optix.android.domain.repository.SettingsRepository
import com.optix.android.domain.usecase.AiBridgeUseCase
import dagger.hilt.android.lifecycle.HiltViewModel
import kotlinx.coroutines.Job
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.asStateFlow
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch
import javax.inject.Inject

data class AiUiState(
    val messages: List<AiMessage> = emptyList(),
    val isStreaming: Boolean = false,
    val streamingMessageId: Long? = null,
    val activeToolCall: String? = null,
    val error: String? = null,
    val pendingChain: ToolChain? = null,
    val chainSkippedSteps: List<String> = emptyList(),
)

@HiltViewModel
class AiViewModel @Inject constructor(
    private val aiBridgeUseCase: AiBridgeUseCase,
    private val nodeRepository: NodeRepository,
    private val settingsRepository: SettingsRepository,
    private val chainRepository: ChainRepository,
) : ViewModel() {

    private val _uiState = MutableStateFlow(AiUiState())
    val uiState: StateFlow<AiUiState> = _uiState.asStateFlow()

    private val _providerConfig = MutableStateFlow(AiProviderConfig())
    val providerConfig: StateFlow<AiProviderConfig> = _providerConfig.asStateFlow()

    private var streamJob: Job? = null

    init {
        viewModelScope.launch {
            settingsRepository.getAiConfig().collect { _providerConfig.value = it }
        }
    }

    fun sendMessage(text: String) {
        if (text.isBlank() || _uiState.value.isStreaming) return

        val userMsg = AiMessage(role = MessageRole.USER, content = text)
        val assistantId = System.currentTimeMillis() + 1
        val assistantMsg = AiMessage(id = assistantId, role = MessageRole.ASSISTANT, content = "", isStreaming = true)

        _uiState.update { state ->
            state.copy(
                messages = state.messages + userMsg + assistantMsg,
                isStreaming = true,
                streamingMessageId = assistantId,
                error = null,
                pendingChain = null,
                chainSkippedSteps = emptyList(),
            )
        }

        streamJob = viewModelScope.launch {
            val config = _providerConfig.value
            val nodes = nodeRepository.observeNodes().first()
            val systemMsg = AiMessage(
                role = MessageRole.SYSTEM,
                content = buildString {
                    append(config.systemPrompt.ifEmpty { DEFAULT_SYSTEM_PROMPT })
                    val networkContext = buildNetworkContext(nodes)
                    if (networkContext.isNotBlank()) {
                        append("\n\n")
                        append(networkContext)
                    }
                },
            )
            val history = listOf(systemMsg) + _uiState.value.messages.filter { !it.isStreaming }
            val firedToolsThisTurn = mutableSetOf<String>()

            aiBridgeUseCase(history, config).collect { event ->
                when (event) {
                    is AiStreamEvent.Token -> appendToken(assistantId, event.text)
                    is AiStreamEvent.ToolStart -> {
                        if (event.toolName in firedToolsThisTurn) return@collect
                        firedToolsThisTurn.add(event.toolName)
                        _uiState.update { it.copy(activeToolCall = event.toolName) }
                    }
                    is AiStreamEvent.ToolResult -> {
                        _uiState.update { it.copy(activeToolCall = null) }
                        val lines = event.result.lines()
                        val preview = lines.take(30).joinToString("\n")
                        val suffix = if (lines.size > 30) "\n... +${lines.size - 30} more" else ""
                        appendToken(assistantId, "\n\n> [${event.toolName}]\n```\n$preview$suffix\n```\n")
                    }
                    is AiStreamEvent.ToolChain -> {
                        _uiState.update {
                            it.copy(
                                pendingChain = event.chain,
                                chainSkippedSteps = event.skippedSteps,
                            )
                        }
                    }
                    is AiStreamEvent.Done -> finalizeStreaming(assistantId)
                }
            }
        }
    }

    fun cancelStreaming() {
        streamJob?.cancel()
        _uiState.update { state ->
            val msgs = state.messages.map { if (it.id == state.streamingMessageId) it.copy(isStreaming = false) else it }
            state.copy(messages = msgs, isStreaming = false, streamingMessageId = null, activeToolCall = null)
        }
    }

    fun clearHistory() {
        cancelStreaming()
        _uiState.update { AiUiState() }
    }

    /** Forward the pending AI-generated chain to the NETWORK tab for execution. */
    fun sendChainToNetwork() {
        val chain = _uiState.value.pendingChain ?: return
        chainRepository.sendChain(chain)
        _uiState.update { it.copy(pendingChain = null, chainSkippedSteps = emptyList()) }
    }

    fun clearPendingChain() {
        _uiState.update { it.copy(pendingChain = null, chainSkippedSteps = emptyList()) }
    }

    fun updateProviderConfig(config: AiProviderConfig) {
        viewModelScope.launch { settingsRepository.saveAiConfig(config) }
    }

    private fun appendToken(msgId: Long, token: String) {
        _uiState.update { state ->
            val msgs = state.messages.map { msg ->
                if (msg.id == msgId) msg.copy(content = msg.content + token) else msg
            }
            state.copy(messages = msgs)
        }
    }

    private fun finalizeStreaming(msgId: Long) {
        _uiState.update { state ->
            val msgs = state.messages.map { msg ->
                if (msg.id == msgId) msg.copy(isStreaming = false) else msg
            }
            state.copy(messages = msgs, isStreaming = false, streamingMessageId = null, activeToolCall = null)
        }
    }

    private fun buildNetworkContext(nodes: List<HostResult>): String {
        if (nodes.isEmpty()) return ""

        val active = nodes.filter { !it.status.equals("STALE", true) }.take(8)
        val staleCount = nodes.count { it.status.equals("STALE", true) }
        return buildString {
            append("NETWORK CONTEXT:\n")
            append("- Known nodes: ${nodes.size}\n")
            append("- Active nodes: ${active.size}\n")
            if (staleCount > 0) append("- Stale nodes: $staleCount\n")
            if (active.isNotEmpty()) {
                append("- Top nodes:\n")
                active.forEach { node ->
                    append("  - ${node.ip}")
                    node.alias?.takeIf { it.isNotBlank() }?.let { append(" [$it]") }
                    node.hostname?.takeIf { it.isNotBlank() && it != node.ip }?.let { append(" ($it)") }
                    if (node.openPorts.isNotEmpty()) append(" ports=${node.openPorts.joinToString(",")}")
                    append("\n")
                }
            }
        }.trim()
    }
}
