// Copyright (c) 2019-2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by a GPL license that can be found in the LICENSE file.
module main

import time
import encoding.base32
import encoding.binary
import crypto.hmac
import crypto.sha1
import crypto.rand as crypto_rand

struct TwoFactor {
    id int @[primary; sql: serial]
mut:
    user_id    int
    secret     string
    is_enabled bool
    created_at int
}

const totp_period = 30
const totp_digits = 6
const totp_issuer = 'Gitly'

fn (mut app App) find_two_factor(user_id int) ?TwoFactor {
    rows := sql app.db {
        select from TwoFactor where user_id == user_id limit 1
    } or { []TwoFactor{} }
    if rows.len == 0 {
        return none
    }
    return rows.first()
}

fn (mut app App) upsert_two_factor(user_id int, secret string, is_enabled bool) ! {
    if existing := app.find_two_factor(user_id) {
        id := existing.id
        sql app.db {
            update TwoFactor set secret = secret, is_enabled = is_enabled where id == id
        }!
        return
    }
    tf := TwoFactor{
        user_id:    user_id
        secret:     secret
        is_enabled: is_enabled
        created_at: int(time.now().unix())
    }
    sql app.db {
        insert tf into TwoFactor
    }!
}

fn (mut app App) delete_two_factor(user_id int) ! {
    sql app.db {
        delete from TwoFactor where user_id == user_id
    }!
}

fn (mut app App) user_has_two_factor(user_id int) bool {
    tf := app.find_two_factor(user_id) or { return false }
    return tf.is_enabled
}

fn generate_totp_secret() string {
    mut buf := []u8{len: 20}
    for i in 0 .. buf.len {
        buf[i] = u8(crypto_rand.int_u64(256) or { 0 })
    }
    enc := base32.encode_to_string(buf)
    return enc.trim_right('=')
}

fn decode_base32_secret(secret string) ![]u8 {
    mut padded := secret.to_upper().replace(' ', '')
    for padded.len % 8 != 0 {
        padded += '='
    }
    return base32.decode(padded.bytes())!
}

fn hotp(key []u8, counter u64) int {
    mut buf := []u8{len: 8}
    binary.big_endian_put_u64(mut buf, counter)
    mac := hmac.new(key, buf, sha1.sum, sha1.block_size)
    offset := int(mac[mac.len - 1] & 0x0f)
    bin := ((u32(mac[offset]) & 0x7f) << 24) | ((u32(mac[offset + 1]) & 0xff) << 16) | ((u32(mac[
        offset + 2]) & 0xff) << 8) | (u32(mac[offset + 3]) & 0xff)
    mut modulus := u32(1)
    for _ in 0 .. totp_digits {
        modulus *= 10
    }
    return int(bin % modulus)
}

fn totp_code_for(secret string, t i64) !string {
    key := decode_base32_secret(secret)!
    counter := u64(t / totp_period)
    code := hotp(key, counter)
    mut s := code.str()
    for s.len < totp_digits {
        s = '0' + s
    }
    return s
}

fn verify_totp(secret string, code string) bool {
    if code.len != totp_digits {
        return false
    }
    now := time.now().unix()
    for offset in [i64(-1), 0, 1] {
        expected := totp_code_for(secret, now + offset * totp_period) or { continue }
        if expected == code {
            return true
        }
    }
    return false
}

fn totp_provisioning_uri(username string, secret string) string {
    label := '${totp_issuer}:${username}'
    return 'otpauth://totp/${label}?secret=${secret}&issuer=${totp_issuer}&algorithm=SHA1&digits=${totp_digits}&period=${totp_period}'
}