package com.correx.infrastructure.artifactscas.index

import com.correx.infrastructure.artifactscas.segment.Location
import com.correx.infrastructure.artifactscas.segment.SegmentLayout
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import java.nio.file.Path
import java.sql.Connection
import java.sql.DriverManager

class SqliteArtifactIndex(dbPath: Path) : ArtifactIndex {
    private val connection: Connection = DriverManager.getConnection("jdbc:sqlite:${dbPath.toAbsolutePath()}")

    init {
        connection.createStatement().use { stmt ->
            stmt.execute("PRAGMA journal_mode=WAL;")
            stmt.execute("PRAGMA synchronous=NORMAL;")
            stmt.execute(
                """
                CREATE TABLE IF NOT EXISTS artifacts (
                    hash BLOB PRIMARY KEY,
                    segment_id INTEGER NOT NULL,
                    offset INTEGER NOT NULL,
                    length INTEGER NOT NULL
                );
                """.trimIndent()
            )
        }
    }

    override suspend fun lookup(hash: ByteArray): Location? = withContext(Dispatchers.IO) {
        connection.prepareStatement(
            "SELECT segment_id, offset, length FROM artifacts WHERE hash = ?"
        ).use { ps ->
            ps.setBytes(PARAM_HASH, hash)
            ps.executeQuery().use { rs ->
                if (!rs.next()) {
                    null
                } else {
                    Location(rs.getLong(COL_SEGMENT_ID), rs.getLong(COL_OFFSET), rs.getInt(COL_LENGTH))
                }
            }
        }
    }

    override suspend fun insert(hash: ByteArray, loc: Location): Unit = withContext(Dispatchers.IO) {
        // Same hash always maps to same location; duplicate inserts are no-ops.
        connection.prepareStatement(
            "INSERT OR IGNORE INTO artifacts (hash, segment_id, offset, length) VALUES (?, ?, ?, ?)"
        ).use { ps ->
            ps.setBytes(PARAM_HASH, hash)
            ps.setLong(PARAM_SEGMENT_ID, loc.segmentId)
            ps.setLong(PARAM_OFFSET, loc.offset)
            ps.setInt(PARAM_LENGTH, loc.length)
            ps.executeUpdate()
        }
    }

    override suspend fun update(hash: ByteArray, loc: Location): Unit = withContext(Dispatchers.IO) {
        connection.prepareStatement(
            """
            INSERT INTO artifacts (hash, segment_id, offset, length) VALUES (?, ?, ?, ?)
            ON CONFLICT(hash) DO UPDATE SET
                segment_id = excluded.segment_id,
                offset = excluded.offset,
                length = excluded.length
            """.trimIndent()
        ).use { ps ->
            ps.setBytes(PARAM_HASH, hash)
            ps.setLong(PARAM_SEGMENT_ID, loc.segmentId)
            ps.setLong(PARAM_OFFSET, loc.offset)
            ps.setInt(PARAM_LENGTH, loc.length)
            ps.executeUpdate()
        }
    }

    override suspend fun entriesIn(segmentId: Long): List<Pair<ByteArray, Location>> = withContext(Dispatchers.IO) {
        connection.prepareStatement(
            "SELECT hash, offset, length FROM artifacts WHERE segment_id = ? ORDER BY offset ASC"
        ).use { ps ->
            ps.setLong(1, segmentId)
            ps.executeQuery().use { rs ->
                buildList {
                    while (rs.next()) {
                        val hash = rs.getBytes(COL_ENTRY_HASH)
                        val loc = Location(segmentId, rs.getLong(COL_ENTRY_OFFSET), rs.getInt(COL_ENTRY_LENGTH))
                        add(hash to loc)
                    }
                }
            }
        }
    }

    override suspend fun deleteEntriesIn(segmentId: Long): Unit = withContext(Dispatchers.IO) {
        connection.prepareStatement("DELETE FROM artifacts WHERE segment_id = ?").use { ps ->
            ps.setLong(1, segmentId)
            ps.executeUpdate()
        }
    }

    override suspend fun maxKnownTailFor(segmentId: Long): Long = withContext(Dispatchers.IO) {
        connection.prepareStatement(
            "SELECT offset, length FROM artifacts WHERE segment_id = ? ORDER BY offset DESC LIMIT 1"
        ).use { ps ->
            ps.setLong(1, segmentId)
            ps.executeQuery().use { rs ->
                if (rs.next()) {
                    rs.getLong(1) + SegmentLayout.HEADER_SIZE + rs.getInt(2)
                } else {
                    0L
                }
            }
        }
    }

    override suspend fun flush(): Unit = withContext(Dispatchers.IO) {
        connection.createStatement().use { it.execute("PRAGMA wal_checkpoint(TRUNCATE);") }
    }

    override fun close() {
        connection.close()
    }

    private companion object {
        const val PARAM_HASH = 1
        const val PARAM_SEGMENT_ID = 2
        const val PARAM_OFFSET = 3
        const val PARAM_LENGTH = 4
        const val COL_SEGMENT_ID = 1
        const val COL_OFFSET = 2
        const val COL_LENGTH = 3
        const val COL_ENTRY_HASH = 1
        const val COL_ENTRY_OFFSET = 2
        const val COL_ENTRY_LENGTH = 3
    }
}
