package net.kentonc.si

import android.app.Application
import android.net.Uri
import android.provider.OpenableColumns
import androidx.lifecycle.AndroidViewModel
import androidx.lifecycle.viewModelScope
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CompletableDeferred
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.Job
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.StateFlow
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext
import net.kentonc.si.agent.Agent
import net.kentonc.si.agent.AgentImage
import net.kentonc.si.agent.IMAGE_MIMES
import net.kentonc.si.agent.Config
import net.kentonc.si.agent.Item
import net.kentonc.si.agent.SettingsStore
import net.kentonc.si.agent.ToolCall
import net.kentonc.si.agent.ToolHost
import net.kentonc.si.agent.summary
import java.io.File
import java.util.concurrent.atomic.AtomicLong

/**
 * Holds conversation state for the Compose UI. The streaming updates arrive one
 * Item at a time from the agent and are folded into an immutable row list; the
 * approval callback suspends until the dialog raised by the UI answers it.
 */
class AgentViewModel(app: Application) : AndroidViewModel(app) {

    /** One entry in the chat transcript. */
    sealed interface UiRow {
        val id: Long
    }
    enum class Role { User, Assistant }

    data class MessageRow(override val id: Long, val role: Role, val markdown: String, val images: Int) : UiRow
    data class ToolRow(override val id: Long, val name: String, val summary: String, val outcome: String?) : UiRow
    data class ErrorRow(override val id: Long, val text: String) : UiRow

    /** A mutating tool call waiting for the user's yes/no. */
    data class ApprovalRequest(val name: String, val detail: String)

    /** An image picked for the next message. */
    data class Attachment(val name: String, val mime: String, val bytes: ByteArray)

    val store = SettingsStore(app)
    val rootDir: File = app.getExternalFilesDir(null) ?: app.filesDir
    private val tools = ToolHost(rootDir)
    private val agent = Agent({ store.config }, tools, ::requestApproval, ::onItem)

    private val nextId = AtomicLong(1)

    private val _rows = MutableStateFlow<List<UiRow>>(emptyList())
    val rows: StateFlow<List<UiRow>> = _rows

    private val _busy = MutableStateFlow(false)
    val busy: StateFlow<Boolean> = _busy

    private val _approval = MutableStateFlow<ApprovalRequest?>(null)
    val approval: StateFlow<ApprovalRequest?> = _approval

    private val _attachments = MutableStateFlow<List<Attachment>>(emptyList())
    val attachments: StateFlow<List<Attachment>> = _attachments

    private val _showSettings = MutableStateFlow(false)
    val showSettings: StateFlow<Boolean> = _showSettings

    // In-app self-update (UpdateManager.kt for the hub round-trip).
    private val updateManager = UpdateManager(app, viewModelScope)
    val update: StateFlow<UpdateStatus> = updateManager.status

    /** Auto-check at startup is silent; the top-bar button retries loudly. */
    fun checkForUpdates(silent: Boolean = true) = updateManager.check(silent)

    /** Top-bar tap: download+install, retry, or resume after the grant. */
    fun startUpdate() = updateManager.start()

    private var sendJob: Job? = null

    // ─── chat ──────────────────────────────────────────────────────────────────

    fun send(text: String) {
        if (_busy.value) return
        val draft = text.trim()
        val images = _attachments.value
        if (draft.isEmpty() && images.isEmpty()) return
        _attachments.value = emptyList()

        append(MessageRow(nextId.getAndIncrement(), Role.User, draft, images.size))
        _busy.value = true
        sendJob = viewModelScope.launch {
            try {
                agent.send(draft, images.map { AgentImage(it.mime, it.bytes) })
            } catch (e: CancellationException) {
                append(ErrorRow(nextId.getAndIncrement(), "stopped"))
                throw e
            } catch (e: Exception) {
                append(ErrorRow(nextId.getAndIncrement(), e.message ?: e.javaClass.simpleName))
            } finally {
                _busy.value = false
            }
        }
    }

    /** Abort the in-flight turn. History was already rolled back by the agent. */
    fun stop() {
        sendJob?.cancel()
    }

    fun clearConversation() {
        if (_busy.value) stop()
        agent.reset()
        _rows.value = emptyList()
    }

    // ─── attachments ───────────────────────────────────────────────────────────

    /** Read a picked image's bytes and stage it for the next message. */
    fun addAttachment(uri: Uri) {
        viewModelScope.launch {
            val ctx = getApplication<Application>()
            try {
                withContext(Dispatchers.IO) {
                    val bytes = ctx.contentResolver.openInputStream(uri)?.use { it.readBytes() }
                        ?: error("couldn't read image")
                    if (bytes.size > MAX_IMAGE_BYTES) error("image is ${bytes.size / 1_000_000} MB, over the $MAX_IMAGE_BYTES MB limit")

                    val mime = imageMimeOf(ctx, uri)
                        ?: error("unsupported image type (supported: png, jpg, jpeg, gif, webp)")
                    Attachment(displayName(ctx, uri), mime, bytes)
                }.let { attachment ->
                    _attachments.update { it + attachment }
                }
            } catch (e: Exception) {
                append(ErrorRow(nextId.getAndIncrement(), "image attachment failed: ${e.message}"))
            }
        }
    }

    fun removeAttachment(name: String) {
        _attachments.update { list -> list.filterNot { it.name == name } }
    }

    private fun imageMimeOf(ctx: Application, uri: Uri): String? {
        val declared = ctx.contentResolver.getType(uri)?.lowercase()
        declared?.takeIf { it in IMAGE_MIMES.values }?.let { return it }
        val name = displayName(ctx, uri).lowercase()
        val ext = name.substringAfterLast('.', "").ifBlank { null }
        return ext?.let { IMAGE_MIMES[it] }
    }

    private fun displayName(ctx: Application, uri: Uri): String {
        ctx.contentResolver.query(uri, arrayOf(OpenableColumns.DISPLAY_NAME), null, null, null)
            ?.use { cursor ->
                if (cursor.moveToFirst()) {
                    cursor.getString(0)?.let { return it }
                }
            }
        return uri.lastPathSegment ?: "image"
    }

    // ─── approval ──────────────────────────────────────────────────────────────

    /** Called by ToolHost before anything mutating runs; parks the tool until
     *  the approval dialog answers. Defaults to nothing mutating without a tap. */
    private suspend fun requestApproval(call: ToolCall): Boolean {
        val deferred = CompletableDeferred<Boolean>()
        // Park the deferred where answerApproval can complete it, so tapping
        // Run it / Decline actually releases the tool.
        pendingApproval = deferred
        try {
            _approval.value = ApprovalRequest(call.function.name, call.summary())
            return deferred.await()
        } catch (ce: CancellationException) {
            throw ce
        } finally {
            _approval.value = null
            pendingApproval = null
        }
    }

    fun answerApproval(allow: Boolean) {
        _approval.value = null
        pendingApproval?.complete(allow)
        pendingApproval = null
    }

    private var pendingApproval: CompletableDeferred<Boolean>? = null

    // ─── settings ──────────────────────────────────────────────────────────────

    fun toggleSettings(open: Boolean) {
        _showSettings.value = open
    }

    fun saveSettings(new: Config) {
        store.save(new)
        toggleSettings(open = false)
    }

    // ─── item folding ──────────────────────────────────────────────────────────

    private fun append(row: UiRow) {
        _rows.update { it + row }
    }

    private fun onItem(item: Item) {
        when (item) {
            is Item.Text -> {
                // Streaming reply: grow the trailing assistant row, or start one
                // after a tool interruption.
                _rows.update { list ->
                    val last = list.lastOrNull()
                    if (last is MessageRow && last.role == Role.Assistant) {
                        list.dropLast(1) + last.copy(markdown = last.markdown + item.delta)
                    } else {
                        list + MessageRow(nextId.getAndIncrement(), Role.Assistant, item.delta, images = 0)
                    }
                }
            }
            is Item.ToolStart ->
                append(ToolRow(nextId.getAndIncrement(), item.name, item.summary, outcome = null))
            is Item.ToolDone -> _rows.update { list ->
                val index = list.indexOfLast { it is ToolRow && (it as ToolRow).outcome == null }
                if (index < 0) list
                else list.mapIndexed { i, row ->
                    if (i == index) (row as ToolRow).copy(outcome = item.outcome) else row
                }
            }
        }
    }

    companion object {
        /** Cap on attached images (base64 inflates ~33%) — the Rust MAX_IMAGE_BYTES. */
        const val MAX_IMAGE_BYTES = 20 * 1_000_000
    }
}