// Copyright (c) 2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
// Based off:   https://github.com/golang/crypto/tree/master/argon2
// Package argon2 implements the Argon2 password hashing and key derivation
// functions described in RFC 9106.
module argon2

import crypto.blake2b
import crypto.hmac
import crypto.rand
import encoding.base64
import encoding.binary
import math.bits

pub const version = u32(0x13)

pub const default_time = u32(3)
pub const default_memory = u32(64 * 1024)
pub const default_threads = u8(4)
pub const default_key_len = u32(32)
pub const default_salt_len = 16

const argon2_d = 0
const argon2_i = 1
const argon2_id = 2
const block_length = 128
const block_bytes = 1024
const sync_points = u32(4)

// Params controls Argon2 password hashing output.
pub struct Params {
pub:
    time     u32
    memory   u32
    threads  u8
    key_len  u32
    salt_len int
}

struct InvalidTimeError {
    Error
    time u32
}

fn (err InvalidTimeError) msg() string {
    return 'the number of rounds ${err.time} must be greater than 0'
}

struct InvalidThreadsError {
    Error
    threads u8
}

fn (err InvalidThreadsError) msg() string {
    return 'the parallelism degree ${err.threads} must be greater than 0'
}

struct InvalidKeyLengthError {
    Error
    length u32
}

fn (err InvalidKeyLengthError) msg() string {
    return 'the output length ${err.length} must be greater than 0'
}

struct InvalidSaltLengthError {
    Error
    length int
}

fn (err InvalidSaltLengthError) msg() string {
    return 'the salt length ${err.length} must be greater than 0'
}

struct InvalidHashError {
    Error
    reason string
}

fn (err InvalidHashError) msg() string {
    return err.reason
}

// default_params returns the recommended default Argon2id password hashing parameters.
pub fn default_params() Params {
    return Params{
        time:     default_time
        memory:   default_memory
        threads:  default_threads
        key_len:  default_key_len
        salt_len: default_salt_len
    }
}

// key derives a key from `password` and `salt` using Argon2i.
pub fn key(password []u8, salt []u8, time u32, memory u32, threads u8, key_len u32) ![]u8 {
    return derive_key(argon2_i, password, salt, []u8{}, []u8{}, time, memory, threads, key_len)
}

// d_key derives a key from `password` and `salt` using Argon2d.
pub fn d_key(password []u8, salt []u8, time u32, memory u32, threads u8, key_len u32) ![]u8 {
    return derive_key(argon2_d, password, salt, []u8{}, []u8{}, time, memory, threads, key_len)
}

// id_key derives a key from `password` and `salt` using Argon2id.
pub fn id_key(password []u8, salt []u8, time u32, memory u32, threads u8, key_len u32) ![]u8 {
    return derive_key(argon2_id, password, salt, []u8{}, []u8{}, time, memory, threads, key_len)
}

// generate_from_password hashes `password` with default Argon2id parameters
// and returns a PHC-formatted encoded string.
pub fn generate_from_password(password []u8) !string {
    return generate_from_password_with_params(password, default_params())
}

// generate_from_password_with_params hashes `password` with Argon2id using `params`
// and returns a PHC-formatted encoded string.
pub fn generate_from_password_with_params(password []u8, params Params) !string {
    normalized := normalize_params(params)
    if normalized.salt_len <= 0 {
        return InvalidSaltLengthError{
            length: normalized.salt_len
        }
    }
    salt := rand.bytes(normalized.salt_len)!
    hash := derive_key(argon2_id, password, salt, []u8{}, []u8{}, normalized.time,
        normalized.memory, normalized.threads, normalized.key_len)!
    return encode_hash(argon2_id, salt, hash, normalized.time, normalized.memory,
        normalized.threads)
}

// compare_hash_and_password verifies a PHC-formatted Argon2 hash against `password`.
pub fn compare_hash_and_password(password []u8, encoded_hash []u8) ! {
    decoded := decode_hash(encoded_hash.bytestr())!
    other_hash := derive_key(decoded.mode, password, decoded.salt, []u8{}, []u8{}, decoded.time,
        decoded.memory, decoded.threads, u32(decoded.hash.len))!
    if !hmac.equal(other_hash, decoded.hash) {
        return error('mismatched hash and password')
    }
}

struct DecodedHash {
    mode    int
    time    u32
    memory  u32
    threads u8
    salt    []u8
    hash    []u8
}

fn derive_key(mode int, password []u8, salt []u8, secret []u8, data []u8, time u32, memory u32, threads u8, key_len u32) ![]u8 {
    if time < 1 {
        return InvalidTimeError{
            time: time
        }
    }
    if threads < 1 {
        return InvalidThreadsError{
            threads: threads
        }
    }
    if key_len < 1 {
        return InvalidKeyLengthError{
            length: key_len
        }
    }
    threads_u32 := u32(threads)
    effective_memory := normalize_memory(memory, threads_u32)
    h0 :=
        init_hash(password, salt, secret, data, time, effective_memory, threads_u32, key_len, mode)!
    mut blocks := init_blocks(h0, effective_memory, threads_u32)!
    process_blocks(mut blocks, time, effective_memory, threads_u32, mode)
    return extract_key(mut blocks, effective_memory, threads_u32, key_len)
}

fn normalize_params(params Params) Params {
    threads := if params.threads == 0 { default_threads } else { params.threads }
    memory := if params.memory == 0 { default_memory } else { params.memory }
    return Params{
        time:     if params.time == 0 { default_time } else { params.time }
        memory:   normalize_memory(memory, u32(threads))
        threads:  threads
        key_len:  if params.key_len == 0 { default_key_len } else { params.key_len }
        salt_len: if params.salt_len == 0 { default_salt_len } else { params.salt_len }
    }
}

@[inline]
fn normalize_memory(memory u32, threads u32) u32 {
    mut effective_memory := memory / (sync_points * threads) * (sync_points * threads)
    if effective_memory < 2 * sync_points * threads {
        effective_memory = 2 * sync_points * threads
    }
    return effective_memory
}

fn init_hash(password []u8, salt []u8, secret []u8, data []u8, time u32, memory u32, threads u32, key_len u32, mode int) ![]u8 {
    mut params := []u8{len: 24, init: 0}
    mut tmp := []u8{len: 4, init: 0}
    mut d := blake2b.new512()!
    binary.little_endian_put_u32_at(mut params, threads, 0)
    binary.little_endian_put_u32_at(mut params, key_len, 4)
    binary.little_endian_put_u32_at(mut params, memory, 8)
    binary.little_endian_put_u32_at(mut params, time, 12)
    binary.little_endian_put_u32_at(mut params, version, 16)
    binary.little_endian_put_u32_at(mut params, u32(mode), 20)
    d.write(params)!
    binary.little_endian_put_u32(mut tmp, u32(password.len))
    d.write(tmp)!
    d.write(password)!
    binary.little_endian_put_u32(mut tmp, u32(salt.len))
    d.write(tmp)!
    d.write(salt)!
    binary.little_endian_put_u32(mut tmp, u32(secret.len))
    d.write(tmp)!
    d.write(secret)!
    binary.little_endian_put_u32(mut tmp, u32(data.len))
    d.write(tmp)!
    d.write(data)!
    mut h0 := []u8{len: blake2b.size512 + 8, init: 0}
    sum := d.checksum()
    copy(mut h0[..blake2b.size512], sum)
    return h0
}

fn init_blocks(h0 []u8, memory u32, threads u32) ![]u64 {
    mut block := []u8{len: block_bytes, init: 0}
    mut blocks := []u64{len: int(memory) * block_length, init: u64(0)}
    mut h0_block := h0.clone()
    for lane := u32(0); lane < threads; lane++ {
        start := lane * (memory / threads)
        binary.little_endian_put_u32_at(mut h0_block, lane, blake2b.size512 + 4)
        binary.little_endian_put_u32_at(mut h0_block, 0, blake2b.size512)
        block = blake2b_hash(block_bytes, h0_block)!
        set_block(mut blocks, start, block)
        binary.little_endian_put_u32_at(mut h0_block, 1, blake2b.size512)
        block = blake2b_hash(block_bytes, h0_block)!
        set_block(mut blocks, start + 1, block)
    }
    return blocks
}

fn process_blocks(mut blocks []u64, time u32, memory u32, threads u32, mode int) {
    lanes := memory / threads
    segments := lanes / sync_points
    for pass := u32(0); pass < time; pass++ {
        for slice := u32(0); slice < sync_points; slice++ {
            for lane := u32(0); lane < threads; lane++ {
                process_segment(mut blocks, pass, slice, lane, time, memory, lanes, segments,
                    threads, mode)
            }
        }
    }
}

fn process_segment(mut blocks []u64, pass u32, slice u32, lane u32, time u32, memory u32, lanes u32, segments u32, threads u32, mode int) {
    mut addresses := [block_length]u64{}
    mut input := [block_length]u64{}
    zero := [block_length]u64{}
    if mode == argon2_i || (mode == argon2_id && pass == 0 && slice < sync_points / 2) {
        input[0] = u64(pass)
        input[1] = u64(lane)
        input[2] = u64(slice)
        input[3] = u64(memory)
        input[4] = u64(time)
        input[5] = u64(mode)
    }
    mut index := u32(0)
    if pass == 0 && slice == 0 {
        index = 2
        if mode == argon2_i || mode == argon2_id {
            input[6]++
            process_block(mut addresses, input, zero)
            process_block(mut addresses, addresses, zero)
        }
    }
    mut offset := lane * lanes + slice * segments + index
    for index < segments {
        mut prev := offset - 1
        if index == 0 && slice == 0 {
            prev += lanes
        }
        mut random := u64(0)
        if mode == argon2_i || (mode == argon2_id && pass == 0 && slice < sync_points / 2) {
            if index % u32(block_length) == 0 {
                input[6]++
                process_block(mut addresses, input, zero)
                process_block(mut addresses, addresses, zero)
            }
            random = addresses[int(index % u32(block_length))]
        } else {
            random = blocks[block_offset(prev)]
        }
        new_offset := index_alpha(random, lanes, segments, threads, pass, slice, lane, index)
        process_block_xor_in_place(mut blocks, offset, prev, new_offset)
        index++
        offset++
    }
}

fn extract_key(mut blocks []u64, memory u32, threads u32, key_len u32) ![]u8 {
    lanes := memory / threads
    last_block := memory - 1
    last_base := block_offset(last_block)
    for lane := u32(0); lane < threads - 1; lane++ {
        base := block_offset(lane * lanes + lanes - 1)
        for i := 0; i < block_length; i++ {
            blocks[last_base + i] ^= blocks[base + i]
        }
    }
    mut block := []u8{len: block_bytes, init: 0}
    for i := 0; i < block_length; i++ {
        binary.little_endian_put_u64_at(mut block, blocks[last_base + i], i * 8)
    }
    return blake2b_hash(int(key_len), block)
}

@[inline]
fn block_offset(index u32) int {
    return int(index) * block_length
}

fn set_block(mut blocks []u64, index u32, data []u8) {
    base := block_offset(index)
    for i := 0; i < block_length; i++ {
        blocks[base + i] = binary.little_endian_u64_at(data, i * 8)
    }
}

fn blake2b_hash(out_len int, input []u8) ![]u8 {
    mut prefix := []u8{len: 4, init: 0}
    binary.little_endian_put_u32(mut prefix, u32(out_len))
    mut initial := prefix.clone()
    initial << input
    if out_len <= blake2b.size512 {
        mut d := blake2b.new_digest(u8(out_len), []u8{})!
        d.write(initial)!
        return d.checksum()
    }
    mut out := []u8{len: out_len, init: 0}
    mut block := blake2b.sum512(initial)
    copy(mut out[..32], block[..32])
    mut pos := 32
    for out_len - pos > blake2b.size512 {
        block = blake2b.sum512(block)
        copy(mut out[pos..pos + 32], block[..32])
        pos += 32
    }
    remaining := out_len - pos
    mut d := blake2b.new_digest(u8(remaining), []u8{})!
    d.write(block)!
    copy(mut out[pos..], d.checksum())
    return out
}

fn process_block(mut out [block_length]u64, in1 [block_length]u64, in2 [block_length]u64) {
    process_block_generic(mut out, in1, in2, false)
}

fn process_block_generic(mut out [block_length]u64, in1 [block_length]u64, in2 [block_length]u64, xor bool) {
    mut t := [block_length]u64{}
    for i := 0; i < block_length; i++ {
        t[i] = in1[i] ^ in2[i]
    }
    for i := 0; i < block_length; i += 16 {
        blamka_generic(mut t, i + 0, i + 1, i + 2, i + 3, i + 4, i + 5, i + 6, i + 7, i + 8, i + 9,

            i + 10, i + 11, i + 12, i + 13, i + 14, i + 15)
    }
    for i := 0; i < block_length / 8; i += 2 {
        blamka_generic(mut t, i, i + 1, 16 + i, 16 + i + 1, 32 + i, 32 + i + 1, 48 + i, 48 + i + 1,

            64 + i, 64 + i + 1, 80 + i, 80 + i + 1, 96 + i, 96 + i + 1, 112 + i, 112 + i + 1)
    }
    for i := 0; i < block_length; i++ {
        if xor {
            out[i] ^= in1[i] ^ in2[i] ^ t[i]
        } else {
            out[i] = in1[i] ^ in2[i] ^ t[i]
        }
    }
}

fn process_block_xor_in_place(mut blocks []u64, out_index u32, in1_index u32, in2_index u32) {
    mut t := [block_length]u64{}
    out_base := block_offset(out_index)
    in1_base := block_offset(in1_index)
    in2_base := block_offset(in2_index)
    for i := 0; i < block_length; i++ {
        t[i] = blocks[in1_base + i] ^ blocks[in2_base + i]
    }
    for i := 0; i < block_length; i += 16 {
        blamka_generic(mut t, i + 0, i + 1, i + 2, i + 3, i + 4, i + 5, i + 6, i + 7, i + 8, i + 9,

            i + 10, i + 11, i + 12, i + 13, i + 14, i + 15)
    }
    for i := 0; i < block_length / 8; i += 2 {
        blamka_generic(mut t, i, i + 1, 16 + i, 16 + i + 1, 32 + i, 32 + i + 1, 48 + i, 48 + i + 1,

            64 + i, 64 + i + 1, 80 + i, 80 + i + 1, 96 + i, 96 + i + 1, 112 + i, 112 + i + 1)
    }
    for i := 0; i < block_length; i++ {
        blocks[out_base + i] ^= blocks[in1_base + i] ^ blocks[in2_base + i] ^ t[i]
    }
}

fn blamka_generic(mut block [block_length]u64, a int, b int, c int, d int, e int, f int, g int, h int, i int, j int, k int, l int, m int, n int, o int, p int) {
    mut v00 := block[a]
    mut v01 := block[b]
    mut v02 := block[c]
    mut v03 := block[d]
    mut v04 := block[e]
    mut v05 := block[f]
    mut v06 := block[g]
    mut v07 := block[h]
    mut v08 := block[i]
    mut v09 := block[j]
    mut v10 := block[k]
    mut v11 := block[l]
    mut v12 := block[m]
    mut v13 := block[n]
    mut v14 := block[o]
    mut v15 := block[p]

    v00 += v04 + 2 * u64(u32(v00)) * u64(u32(v04))
    v12 ^= v00
    v12 = bits.rotate_left_64(v12, -32)
    v08 += v12 + 2 * u64(u32(v08)) * u64(u32(v12))
    v04 ^= v08
    v04 = bits.rotate_left_64(v04, -24)

    v00 += v04 + 2 * u64(u32(v00)) * u64(u32(v04))
    v12 ^= v00
    v12 = bits.rotate_left_64(v12, -16)
    v08 += v12 + 2 * u64(u32(v08)) * u64(u32(v12))
    v04 ^= v08
    v04 = bits.rotate_left_64(v04, -63)

    v01 += v05 + 2 * u64(u32(v01)) * u64(u32(v05))
    v13 ^= v01
    v13 = bits.rotate_left_64(v13, -32)
    v09 += v13 + 2 * u64(u32(v09)) * u64(u32(v13))
    v05 ^= v09
    v05 = bits.rotate_left_64(v05, -24)

    v01 += v05 + 2 * u64(u32(v01)) * u64(u32(v05))
    v13 ^= v01
    v13 = bits.rotate_left_64(v13, -16)
    v09 += v13 + 2 * u64(u32(v09)) * u64(u32(v13))
    v05 ^= v09
    v05 = bits.rotate_left_64(v05, -63)

    v02 += v06 + 2 * u64(u32(v02)) * u64(u32(v06))
    v14 ^= v02
    v14 = bits.rotate_left_64(v14, -32)
    v10 += v14 + 2 * u64(u32(v10)) * u64(u32(v14))
    v06 ^= v10
    v06 = bits.rotate_left_64(v06, -24)

    v02 += v06 + 2 * u64(u32(v02)) * u64(u32(v06))
    v14 ^= v02
    v14 = bits.rotate_left_64(v14, -16)
    v10 += v14 + 2 * u64(u32(v10)) * u64(u32(v14))
    v06 ^= v10
    v06 = bits.rotate_left_64(v06, -63)

    v03 += v07 + 2 * u64(u32(v03)) * u64(u32(v07))
    v15 ^= v03
    v15 = bits.rotate_left_64(v15, -32)
    v11 += v15 + 2 * u64(u32(v11)) * u64(u32(v15))
    v07 ^= v11
    v07 = bits.rotate_left_64(v07, -24)

    v03 += v07 + 2 * u64(u32(v03)) * u64(u32(v07))
    v15 ^= v03
    v15 = bits.rotate_left_64(v15, -16)
    v11 += v15 + 2 * u64(u32(v11)) * u64(u32(v15))
    v07 ^= v11
    v07 = bits.rotate_left_64(v07, -63)

    v00 += v05 + 2 * u64(u32(v00)) * u64(u32(v05))
    v15 ^= v00
    v15 = bits.rotate_left_64(v15, -32)
    v10 += v15 + 2 * u64(u32(v10)) * u64(u32(v15))
    v05 ^= v10
    v05 = bits.rotate_left_64(v05, -24)

    v00 += v05 + 2 * u64(u32(v00)) * u64(u32(v05))
    v15 ^= v00
    v15 = bits.rotate_left_64(v15, -16)
    v10 += v15 + 2 * u64(u32(v10)) * u64(u32(v15))
    v05 ^= v10
    v05 = bits.rotate_left_64(v05, -63)

    v01 += v06 + 2 * u64(u32(v01)) * u64(u32(v06))
    v12 ^= v01
    v12 = bits.rotate_left_64(v12, -32)
    v11 += v12 + 2 * u64(u32(v11)) * u64(u32(v12))
    v06 ^= v11
    v06 = bits.rotate_left_64(v06, -24)

    v01 += v06 + 2 * u64(u32(v01)) * u64(u32(v06))
    v12 ^= v01
    v12 = bits.rotate_left_64(v12, -16)
    v11 += v12 + 2 * u64(u32(v11)) * u64(u32(v12))
    v06 ^= v11
    v06 = bits.rotate_left_64(v06, -63)

    v02 += v07 + 2 * u64(u32(v02)) * u64(u32(v07))
    v13 ^= v02
    v13 = bits.rotate_left_64(v13, -32)
    v08 += v13 + 2 * u64(u32(v08)) * u64(u32(v13))
    v07 ^= v08
    v07 = bits.rotate_left_64(v07, -24)

    v02 += v07 + 2 * u64(u32(v02)) * u64(u32(v07))
    v13 ^= v02
    v13 = bits.rotate_left_64(v13, -16)
    v08 += v13 + 2 * u64(u32(v08)) * u64(u32(v13))
    v07 ^= v08
    v07 = bits.rotate_left_64(v07, -63)

    v03 += v04 + 2 * u64(u32(v03)) * u64(u32(v04))
    v14 ^= v03
    v14 = bits.rotate_left_64(v14, -32)
    v09 += v14 + 2 * u64(u32(v09)) * u64(u32(v14))
    v04 ^= v09
    v04 = bits.rotate_left_64(v04, -24)

    v03 += v04 + 2 * u64(u32(v03)) * u64(u32(v04))
    v14 ^= v03
    v14 = bits.rotate_left_64(v14, -16)
    v09 += v14 + 2 * u64(u32(v09)) * u64(u32(v14))
    v04 ^= v09
    v04 = bits.rotate_left_64(v04, -63)

    block[a] = v00
    block[b] = v01
    block[c] = v02
    block[d] = v03
    block[e] = v04
    block[f] = v05
    block[g] = v06
    block[h] = v07
    block[i] = v08
    block[j] = v09
    block[k] = v10
    block[l] = v11
    block[m] = v12
    block[n] = v13
    block[o] = v14
    block[p] = v15
}

fn index_alpha(random u64, lanes u32, segments u32, threads u32, pass u32, slice u32, lane u32, index u32) u32 {
    mut ref_lane := u32(random >> 32) % threads
    if pass == 0 && slice == 0 {
        ref_lane = lane
    }
    mut m := u32(3) * segments
    mut s := ((slice + 1) % sync_points) * segments
    if lane == ref_lane {
        m += index
    }
    if pass == 0 {
        m = slice * segments
        s = 0
        if slice == 0 || lane == ref_lane {
            m += index
        }
    }
    if index == 0 || lane == ref_lane {
        m--
    }
    return phi(random, u64(m), u64(s), ref_lane, lanes)
}

fn phi(random u64, m u64, s u64, lane u32, lanes u32) u32 {
    mut p := random & u64(0xffffffff)
    p = (p * p) >> 32
    p = (p * m) >> 32
    return lane * lanes + u32((s + m - (p + 1)) % u64(lanes))
}

fn encode_hash(mode int, salt []u8, hash []u8, time u32, memory u32, threads u8) string {
    salt_b64 := base64.encode(salt).replace_each(['=', ''])
    hash_b64 := base64.encode(hash).replace_each(['=', ''])
    return '\$${mode_name(mode)}\$v=${version}\$m=${memory},t=${time},p=${threads}\$${salt_b64}\$${hash_b64}'
}

fn decode_hash(encoded string) !DecodedHash {
    parts := encoded.split('\$')
    if parts.len != 6 || parts[0] != '' {
        return InvalidHashError{
            reason: 'invalid argon2 hash format'
        }
    }
    if parts[2] != 'v=${version}' {
        return InvalidHashError{
            reason: 'unsupported argon2 hash version'
        }
    }
    mode := parse_mode(parts[1])!
    mut memory := u32(0)
    mut time := u32(0)
    mut threads := u8(0)
    for param in parts[3].split(',') {
        name := param.all_before('=')
        value := param.all_after('=')
        match name {
            'm' {
                memory = value.u32()
            }
            't' {
                time = value.u32()
            }
            'p' {
                threads = u8(value.int())
            }
            else {
                return InvalidHashError{
                    reason: 'unsupported argon2 parameter ${name}'
                }
            }
        }
    }
    if time == 0 || threads == 0 {
        return InvalidHashError{
            reason: 'invalid argon2 parameters'
        }
    }
    salt := decode_base64(parts[4])
    hash := decode_base64(parts[5])
    if salt.len == 0 || hash.len == 0 {
        return InvalidHashError{
            reason: 'invalid argon2 hash payload'
        }
    }
    return DecodedHash{
        mode:    mode
        time:    time
        memory:  normalize_memory(memory, u32(threads))
        threads: threads
        salt:    salt
        hash:    hash
    }
}

fn decode_base64(value string) []u8 {
    mut padded := value
    match padded.len % 4 {
        2 { padded += '==' }
        3 { padded += '=' }
        else {}
    }

    return base64.decode(padded)
}

fn mode_name(mode int) string {
    return match mode {
        argon2_d { 'argon2d' }
        argon2_i { 'argon2i' }
        else { 'argon2id' }
    }
}

fn parse_mode(name string) !int {
    return match name {
        'argon2d' {
            argon2_d
        }
        'argon2i' {
            argon2_i
        }
        'argon2id' {
            argon2_id
        }
        else {
            return InvalidHashError{
                reason: 'unsupported argon2 algorithm ${name}'
            }
        }
    }
}