AeroToss/app/src/androidMain/kotlin/com/aerotoss/transfer/AndroidTransferManager.kt

package com.aerotoss.transfer

import android.content.Context
import android.os.Environment
import com.aerotoss.model.TransferProgress
import com.aerotoss.model.TransferRequest
import com.aerotoss.model.TransferState
import com.aerotoss.util.FileUtils
import com.aerotoss.util.getDeviceName
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.*
import kotlinx.serialization.json.Json
import java.io.*
import java.net.InetAddress
import java.net.ServerSocket
import java.net.Socket
import java.security.MessageDigest
import java.util.UUID
import java.util.concurrent.ConcurrentHashMap
import java.util.concurrent.Semaphore
import java.util.concurrent.atomic.AtomicBoolean

@OptIn(ExperimentalCoroutinesApi::class)

class AndroidFileTransferManager(private val context: Context) : TransferManager {
    private val _incomingTransfers = MutableStateFlow<List<TransferProgress>>(emptyList())
    override val incomingTransfers: Flow<TransferProgress> = _incomingTransfers.mapLatest { list ->
        list.lastOrNull()
    }.filterNotNull()

    private val _outgoingTransfers = MutableStateFlow<List<TransferProgress>>(emptyList())
    override val outgoingTransfers: Flow<TransferProgress> = _outgoingTransfers.mapLatest { list ->
        list.lastOrNull()
    }.filterNotNull()

    private var serverSocket: ServerSocket? = null
    private var serverThread: Thread? = null
    private val scope = CoroutineScope(Dispatchers.IO + SupervisorJob())
    private val activeJobs = ConcurrentHashMap<String, Job>()
    private val activeSockets = ConcurrentHashMap<String, Socket>()
    private val running = AtomicBoolean(false)
    private val connectionLimiter = Semaphore(10)

    private val downloadsDir: File = run {
        val dir = context.getExternalFilesDir(Environment.DIRECTORY_DOWNLOADS)
            ?: context.filesDir
        File(dir, "AeroToss").apply { mkdirs() }
    }

    override fun getServerPort(): Int = serverSocket?.localPort ?: 0

    fun startServer(port: Int = 0): Int {
        if (!running.compareAndSet(false, true)) return serverSocket?.localPort ?: 0

        try {
            val socket = ServerSocket(port, 50, InetAddress.getByName("0.0.0.0"))
            serverSocket = socket

            serverThread = Thread {
                while (running.get() && !socket.isClosed) {
                    try {
                        val clientSocket = socket.accept()
                        scope.launch { handleIncomingConnection(clientSocket) }
                    } catch (_: Exception) {
                        if (running.get()) break
                    }
                }
            }.apply {
                isDaemon = true
                name = "aerotoss-server"
                start()
            }

            return socket.localPort
        } catch (e: Exception) {
            e.printStackTrace()
            return 0
        }
    }

    private suspend fun handleIncomingConnection(socket: Socket) {
        if (!connectionLimiter.tryAcquire()) {
            try { socket.close() } catch (_: Exception) {}
            return
        }
        withContext(Dispatchers.IO) {
            val progressId = UUID.randomUUID().toString()
            var actualFile: File? = null
            var requestFileName: String = ""
            try {
                socket.use { sock ->
                    sock.soTimeout = 60_000
                    val input = DataInputStream(sock.getInputStream())
                    val output = DataOutputStream(sock.getOutputStream())

                    val requestJson = input.readUTF()
                    val request = try {
                        Json.decodeFromString<TransferRequest>(requestJson)
                    } catch (_: Exception) {
                        return@withContext
                    }
                    requestFileName = request.fileName

                    if (request.fileSize <= 0 || request.fileSize > FileUtils.MAX_FILE_SIZE) {
                        output.writeBoolean(false)
                        return@withContext
                    }

                    val progress = TransferProgress(
                        id = progressId,
                        request = request,
                        bytesTransferred = 0,
                        totalBytes = request.fileSize,
                        state = TransferState.PENDING
                    )
                    updateIncoming(progress)

                    output.writeUTF(Json.encodeToString(TransferRequest.serializer(), request))

                    val accepted = input.readBoolean()
                    if (!accepted) {
                        updateIncoming(progress.copy(state = TransferState.CANCELLED))
                        return@withContext
                    }

                    val file = FileUtils.resolveUniqueFile(downloadsDir, request.fileName)
                    if (!FileUtils.validateFilePath(file, downloadsDir)) {
                        output.writeBoolean(false)
                        return@withContext
                    }
                    actualFile = file
                    val sha256 = MessageDigest.getInstance("SHA-256")
                    var bytesWritten = 0L

                    updateIncoming(progress.copy(state = TransferState.TRANSFERRING))

                    file.outputStream().use { fos ->
                        val buffer = ByteArray(65536)
                        var remaining = request.fileSize
                        while (remaining > 0) {
                            val toRead = minOf(buffer.size.toLong(), remaining).toInt()
                            val read = input.read(buffer, 0, toRead)
                            if (read == -1) break
                            fos.write(buffer, 0, read)
                            sha256.update(buffer, 0, read)
                            bytesWritten += read
                            remaining -= read
                            updateIncomingById(progressId, TransferProgress(
                                id = progressId,
                                request = request,
                                bytesTransferred = bytesWritten,
                                totalBytes = request.fileSize,
                                state = TransferState.TRANSFERRING
                            ))
                        }
                    }

                    val checksum = "sha256:${sha256.digest().joinToString("") { "%02x".format(it) }}"
                    output.writeUTF(checksum)
                    output.writeLong(bytesWritten)

                    if (bytesWritten == request.fileSize) {
                        updateIncomingById(progressId, TransferProgress(
                            id = progressId,
                            request = request,
                            bytesTransferred = bytesWritten,
                            totalBytes = request.fileSize,
                            state = TransferState.COMPLETED
                        ))
                    } else {
                        FileUtils.deleteIfExists(file)
                        updateIncomingById(progressId, TransferProgress(
                            id = progressId,
                            request = request,
                            bytesTransferred = bytesWritten,
                            totalBytes = request.fileSize,
                            state = TransferState.FAILED,
                            error = "Incomplete transfer: expected ${request.fileSize}, got $bytesWritten"
                        ))
                    }
                }
            } catch (_: Exception) {
                actualFile?.let { FileUtils.deleteIfExists(it) }
                val current = _incomingTransfers.value.find { it.id == progressId }
                if (current != null && current.state != TransferState.COMPLETED &&
                    current.state != TransferState.FAILED && current.state != TransferState.CANCELLED
                ) {
                    updateIncomingById(progressId, current.copy(
                        state = TransferState.FAILED,
                        error = "Transfer failed"
                    ))
                }
            } finally {
                connectionLimiter.release()
            }
        }
    }

    override suspend fun sendFile(file: File, targetHost: String, targetPort: Int): Flow<TransferProgress> {
        require(file.exists() && file.isFile) { "File does not exist or is not a regular file" }
        require(file.length() <= FileUtils.MAX_FILE_SIZE) { "File exceeds maximum size" }

        val requestId = UUID.randomUUID().toString()
        val request = TransferRequest(
            fileName = file.name,
            fileSize = file.length(),
            mimeType = "application/octet-stream",
            senderId = "",
            senderName = getDeviceName()
        )

        val initialProgress = TransferProgress(
            id = requestId,
            request = request,
            bytesTransferred = 0,
            totalBytes = file.length(),
            state = TransferState.PENDING
        )
        updateOutgoing(initialProgress)

        val job = scope.launch {
            try {
                val socket = Socket(targetHost, targetPort)
                activeSockets[requestId] = socket
                socket.use { sock ->
                    sock.soTimeout = 30_000
                    val input = DataInputStream(sock.getInputStream())
                    val output = DataOutputStream(sock.getOutputStream())

                    output.writeUTF(Json.encodeToString(TransferRequest.serializer(), request))

                    val serverResponseJson = input.readUTF()
                    val serverRequest = try {
                        Json.decodeFromString<TransferRequest>(serverResponseJson)
                    } catch (_: Exception) {
                        request
                    }

                    output.writeBoolean(true)

                    updateOutgoingById(requestId, initialProgress.copy(state = TransferState.TRANSFERRING))

                    val sha256 = MessageDigest.getInstance("SHA-256")
                    var bytesSent = 0L
                    val buffer = ByteArray(65536)
                    file.inputStream().use { fis ->
                        while (true) {
                            val read = fis.read(buffer)
                            if (read == -1) break
                            output.write(buffer, 0, read)
                            sha256.update(buffer, 0, read)
                            bytesSent += read
                            updateOutgoingById(requestId, TransferProgress(
                                id = requestId,
                                request = request,
                                bytesTransferred = bytesSent,
                                totalBytes = file.length(),
                                state = TransferState.TRANSFERRING
                            ))
                        }
                    }
                    output.flush()

                    val serverChecksum = input.readUTF()
                    val bytesReceived = input.readLong()

                    val localChecksum = "sha256:${sha256.digest().joinToString("") { "%02x".format(it) }}"

                    if (bytesReceived == file.length() && serverChecksum == localChecksum) {
                        updateOutgoingById(requestId, TransferProgress(
                            id = requestId,
                            request = request,
                            bytesTransferred = file.length(),
                            totalBytes = file.length(),
                            state = TransferState.COMPLETED
                        ))
                    } else {
                        updateOutgoingById(requestId, TransferProgress(
                            id = requestId,
                            request = request,
                            bytesTransferred = bytesSent,
                            totalBytes = file.length(),
                            state = TransferState.FAILED,
                            error = "Checksum mismatch or incomplete transfer"
                        ))
                    }
                }
            } catch (_: Exception) {
                updateOutgoingById(requestId, initialProgress.copy(
                    state = TransferState.FAILED,
                    error = "Transfer failed"
                ))
            } finally {
                activeJobs.remove(requestId)
            }
        }

        activeJobs[requestId] = job
        return outgoingTransfers.filter { it.id == requestId }
    }

    override fun cancelTransfer(requestId: String) {
        activeJobs[requestId]?.cancel()
        activeJobs.remove(requestId)
        activeSockets[requestId]?.let { try { it.close() } catch (_: Exception) {} }
        activeSockets.remove(requestId)
        updateOutgoingById(requestId, _outgoingTransfers.value.find { it.id == requestId }?.copy(
            state = TransferState.CANCELLED
        ) ?: return)
    }

    override fun dispose() {
        running.set(false)
        activeJobs.values.forEach { it.cancel() }
        activeJobs.clear()
        scope.cancel()
        try { serverSocket?.close() } catch (_: Exception) {}
        serverThread?.interrupt()
    }

    private fun updateIncoming(progress: TransferProgress) {
        _incomingTransfers.update { list ->
            list.filter { it.id != progress.id }.plus(progress).takeLast(50)
        }
    }

    private fun updateIncomingById(id: String, progress: TransferProgress) {
        _incomingTransfers.update { list ->
            list.filter { it.id != id }.plus(progress).takeLast(50)
        }
    }

    private fun updateOutgoing(progress: TransferProgress) {
        _outgoingTransfers.update { list ->
            list.filter { it.id != progress.id }.plus(progress).takeLast(50)
        }
    }

    private fun updateOutgoingById(id: String, progress: TransferProgress) {
        _outgoingTransfers.update { list ->
            list.filter { it.id != id }.plus(progress)
        }
    }
}