import java.io.ByteArrayInputStream
import java.io.File
import java.io.RandomAccessFile
import java.nio.ByteBuffer
import java.nio.ByteOrder
import java.security.MessageDigest
import java.util.zip.GZIPInputStream
import javax.crypto.Cipher
import javax.crypto.Mac
import javax.crypto.spec.IvParameterSpec
import javax.crypto.spec.SecretKeySpec
// Bouncy Castle — required only for --password (Argon2id KDF).
// Download: https://www.bouncycastle.org/download/bouncy-castle-java/#latest
// Run: java -cp bcprov-jdk18on-1.79.jar:ArchiveDecoder.jar ArchiveDecoderKt ...
import org.bouncycastle.crypto.generators.Argon2BytesGenerator
import org.bouncycastle.crypto.params.Argon2Parameters

// ---------------------------------------------------------------------------
// Constants (matching FORMAT.md Section 4 and src/key.rs)
// ---------------------------------------------------------------------------

/** Custom magic bytes: 0x00 0xEA 0x72 0x63 (FORMAT.md Section 4). */
val MAGIC = byteArrayOf(0x00, 0xEA.toByte(), 0x72, 0x63)

/** Fixed header size in bytes (FORMAT.md Section 4). */
const val HEADER_SIZE = 40

/**
 * Fixed 8-byte XOR obfuscation key (FORMAT.md Section 9.1).
 * Applied cyclically across the 40-byte header for obfuscation/de-obfuscation.
 */
val XOR_KEY = byteArrayOf(
    0xA5.toByte(), 0x3C, 0x96.toByte(), 0x0F,
    0xE1.toByte(), 0x7B, 0x4D, 0xC8.toByte()
)

// ---------------------------------------------------------------------------
// Data classes
// ---------------------------------------------------------------------------

/** Archive header (40 bytes fixed at offset 0x00). FORMAT.md Section 4. */
data class ArchiveHeader(
    val version: Int,
    val flags: Int,
    val fileCount: Int,
    val tocOffset: Long,
    val tocSize: Long,
    val tocIv: ByteArray,
)

/** Entry table entry (variable length: 104 + name_length bytes). FORMAT.md Section 5 (v1.1). */
data class TocEntry(
    val name: String,
    val entryType: Int,      // 0x00=file, 0x01=directory
    val permissions: Int,    // Lower 12 bits of POSIX mode_t
    val originalSize: Long,
    val compressedSize: Long,
    val encryptedSize: Int,
    val dataOffset: Long,
    val iv: ByteArray,
    val hmac: ByteArray,
    val sha256: ByteArray,
    val compressionFlag: Int,
    val paddingAfter: Int,
)

// ---------------------------------------------------------------------------
// Little-endian integer helpers (using ByteBuffer)
// ---------------------------------------------------------------------------

/** Read an unsigned 16-bit little-endian integer from [data] at [offset]. */
fun readLeU16(data: ByteArray, offset: Int): Int {
    return ByteBuffer.wrap(data, offset, 2)
        .order(ByteOrder.LITTLE_ENDIAN)
        .short.toInt() and 0xFFFF
}

/** Read an unsigned 32-bit little-endian integer from [data] at [offset]. */
fun readLeU32(data: ByteArray, offset: Int): Long {
    return ByteBuffer.wrap(data, offset, 4)
        .order(ByteOrder.LITTLE_ENDIAN)
        .int.toLong() and 0xFFFFFFFFL
}

// ---------------------------------------------------------------------------
// Header parsing (FORMAT.md Section 4)
// ---------------------------------------------------------------------------

/**
 * Parse the 40-byte archive header.
 *
 * Verifies: magic bytes, version == 2 (v1.1 format), reserved flag bits 5-7 are zero.
 */
fun parseHeader(data: ByteArray): ArchiveHeader {
    require(data.size >= HEADER_SIZE) { "Header too short: ${data.size} bytes" }

    // Verify magic bytes
    require(
        data[0] == MAGIC[0] && data[1] == MAGIC[1] &&
        data[2] == MAGIC[2] && data[3] == MAGIC[3]
    ) { "Invalid magic bytes" }

    // Version check
    val version = data[4].toInt() and 0xFF
    require(version == 2) { "Unsupported version: $version (expected v1.1 format, version=2)" }

    // Flags validation
    val flags = data[5].toInt() and 0xFF
    require(flags and 0xE0 == 0) { "Unknown flags set: 0x${flags.toString(16)} (bits 5-7 must be zero)" }

    // Read remaining fields
    val fileCount = readLeU16(data, 6)
    val tocOffset = readLeU32(data, 8)
    val tocSize = readLeU32(data, 12)
    val tocIv = data.copyOfRange(16, 32)

    return ArchiveHeader(version, flags, fileCount, tocOffset, tocSize, tocIv)
}

// ---------------------------------------------------------------------------
// TOC parsing (FORMAT.md Section 5)
// ---------------------------------------------------------------------------

/**
 * Parse a single TOC entry from [data] starting at [offset].
 *
 * Returns a Pair of the parsed entry and the new offset after the entry.
 * Entry size formula: 104 + name_length bytes (v1.1).
 */
fun parseTocEntry(data: ByteArray, offset: Int): Pair<TocEntry, Int> {
    var pos = offset

    // name_length (u16 LE)
    val nameLength = readLeU16(data, pos)
    pos += 2

    // name (UTF-8 bytes)
    val name = String(data, pos, nameLength, Charsets.UTF_8)
    pos += nameLength

    // entry_type (u8): 0x00=file, 0x01=directory (v1.1)
    val entryType = data[pos].toInt() and 0xFF; pos += 1

    // permissions (u16 LE): lower 12 bits of POSIX mode_t (v1.1)
    val permissions = readLeU16(data, pos); pos += 2

    // Fixed fields: original_size, compressed_size, encrypted_size, data_offset (all u32 LE)
    val originalSize = readLeU32(data, pos); pos += 4
    val compressedSize = readLeU32(data, pos); pos += 4
    val encryptedSize = readLeU32(data, pos).toInt(); pos += 4
    val dataOffset = readLeU32(data, pos); pos += 4

    // iv (16 bytes)
    val iv = data.copyOfRange(pos, pos + 16); pos += 16

    // hmac (32 bytes)
    val hmac = data.copyOfRange(pos, pos + 32); pos += 32

    // sha256 (32 bytes)
    val sha256 = data.copyOfRange(pos, pos + 32); pos += 32

    // compression_flag (u8)
    val compressionFlag = data[pos].toInt() and 0xFF; pos += 1

    // padding_after (u16 LE)
    val paddingAfter = readLeU16(data, pos); pos += 2

    val entry = TocEntry(
        name, entryType, permissions, originalSize, compressedSize, encryptedSize,
        dataOffset, iv, hmac, sha256, compressionFlag, paddingAfter
    )
    return Pair(entry, pos)
}

/**
 * Parse all TOC entries sequentially.
 *
 * Asserts that after parsing all [fileCount] entries, the cursor equals [data].size.
 */
fun parseToc(data: ByteArray, fileCount: Int): List<TocEntry> {
    val entries = mutableListOf<TocEntry>()
    var pos = 0
    for (i in 0 until fileCount) {
        val (entry, newPos) = parseTocEntry(data, pos)
        entries.add(entry)
        pos = newPos
    }
    require(pos == data.size) {
        "TOC parsing error: consumed $pos bytes but TOC size is ${data.size}"
    }
    return entries
}

// ---------------------------------------------------------------------------
// Crypto utility functions (FORMAT.md Section 7, Section 13.6)
// ---------------------------------------------------------------------------

/**
 * Verify HMAC-SHA-256 over IV || ciphertext.
 *
 * @param iv The 16-byte IV from the TOC entry.
 * @param ciphertext The encrypted data (encrypted_size bytes).
 * @param key The 32-byte key (same as AES key in v1).
 * @param expectedHmac The 32-byte HMAC from the TOC entry.
 * @return true if HMAC matches.
 */
fun verifyHmac(iv: ByteArray, ciphertext: ByteArray, key: ByteArray, expectedHmac: ByteArray): Boolean {
    val mac = Mac.getInstance("HmacSHA256")
    mac.init(SecretKeySpec(key, "HmacSHA256"))
    mac.update(iv)
    mac.update(ciphertext)
    val computed = mac.doFinal()
    return computed.contentEquals(expectedHmac)
}

/**
 * Decrypt AES-256-CBC ciphertext.
 *
 * Uses PKCS5Padding which is functionally identical to PKCS7 for 16-byte AES blocks.
 * cipher.doFinal() automatically removes PKCS7 padding.
 *
 * @param ciphertext The encrypted data.
 * @param iv The 16-byte IV.
 * @param key The 32-byte AES key.
 * @return Decrypted plaintext (padding already removed).
 */
fun decryptAesCbc(ciphertext: ByteArray, iv: ByteArray, key: ByteArray): ByteArray {
    val cipher = Cipher.getInstance("AES/CBC/PKCS5Padding")
    cipher.init(Cipher.DECRYPT_MODE, SecretKeySpec(key, "AES"), IvParameterSpec(iv))
    return cipher.doFinal(ciphertext)
}

/**
 * Decompress gzip data.
 *
 * @param compressed Gzip-compressed data.
 * @return Decompressed data.
 */
fun decompressGzip(compressed: ByteArray): ByteArray {
    return GZIPInputStream(ByteArrayInputStream(compressed)).readBytes()
}

/**
 * Verify SHA-256 checksum of data.
 *
 * @param data The decompressed file content.
 * @param expectedSha256 The 32-byte SHA-256 from the TOC entry.
 * @return true if checksum matches.
 */
fun verifySha256(data: ByteArray, expectedSha256: ByteArray): Boolean {
    val digest = MessageDigest.getInstance("SHA-256")
    val computed = digest.digest(data)
    return computed.contentEquals(expectedSha256)
}

// ---------------------------------------------------------------------------
// XOR header de-obfuscation (FORMAT.md Section 9.1)
// ---------------------------------------------------------------------------

/**
 * XOR-obfuscate or de-obfuscate a header buffer in-place.
 *
 * XOR is its own inverse, so the same function encodes and decodes.
 * Applies the 8-byte XOR_KEY cyclically across the first 40 bytes.
 * Uses `and 0xFF` on BOTH operands to avoid Kotlin signed byte issues.
 */
fun xorHeader(buf: ByteArray) {
    for (i in 0 until minOf(buf.size, 40)) {
        buf[i] = ((buf[i].toInt() and 0xFF) xor (XOR_KEY[i % 8].toInt() and 0xFF)).toByte()
    }
}

// ---------------------------------------------------------------------------
// Key source types and resolution
// ---------------------------------------------------------------------------

/** How the user supplies the decryption key. */
sealed class KeySource {
    data class Hex(val hex: String) : KeySource()
    data class KeyFile(val path: String) : KeySource()
    data class Password(val password: String) : KeySource()
}

/** Size of the KDF salt appended after the 40-byte header (FORMAT.md Section 4). */
const val SALT_SIZE = 16

/**
 * Read the 16-byte KDF salt from offset 40 if the KDF flag (bit 4) is set.
 * Returns null when the archive uses a raw key (no salt present).
 */
fun readSalt(raf: RandomAccessFile, header: ArchiveHeader): ByteArray? {
    if (header.flags and 0x10 == 0) return null
    raf.seek(HEADER_SIZE.toLong())
    val salt = ByteArray(SALT_SIZE)
    raf.readFully(salt)
    return salt
}

/**
 * Derive a 32-byte key from a password and salt using Argon2id.
 *
 * Parameters match the Rust implementation (src/kdf.rs) exactly:
 *   - Argon2id v19
 *   - memory = 19456 KiB (19 MiB)
 *   - iterations = 2
 *   - parallelism = 1
 *   - output length = 32 bytes
 *
 * Requires Bouncy Castle on the classpath.
 */
fun deriveKeyFromPassword(password: String, salt: ByteArray): ByteArray {
    val params = Argon2Parameters.Builder(Argon2Parameters.ARGON2_id)
        .withVersion(Argon2Parameters.ARGON2_VERSION_13)
        .withMemoryAsKB(19456)
        .withIterations(2)
        .withParallelism(1)
        .withSalt(salt)
        .build()

    val generator = Argon2BytesGenerator()
    generator.init(params)

    val key = ByteArray(32)
    generator.generateBytes(password.toByteArray(Charsets.UTF_8), key)
    return key
}

/**
 * Parse a hex string into a ByteArray.
 * Accepts lowercase, uppercase, or mixed hex. Must be exactly 64 hex chars (32 bytes).
 */
fun hexToBytes(hex: String): ByteArray {
    require(hex.length == 64) { "Hex key must be exactly 64 hex characters (32 bytes), got ${hex.length}" }
    return ByteArray(32) { i ->
        hex.substring(i * 2, i * 2 + 2).toInt(16).toByte()
    }
}

/**
 * Resolve a [KeySource] into a 32-byte key.
 *
 * @param source How the key was supplied (hex, file, or password).
 * @param salt Optional 16-byte salt read from the archive (required for Password source).
 * @return 32-byte key suitable for AES-256 and HMAC-SHA-256.
 */
fun resolveKey(source: KeySource, salt: ByteArray?): ByteArray {
    return when (source) {
        is KeySource.Hex -> hexToBytes(source.hex)
        is KeySource.KeyFile -> {
            val bytes = File(source.path).readBytes()
            require(bytes.size == 32) { "Key file must be exactly 32 bytes, got ${bytes.size}" }
            bytes
        }
        is KeySource.Password -> {
            requireNotNull(salt) {
                "Archive does not contain a KDF salt (flag bit 4 not set). " +
                "This archive was not created with --password. Use --key or --key-file instead."
            }
            deriveKeyFromPassword(source.password, salt)
        }
    }
}

// ---------------------------------------------------------------------------
// Permissions restoration (v1.1)
// ---------------------------------------------------------------------------

/**
 * Apply POSIX permissions to a file or directory using Java File API.
 *
 * Java's File API can only distinguish "owner" vs "everyone" permissions
 * (not owner/group/others separately). This is acceptable per KOT-07.
 *
 * @param file The file or directory to apply permissions to.
 * @param permissions Lower 12 bits of POSIX mode_t (e.g., 0o755 = 0x01ED).
 */
fun applyPermissions(file: File, permissions: Int) {
    val ownerRead = (permissions shr 8) and 1 != 0     // bit 8
    val ownerWrite = (permissions shr 7) and 1 != 0    // bit 7
    val ownerExec = (permissions shr 6) and 1 != 0     // bit 6
    val othersRead = (permissions shr 2) and 1 != 0    // bit 2
    val othersWrite = (permissions shr 1) and 1 != 0   // bit 1
    val othersExec = permissions and 1 != 0             // bit 0

    // Set "everyone" permissions first (ownerOnly=false), then override owner-only
    file.setReadable(othersRead, false)
    file.setWritable(othersWrite, false)
    file.setExecutable(othersExec, false)
    // Owner-only overrides (ownerOnly=true)
    file.setReadable(ownerRead, true)
    file.setWritable(ownerWrite, true)
    file.setExecutable(ownerExec, true)
}

// ---------------------------------------------------------------------------
// Main decode orchestration (FORMAT.md Section 10)
// ---------------------------------------------------------------------------

/**
 * Decode an encrypted archive, extracting all files to [outputDir].
 *
 * Follows FORMAT.md Section 10 decode order:
 *   1. Read and parse 40-byte header
 *   2. Read KDF salt if present (flag bit 4)
 *   3. Resolve key from [keySource] (hex, file, or password+salt)
 *   4. Seek to tocOffset, read and parse TOC entries
 *   5. For each file: verify HMAC, decrypt, decompress, verify SHA-256, write
 */
fun decode(archivePath: String, outputDir: String, keySource: KeySource) {
    val raf = RandomAccessFile(archivePath, "r")

    // Read 40-byte header
    val headerBytes = ByteArray(HEADER_SIZE)
    raf.readFully(headerBytes)

    // XOR bootstrapping (FORMAT.md Section 10, step 2):
    // Check if first 4 bytes match MAGIC; if not, attempt XOR de-obfuscation
    if (!(headerBytes[0] == MAGIC[0] && headerBytes[1] == MAGIC[1] &&
          headerBytes[2] == MAGIC[2] && headerBytes[3] == MAGIC[3])) {
        xorHeader(headerBytes)
    }

    val header = parseHeader(headerBytes)

    // Read KDF salt if present (flag bit 4)
    val salt = readSalt(raf, header)

    // Resolve the key from the supplied source
    val key = resolveKey(keySource, salt)

    // Read TOC bytes -- decrypt if TOC encryption flag is set (bit 1)
    val entries: List<TocEntry>
    if (header.flags and 0x02 != 0) {
        // TOC is encrypted: read encrypted bytes, decrypt, then parse
        raf.seek(header.tocOffset)
        val encryptedToc = ByteArray(header.tocSize.toInt())
        raf.readFully(encryptedToc)
        val decryptedToc = decryptAesCbc(encryptedToc, header.tocIv, key)
        entries = parseToc(decryptedToc, header.fileCount)
    } else {
        // TOC is plaintext (backward compatibility)
        raf.seek(header.tocOffset)
        val tocBytes = ByteArray(header.tocSize.toInt())
        raf.readFully(tocBytes)
        entries = parseToc(tocBytes, header.fileCount)
    }

    var successCount = 0

    for (entry in entries) {
        if (entry.entryType == 1) {
            // Directory entry: create the directory, apply permissions, no decryption
            val dir = File(outputDir, entry.name)
            dir.mkdirs()
            applyPermissions(dir, entry.permissions)
            println("Created dir: ${entry.name}")
            successCount++
            continue
        }

        // File entry (entryType == 0): standard crypto pipeline

        // Ensure parent directories exist (for files with relative paths)
        val outFile = File(outputDir, entry.name)
        outFile.parentFile?.mkdirs()

        // Step 1: Seek to data_offset and read ciphertext
        raf.seek(entry.dataOffset)
        val ciphertext = ByteArray(entry.encryptedSize)
        raf.readFully(ciphertext)

        // Step 2: Verify HMAC FIRST (Encrypt-then-MAC -- FORMAT.md Section 7)
        if (!verifyHmac(entry.iv, ciphertext, key, entry.hmac)) {
            System.err.println("HMAC failed for ${entry.name}, skipping")
            continue
        }

        // Step 3: Decrypt (PKCS5Padding auto-removes PKCS7 padding)
        val decrypted = decryptAesCbc(ciphertext, entry.iv, key)

        // Step 4: Decompress if compression_flag == 1
        val original = if (entry.compressionFlag == 1) {
            decompressGzip(decrypted)
        } else {
            decrypted
        }

        // Step 5: Verify SHA-256 checksum
        if (!verifySha256(original, entry.sha256)) {
            System.err.println("WARNING: SHA-256 mismatch for ${entry.name}")
            // Still write the file (matching Rust behavior)
        }

        // Step 6: Write output file and apply permissions
        outFile.writeBytes(original)
        applyPermissions(outFile, entry.permissions)
        println("Extracted: ${entry.name} (${original.size} bytes)")
        successCount++
    }

    raf.close()

    println("Done: $successCount entries extracted")
}

// ---------------------------------------------------------------------------
// CLI entry point
// ---------------------------------------------------------------------------

fun main(args: Array<String>) {
    val usage = """
        |Usage: java -jar ArchiveDecoder.jar [OPTIONS] <archive> <output_dir>
        |
        |Key options (exactly one required):
        |  --key <hex>          64-char hex key (32 bytes)
        |  --key-file <path>    Path to 32-byte raw key file
        |  --password <pass>    Password (requires Bouncy Castle on classpath for Argon2id)
        |
        |For --password, run with Bouncy Castle:
        |  java -cp bcprov-jdk18on-1.79.jar:ArchiveDecoder.jar ArchiveDecoderKt --password <pass> <archive> <output_dir>
    """.trimMargin()

    // Parse arguments
    var keySource: KeySource? = null
    val positional = mutableListOf<String>()
    var i = 0
    while (i < args.size) {
        when (args[i]) {
            "--key" -> {
                require(i + 1 < args.size) { "--key requires a hex argument" }
                keySource = KeySource.Hex(args[i + 1])
                i += 2
            }
            "--key-file" -> {
                require(i + 1 < args.size) { "--key-file requires a path argument" }
                keySource = KeySource.KeyFile(args[i + 1])
                i += 2
            }
            "--password" -> {
                require(i + 1 < args.size) { "--password requires a password argument" }
                keySource = KeySource.Password(args[i + 1])
                i += 2
            }
            "--help", "-h" -> {
                println(usage)
                return
            }
            else -> {
                positional.add(args[i])
                i++
            }
        }
    }

    if (keySource == null || positional.size != 2) {
        System.err.println(usage)
        System.exit(1)
    }

    val archivePath = positional[0]
    val outputDir = positional[1]

    // Validate archive exists
    require(File(archivePath).exists()) { "Archive not found: $archivePath" }

    // Create output directory if needed
    File(outputDir).mkdirs()

    decode(archivePath, outputDir, keySource!!)
}