package com.example.aiapp

import androidx.compose.foundation.layout.Column
import androidx.compose.foundation.layout.Row
import androidx.compose.foundation.layout.Spacer
import androidx.compose.foundation.layout.fillMaxSize
import androidx.compose.foundation.layout.fillMaxWidth
import androidx.compose.foundation.layout.height
import androidx.compose.foundation.layout.padding
import androidx.compose.foundation.lazy.LazyColumn
import androidx.compose.material3.Card
import androidx.compose.material3.CircularProgressIndicator
import androidx.compose.material3.LinearProgressIndicator
import androidx.compose.material3.MaterialTheme
import androidx.compose.material3.OutlinedTextField
import androidx.compose.material3.Text
import androidx.compose.material3.TextButton
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.rememberCoroutineScope
import androidx.compose.runtime.setValue
import androidx.compose.ui.Alignment
import androidx.compose.ui.Modifier
import androidx.compose.ui.text.style.TextOverflow
import androidx.compose.ui.unit.dp
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch
import kotlinx.coroutines.withContext

/**
 * Models on the backend, and HuggingFace to get more from.
 *
 * Everything here is the server's state rather than this screen's: what is downloaded, and what is
 * downloading, are the same answers on every enrolled device, and a download started here keeps
 * going when this screen closes.
 */
@Composable
fun ModelsScreen(settings: ServerSettings, reloadToken: Int) {
    val scope = rememberCoroutineScope()
    var state by remember { mutableStateOf<LoadState<Models>>(LoadState.Loading) }
    var query by remember { mutableStateOf("") }
    var results by remember { mutableStateOf<LoadState<List<RemoteRepo>>?>(null) }
    var openRepo by remember { mutableStateOf<String?>(null) }
    var repoFiles by remember { mutableStateOf<LoadState<List<RemoteFile>>?>(null) }
    var actionError by remember { mutableStateOf<String?>(null) }

    suspend fun reload() {
        state =
            try {
                withContext(Dispatchers.IO) { LoadState.Loaded(fetchModels(settings)) }
            } catch (e: ApiException) {
                LoadState.failed(e)
            }
    }

    // Polled rather than pushed: a download belongs to the machine, not to any session, so it has
    // no event stream of its own. Keyed on the token as well, so the header's Refresh restarts the
    // loop with a read now rather than leaving the reader watching for a second and a half.
    LaunchedEffect(reloadToken) {
        while (true) {
            reload()
            delay(1500)
        }
    }

    Column(Modifier.fillMaxSize().padding(16.dp)) {
        actionError?.let {
            Text(it, color = MaterialTheme.colorScheme.error)
            Spacer(Modifier.height(8.dp))
        }

        OutlinedTextField(
            value = query,
            onValueChange = { query = it },
            label = { Text("Search HuggingFace") },
            singleLine = true,
            modifier = Modifier.fillMaxWidth(),
        )
        Spacer(Modifier.height(8.dp))
        TextButton(
            enabled = query.isNotBlank(),
            onClick = {
                openRepo = null
                results = LoadState.Loading
                scope.launch {
                    results =
                        try {
                            withContext(Dispatchers.IO) {
                                LoadState.Loaded(searchModels(settings, query))
                            }
                        } catch (e: ApiException) {
                            LoadState.failed(e)
                        }
                }
            },
        ) {
            Text("Search")
        }

        Spacer(Modifier.height(8.dp))
        LazyColumn(Modifier.fillMaxSize()) {
            when (val current = state) {
                is LoadState.Loading -> item { CircularProgressIndicator() }
                is LoadState.Error ->
                    item { Text(current.message, color = MaterialTheme.colorScheme.error) }
                is LoadState.Loaded -> {
                    if (current.value.downloads.isNotEmpty()) {
                        item { SectionLabel("Downloading") }
                        uniqueItems(current.value.downloads, key = { it.key + it.run }) { download
                            ->
                            DownloadCard(download) {
                                scope.launch {
                                    actionError =
                                        runCatching {
                                            withContext(Dispatchers.IO) {
                                                cancelDownload(settings, download.key)
                                            }
                                        }
                                            .exceptionOrNull()
                                            ?.message
                                }
                            }
                        }
                    }
                    item { SectionLabel("On the backend") }
                    if (current.value.local.isEmpty()) {
                        item {
                            Text(
                                "None yet. Search above to find one.",
                                style = MaterialTheme.typography.bodyMedium,
                                color = MaterialTheme.colorScheme.onSurfaceVariant,
                            )
                        }
                    }
                    uniqueItems(current.value.local, key = { it.key }) { model ->
                        LocalModelCard(model) {
                            scope.launch {
                                actionError =
                                    runCatching {
                                        withContext(Dispatchers.IO) {
                                            deleteModel(settings, model.key)
                                        }
                                    }
                                        .exceptionOrNull()
                                        ?.message
                                reload()
                            }
                        }
                    }
                }
            }

            results?.let { found ->
                item { SectionLabel("HuggingFace") }
                when (found) {
                    is LoadState.Loading -> item { CircularProgressIndicator() }
                    is LoadState.Error ->
                        item { Text(found.message, color = MaterialTheme.colorScheme.error) }
                    is LoadState.Loaded ->
                        uniqueItems(found.value, key = { it.id }) { repo ->
                            val open = openRepo == repo.id
                            RepoRow(repo, expanded = open) {
                                if (open) {
                                    openRepo = null
                                } else {
                                    openRepo = repo.id
                                    repoFiles = LoadState.Loading
                                    scope.launch {
                                        repoFiles =
                                            try {
                                                withContext(Dispatchers.IO) {
                                                    LoadState.Loaded(
                                                        fetchRepoFiles(settings, repo.id)
                                                    )
                                                }
                                            } catch (e: ApiException) {
                                                LoadState.failed(e)
                                            }
                                    }
                                }
                            }
                            // Inside the expanded repository's own item rather than as a section
                            // after the list: drawn after every card, a repository's files read as
                            // belonging to whichever card happened to be last.
                            if (open) {
                                when (val files = repoFiles) {
                                    null -> {}
                                    is LoadState.Loading -> CircularProgressIndicator()
                                    is LoadState.Error ->
                                        Text(files.message, color = MaterialTheme.colorScheme.error)
                                    is LoadState.Loaded ->
                                        Column {
                                            val busy =
                                                (state as? LoadState.Loaded)
                                                    ?.value
                                                    ?.downloads
                                                    .orEmpty()
                                                    .filter { it.state == "running" }
                                                    .map { it.key }
                                                    .toSet()
                                            files.value.forEach { file ->
                                                RepoFileRow(
                                                    file,
                                                    downloading = "${repo.id}/${file.path}" in busy,
                                                ) {
                                                    scope.launch {
                                                        actionError =
                                                            runCatching {
                                                                withContext(Dispatchers.IO) {
                                                                    startDownload(
                                                                        settings,
                                                                        repo.id,
                                                                        file.path,
                                                                    )
                                                                }
                                                            }
                                                                .exceptionOrNull()
                                                                ?.message
                                                        reload()
                                                    }
                                                }
                                            }
                                        }
                                }
                            }
                        }
                }
            }
        }
    }
}

@Composable
private fun SectionLabel(text: String) {
    Spacer(Modifier.height(12.dp))
    Text(text, style = MaterialTheme.typography.titleSmall)
    Spacer(Modifier.height(4.dp))
}

@Composable
private fun DownloadCard(download: Download, onCancel: () -> Unit) {
    Card(Modifier.fillMaxWidth().padding(vertical = 4.dp)) {
        Column(Modifier.padding(12.dp)) {
            Text(download.file, style = MaterialTheme.typography.titleSmall)
            Text(
                download.repo,
                style = MaterialTheme.typography.bodySmall,
                color = MaterialTheme.colorScheme.onSurfaceVariant,
            )
            Spacer(Modifier.height(8.dp))
            // A determinate bar only when the size is known. The server sends no total when it was
            // never told one, and a bar drawn from a guess is worse than one that admits it is
            // counting.
            if (download.total != null && download.total > 0) {
                LinearProgressIndicator(
                    progress = { download.done.toFloat() / download.total.toFloat() },
                    // Blue at every value, unlike a quota bar: a download nearing its end is
                    // nearing success, and colouring it like a limit being approached would say the
                    // opposite.
                    color = progressColor,
                    modifier = Modifier.fillMaxWidth(),
                )
                Text(
                    "${gigabytes(download.done)} of ${gigabytes(download.total)}",
                    style = MaterialTheme.typography.bodySmall,
                )
            } else {
                LinearProgressIndicator(color = progressColor, modifier = Modifier.fillMaxWidth())
                Text(
                    "${gigabytes(download.done)} so far, total size unknown",
                    style = MaterialTheme.typography.bodySmall,
                )
            }
            download.error?.let {
                Text(
                    it,
                    style = MaterialTheme.typography.bodySmall,
                    color = MaterialTheme.colorScheme.error,
                )
            }
            Row {
                Text(
                    download.state,
                    style = MaterialTheme.typography.bodySmall,
                    modifier = Modifier.weight(1f),
                )
                if (download.state == "running") {
                    TextButton(onClick = onCancel) { Text("Cancel") }
                }
            }
        }
    }
}

@Composable
private fun LocalModelCard(model: LocalModel, onDelete: () -> Unit) {
    Card(Modifier.fillMaxWidth().padding(vertical = 4.dp)) {
        Row(Modifier.padding(12.dp), verticalAlignment = Alignment.CenterVertically) {
            Column(Modifier.weight(1f)) {
                Text(model.file, style = MaterialTheme.typography.titleSmall)
                Text(
                    "${model.repo} · ${gigabytes(model.bytes)}",
                    style = MaterialTheme.typography.bodySmall,
                    color = MaterialTheme.colorScheme.onSurfaceVariant,
                )
            }
            TextButton(onClick = onDelete) { Text("Delete") }
        }
    }
}

@Composable
private fun RepoRow(repo: RemoteRepo, expanded: Boolean, onToggle: () -> Unit) {
    Card(Modifier.fillMaxWidth().padding(vertical = 4.dp)) {
        Row(Modifier.padding(12.dp), verticalAlignment = Alignment.CenterVertically) {
            Column(Modifier.weight(1f)) {
                Text(
                    repo.id,
                    style = MaterialTheme.typography.titleSmall,
                    maxLines = 1,
                    // The owner is the part that repeats; the model name at the end is what tells
                    // two entries apart.
                    overflow = TextOverflow.StartEllipsis,
                )
                Text(
                    "${repo.downloads} downloads · ${repo.likes} likes",
                    style = MaterialTheme.typography.bodySmall,
                    color = MaterialTheme.colorScheme.onSurfaceVariant,
                )
            }
            TextButton(onClick = onToggle) { Text(if (expanded) "Hide" else "Files") }
        }
    }
}

@Composable
private fun RepoFileRow(file: RemoteFile, downloading: Boolean, onDownload: () -> Unit) {
    Row(
        Modifier.fillMaxWidth().padding(start = 16.dp, top = 4.dp, bottom = 4.dp),
        verticalAlignment = Alignment.CenterVertically,
    ) {
        Column(Modifier.weight(1f)) {
            Text(file.path, style = MaterialTheme.typography.bodyMedium)
            Text(
                gigabytes(file.bytes),
                style = MaterialTheme.typography.bodySmall,
                color = MaterialTheme.colorScheme.onSurfaceVariant,
            )
        }
        // Disabled rather than absent, so the row reads the same whether this one is absent,
        // already here, or on its way. Offering "Download" for a file that is downloading would be
        // a button that does nothing anyone can see.
        TextButton(enabled = !file.have && !downloading, onClick = onDownload) {
            Text(
                when {
                    file.have -> "Downloaded"
                    downloading -> "Downloading"
                    else -> "Download"
                }
            )
        }
    }
}

private fun gigabytes(bytes: Long): String =
    if (bytes >= 1_000_000_000) {
        "%.2f GB".format(bytes / 1_000_000_000.0)
    } else {
        "%.0f MB".format(bytes / 1_000_000.0)
    }