//! Módulo de criptografía - SHA256, ChaCha20-Poly1305
//!
//! Implementación pura en Zig sin dependencias externas.

const std = @import("std");

// =============================================================================
// SHA-256
// =============================================================================

/// Longitud del hash SHA256 en bytes
pub const SHA256_DIGEST_LENGTH: usize = 32;

/// Calcula SHA256 de los datos
pub fn sha256(data: []const u8) [SHA256_DIGEST_LENGTH]u8 {
    var hasher = Sha256.init();
    hasher.update(data);
    return hasher.final();
}

/// Estado del hasher SHA256
pub const Sha256 = struct {
    state: [8]u32,
    buf: [64]u8,
    buf_len: usize,
    total_len: u64,

    const K: [64]u32 = .{
        0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5,
        0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
        0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3,
        0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
        0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc,
        0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
        0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7,
        0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
        0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13,
        0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
        0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3,
        0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
        0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5,
        0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
        0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208,
        0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2,
    };

    pub fn init() Sha256 {
        return .{
            .state = .{
                0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a,
                0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19,
            },
            .buf = undefined,
            .buf_len = 0,
            .total_len = 0,
        };
    }

    pub fn update(self: *Sha256, data: []const u8) void {
        var input = data;
        self.total_len += data.len;

        // Si hay datos en el buffer, intentar completar un bloque
        if (self.buf_len > 0) {
            const space = 64 - self.buf_len;
            const to_copy = @min(space, input.len);
            @memcpy(self.buf[self.buf_len .. self.buf_len + to_copy], input[0..to_copy]);
            self.buf_len += to_copy;
            input = input[to_copy..];

            if (self.buf_len == 64) {
                self.processBlock(&self.buf);
                self.buf_len = 0;
            }
        }

        // Procesar bloques completos
        while (input.len >= 64) {
            self.processBlock(input[0..64]);
            input = input[64..];
        }

        // Guardar resto en buffer
        if (input.len > 0) {
            @memcpy(self.buf[0..input.len], input);
            self.buf_len = input.len;
        }
    }

    pub fn final(self: *Sha256) [SHA256_DIGEST_LENGTH]u8 {
        // Padding
        const total_bits = self.total_len * 8;
        self.buf[self.buf_len] = 0x80;
        self.buf_len += 1;

        // Si no hay espacio para la longitud, procesar bloque extra
        if (self.buf_len > 56) {
            @memset(self.buf[self.buf_len..64], 0);
            self.processBlock(&self.buf);
            self.buf_len = 0;
        }

        @memset(self.buf[self.buf_len..56], 0);
        std.mem.writeInt(u64, self.buf[56..64], total_bits, .big);
        self.processBlock(&self.buf);

        // Convertir estado a bytes
        var result: [32]u8 = undefined;
        for (self.state, 0..) |s, i| {
            std.mem.writeInt(u32, result[i * 4 ..][0..4], s, .big);
        }
        return result;
    }

    fn processBlock(self: *Sha256, block: *const [64]u8) void {
        var w: [64]u32 = undefined;

        // Preparar schedule
        for (0..16) |i| {
            w[i] = std.mem.readInt(u32, block[i * 4 ..][0..4], .big);
        }
        for (16..64) |i| {
            const s0 = rotr(w[i - 15], 7) ^ rotr(w[i - 15], 18) ^ (w[i - 15] >> 3);
            const s1 = rotr(w[i - 2], 17) ^ rotr(w[i - 2], 19) ^ (w[i - 2] >> 10);
            w[i] = w[i - 16] +% s0 +% w[i - 7] +% s1;
        }

        var a = self.state[0];
        var b = self.state[1];
        var c = self.state[2];
        var d = self.state[3];
        var e = self.state[4];
        var f = self.state[5];
        var g = self.state[6];
        var h = self.state[7];

        for (0..64) |i| {
            const S1 = rotr(e, 6) ^ rotr(e, 11) ^ rotr(e, 25);
            const ch = (e & f) ^ (~e & g);
            const temp1 = h +% S1 +% ch +% K[i] +% w[i];
            const S0 = rotr(a, 2) ^ rotr(a, 13) ^ rotr(a, 22);
            const maj = (a & b) ^ (a & c) ^ (b & c);
            const temp2 = S0 +% maj;

            h = g;
            g = f;
            f = e;
            e = d +% temp1;
            d = c;
            c = b;
            b = a;
            a = temp1 +% temp2;
        }

        self.state[0] +%= a;
        self.state[1] +%= b;
        self.state[2] +%= c;
        self.state[3] +%= d;
        self.state[4] +%= e;
        self.state[5] +%= f;
        self.state[6] +%= g;
        self.state[7] +%= h;
    }

    fn rotr(x: u32, comptime n: u5) u32 {
        return std.math.rotr(u32, x, n);
    }
};

// =============================================================================
// ChaCha20-Poly1305
// =============================================================================

pub const CHACHA20_KEY_SIZE: usize = 32;
pub const CHACHA20_NONCE_SIZE: usize = 12;
pub const XCHACHA20_NONCE_SIZE: usize = 24;
pub const POLY1305_TAG_SIZE: usize = 16;

/// ChaCha20 block cipher
pub const ChaCha20 = struct {
    state: [16]u32,

    pub fn init(key: *const [32]u8, nonce: *const [12]u8, counter: u32) ChaCha20 {
        return .{
            .state = .{
                0x61707865, 0x3320646e, 0x79622d32, 0x6b206574,
                std.mem.readInt(u32, key[0..4], .little),
                std.mem.readInt(u32, key[4..8], .little),
                std.mem.readInt(u32, key[8..12], .little),
                std.mem.readInt(u32, key[12..16], .little),
                std.mem.readInt(u32, key[16..20], .little),
                std.mem.readInt(u32, key[20..24], .little),
                std.mem.readInt(u32, key[24..28], .little),
                std.mem.readInt(u32, key[28..32], .little),
                counter,
                std.mem.readInt(u32, nonce[0..4], .little),
                std.mem.readInt(u32, nonce[4..8], .little),
                std.mem.readInt(u32, nonce[8..12], .little),
            },
        };
    }

    pub fn xor(self: *ChaCha20, out: []u8, in: []const u8) void {
        var remaining = in;
        var out_ptr = out;

        while (remaining.len > 0) {
            const keystream = self.block();

            const to_process = @min(remaining.len, 64);
            for (0..to_process) |i| {
                out_ptr[i] = remaining[i] ^ keystream[i];
            }

            remaining = remaining[to_process..];
            out_ptr = out_ptr[to_process..];

            // Incrementar contador
            self.state[12] +%= 1;
        }
    }

    fn block(self: *ChaCha20) [64]u8 {
        var working = self.state;

        // 20 rounds (10 double rounds)
        for (0..10) |_| {
            quarterRound(&working, 0, 4, 8, 12);
            quarterRound(&working, 1, 5, 9, 13);
            quarterRound(&working, 2, 6, 10, 14);
            quarterRound(&working, 3, 7, 11, 15);
            quarterRound(&working, 0, 5, 10, 15);
            quarterRound(&working, 1, 6, 11, 12);
            quarterRound(&working, 2, 7, 8, 13);
            quarterRound(&working, 3, 4, 9, 14);
        }

        // Add original state
        for (&working, self.state) |*w, s| {
            w.* +%= s;
        }

        // Convert to bytes
        var result: [64]u8 = undefined;
        for (working, 0..) |w, i| {
            std.mem.writeInt(u32, result[i * 4 ..][0..4], w, .little);
        }
        return result;
    }

    fn quarterRound(state: *[16]u32, a: usize, b: usize, c: usize, d: usize) void {
        state[a] +%= state[b];
        state[d] ^= state[a];
        state[d] = rotl(state[d], 16);
        state[c] +%= state[d];
        state[b] ^= state[c];
        state[b] = rotl(state[b], 12);
        state[a] +%= state[b];
        state[d] ^= state[a];
        state[d] = rotl(state[d], 8);
        state[c] +%= state[d];
        state[b] ^= state[c];
        state[b] = rotl(state[b], 7);
    }

    fn rotl(x: u32, comptime n: u5) u32 {
        return std.math.rotl(u32, x, n);
    }
};

/// XChaCha20 - ChaCha20 con nonce extendido de 24 bytes
pub const XChaCha20 = struct {
    chacha: ChaCha20,

    pub fn init(key: *const [32]u8, nonce: *const [24]u8, counter: u32) XChaCha20 {
        // HChaCha20 para derivar subkey
        const subkey = hchacha20(key, nonce[0..16]);
        var short_nonce: [12]u8 = undefined;
        @memset(short_nonce[0..4], 0);
        @memcpy(short_nonce[4..12], nonce[16..24]);

        return .{
            .chacha = ChaCha20.init(&subkey, &short_nonce, counter),
        };
    }

    pub fn xor(self: *XChaCha20, out: []u8, in: []const u8) void {
        self.chacha.xor(out, in);
    }
};

fn hchacha20(key: *const [32]u8, nonce: *const [16]u8) [32]u8 {
    var state: [16]u32 = .{
        0x61707865, 0x3320646e, 0x79622d32, 0x6b206574,
        std.mem.readInt(u32, key[0..4], .little),
        std.mem.readInt(u32, key[4..8], .little),
        std.mem.readInt(u32, key[8..12], .little),
        std.mem.readInt(u32, key[12..16], .little),
        std.mem.readInt(u32, key[16..20], .little),
        std.mem.readInt(u32, key[20..24], .little),
        std.mem.readInt(u32, key[24..28], .little),
        std.mem.readInt(u32, key[28..32], .little),
        std.mem.readInt(u32, nonce[0..4], .little),
        std.mem.readInt(u32, nonce[4..8], .little),
        std.mem.readInt(u32, nonce[8..12], .little),
        std.mem.readInt(u32, nonce[12..16], .little),
    };

    for (0..10) |_| {
        ChaCha20.quarterRound(&state, 0, 4, 8, 12);
        ChaCha20.quarterRound(&state, 1, 5, 9, 13);
        ChaCha20.quarterRound(&state, 2, 6, 10, 14);
        ChaCha20.quarterRound(&state, 3, 7, 11, 15);
        ChaCha20.quarterRound(&state, 0, 5, 10, 15);
        ChaCha20.quarterRound(&state, 1, 6, 11, 12);
        ChaCha20.quarterRound(&state, 2, 7, 8, 13);
        ChaCha20.quarterRound(&state, 3, 4, 9, 14);
    }

    var result: [32]u8 = undefined;
    std.mem.writeInt(u32, result[0..4], state[0], .little);
    std.mem.writeInt(u32, result[4..8], state[1], .little);
    std.mem.writeInt(u32, result[8..12], state[2], .little);
    std.mem.writeInt(u32, result[12..16], state[3], .little);
    std.mem.writeInt(u32, result[16..20], state[12], .little);
    std.mem.writeInt(u32, result[20..24], state[13], .little);
    std.mem.writeInt(u32, result[24..28], state[14], .little);
    std.mem.writeInt(u32, result[28..32], state[15], .little);
    return result;
}

// =============================================================================
// Poly1305 MAC
// =============================================================================

/// Poly1305 message authentication code
pub const Poly1305 = struct {
    r: [3]u64,
    h: [3]u64,
    pad: [2]u64,
    leftover: usize,
    buffer: [16]u8,

    pub fn init(key: *const [32]u8) Poly1305 {
        // r = key[0..16] con bits específicos enmascarados
        var r0 = std.mem.readInt(u64, key[0..8], .little);
        var r1 = std.mem.readInt(u64, key[8..16], .little);

        // Clamp r
        r0 &= 0x0ffffffc0fffffff;
        r1 &= 0x0ffffffc0ffffffc;

        return .{
            .r = .{ r0 & 0xfffffffffff, (r0 >> 44) | ((r1 & 0xffffff) << 20), r1 >> 24 },
            .h = .{ 0, 0, 0 },
            .pad = .{
                std.mem.readInt(u64, key[16..24], .little),
                std.mem.readInt(u64, key[24..32], .little),
            },
            .leftover = 0,
            .buffer = undefined,
        };
    }

    pub fn update(self: *Poly1305, data: []const u8) void {
        var input = data;

        // Procesar datos en buffer primero
        if (self.leftover > 0) {
            const want = 16 - self.leftover;
            const have = @min(want, input.len);
            @memcpy(self.buffer[self.leftover .. self.leftover + have], input[0..have]);
            input = input[have..];
            self.leftover += have;

            if (self.leftover == 16) {
                self.processBlock(&self.buffer, false);
                self.leftover = 0;
            }
        }

        // Procesar bloques completos
        while (input.len >= 16) {
            self.processBlock(input[0..16], false);
            input = input[16..];
        }

        // Guardar resto
        if (input.len > 0) {
            @memcpy(self.buffer[0..input.len], input);
            self.leftover = input.len;
        }
    }

    pub fn final(self: *Poly1305) [16]u8 {
        // Procesar último bloque parcial
        if (self.leftover > 0) {
            self.buffer[self.leftover] = 1;
            @memset(self.buffer[self.leftover + 1 .. 16], 0);
            self.processBlock(&self.buffer, true);
        }

        // Finalizar
        var h0 = self.h[0];
        var h1 = self.h[1];
        var h2 = self.h[2];

        // Reducción final
        var c: u64 = 0;
        c = h0 >> 44;
        h0 &= 0xfffffffffff;
        h1 += c;
        c = h1 >> 44;
        h1 &= 0xfffffffffff;
        h2 += c;
        c = h2 >> 42;
        h2 &= 0x3ffffffffff;
        h0 += c * 5;
        c = h0 >> 44;
        h0 &= 0xfffffffffff;
        h1 += c;

        // Computar h + -p
        var g0 = h0 +% 5;
        c = g0 >> 44;
        g0 &= 0xfffffffffff;
        var g1 = h1 +% c;
        c = g1 >> 44;
        g1 &= 0xfffffffffff;
        var g2 = h2 +% c -% (1 << 42);

        // Seleccionar h o h + -p
        c = (g2 >> 63) -% 1;
        g0 &= c;
        g1 &= c;
        g2 &= c;
        c = ~c;
        h0 = (h0 & c) | g0;
        h1 = (h1 & c) | g1;
        h2 = (h2 & c) | g2;

        // h = h + pad
        const t0 = self.pad[0];
        const t1 = self.pad[1];

        h0 +%= t0 & 0xfffffffffff;
        c = h0 >> 44;
        h0 &= 0xfffffffffff;
        h1 +%= ((t0 >> 44) | (t1 << 20)) & 0xfffffffffff;
        h1 += c;
        c = h1 >> 44;
        h1 &= 0xfffffffffff;
        h2 +%= (t1 >> 24) & 0x3ffffffffff;
        h2 += c;
        h2 &= 0x3ffffffffff;

        // Convertir a bytes
        var result: [16]u8 = undefined;
        const full: u128 = @as(u128, h0) | (@as(u128, h1) << 44) | (@as(u128, h2) << 88);
        std.mem.writeInt(u128, &result, full, .little);
        return result;
    }

    fn processBlock(self: *Poly1305, block: *const [16]u8, is_final: bool) void {
        const hibit: u64 = if (is_final) 0 else (1 << 40);

        // h += m
        const t0 = std.mem.readInt(u64, block[0..8], .little);
        const t1 = std.mem.readInt(u64, block[8..16], .little);

        self.h[0] += t0 & 0xfffffffffff;
        self.h[1] += ((t0 >> 44) | (t1 << 20)) & 0xfffffffffff;
        self.h[2] += ((t1 >> 24) & 0x3ffffffffff) | hibit;

        // h *= r (mod 2^130 - 5)
        const r0 = self.r[0];
        const r1 = self.r[1];
        const r2 = self.r[2];
        const s1 = r1 * 5;
        const s2 = r2 * 5;

        var d0: u128 = @as(u128, self.h[0]) * r0;
        d0 += @as(u128, self.h[1]) * s2;
        d0 += @as(u128, self.h[2]) * s1;
        var d1: u128 = @as(u128, self.h[0]) * r1;
        d1 += @as(u128, self.h[1]) * r0;
        d1 += @as(u128, self.h[2]) * s2;
        var d2: u128 = @as(u128, self.h[0]) * r2;
        d2 += @as(u128, self.h[1]) * r1;
        d2 += @as(u128, self.h[2]) * r0;

        // Reducción parcial
        var c: u64 = @truncate(d0 >> 44);
        self.h[0] = @truncate(d0 & 0xfffffffffff);
        d1 += c;
        c = @truncate(d1 >> 44);
        self.h[1] = @truncate(d1 & 0xfffffffffff);
        d2 += c;
        c = @truncate(d2 >> 42);
        self.h[2] = @truncate(d2 & 0x3ffffffffff);
        self.h[0] += c * 5;
        c = self.h[0] >> 44;
        self.h[0] &= 0xfffffffffff;
        self.h[1] += c;
    }
};

// =============================================================================
// ChaCha20-Poly1305 AEAD
// =============================================================================

/// Cifra datos con ChaCha20-Poly1305
/// Retorna: nonce || ciphertext || tag
pub fn chachaPoly1305Encrypt(
    key: *const [32]u8,
    nonce: *const [12]u8,
    plaintext: []const u8,
    aad: []const u8,
    allocator: std.mem.Allocator,
) ![]u8 {
    const result = try allocator.alloc(u8, 12 + plaintext.len + 16);
    errdefer allocator.free(result);

    @memcpy(result[0..12], nonce);

    // Generar keystream para Poly1305
    var chacha = ChaCha20.init(key, nonce, 0);
    var poly_key: [64]u8 = undefined;
    var zeros: [64]u8 = [_]u8{0} ** 64;
    chacha.xor(&poly_key, &zeros);

    // Cifrar
    chacha = ChaCha20.init(key, nonce, 1);
    chacha.xor(result[12 .. 12 + plaintext.len], plaintext);

    // Calcular tag
    var poly = Poly1305.init(poly_key[0..32]);
    poly.update(aad);
    if (aad.len % 16 != 0) {
        var pad: [16]u8 = [_]u8{0} ** 16;
        poly.update(pad[0 .. 16 - (aad.len % 16)]);
    }
    poly.update(result[12 .. 12 + plaintext.len]);
    if (plaintext.len % 16 != 0) {
        var pad: [16]u8 = [_]u8{0} ** 16;
        poly.update(pad[0 .. 16 - (plaintext.len % 16)]);
    }
    var lens: [16]u8 = undefined;
    std.mem.writeInt(u64, lens[0..8], aad.len, .little);
    std.mem.writeInt(u64, lens[8..16], plaintext.len, .little);
    poly.update(&lens);

    const tag = poly.final();
    @memcpy(result[12 + plaintext.len ..][0..16], &tag);

    return result;
}

/// Descifra datos con ChaCha20-Poly1305
/// Input format: nonce (12) || ciphertext || tag (16)
pub fn chachaPoly1305Decrypt(
    key: *const [32]u8,
    data: []const u8,
    aad: []const u8,
    allocator: std.mem.Allocator,
) ![]u8 {
    if (data.len < 12 + 16) return error.InvalidData;

    const nonce = data[0..12];
    const ciphertext = data[12 .. data.len - 16];
    const tag = data[data.len - 16 ..][0..16];

    // Verificar tag
    var chacha = ChaCha20.init(key, nonce, 0);
    var poly_key: [64]u8 = undefined;
    var zeros: [64]u8 = [_]u8{0} ** 64;
    chacha.xor(&poly_key, &zeros);

    var poly = Poly1305.init(poly_key[0..32]);
    poly.update(aad);
    if (aad.len % 16 != 0) {
        var pad: [16]u8 = [_]u8{0} ** 16;
        poly.update(pad[0 .. 16 - (aad.len % 16)]);
    }
    poly.update(ciphertext);
    if (ciphertext.len % 16 != 0) {
        var pad: [16]u8 = [_]u8{0} ** 16;
        poly.update(pad[0 .. 16 - (ciphertext.len % 16)]);
    }
    var lens: [16]u8 = undefined;
    std.mem.writeInt(u64, lens[0..8], aad.len, .little);
    std.mem.writeInt(u64, lens[8..16], ciphertext.len, .little);
    poly.update(&lens);

    const computed_tag = poly.final();
    if (!std.mem.eql(u8, &computed_tag, tag)) {
        return error.AuthenticationFailed;
    }

    // Descifrar
    const result = try allocator.alloc(u8, ciphertext.len);
    chacha = ChaCha20.init(key, nonce, 1);
    chacha.xor(result, ciphertext);

    return result;
}

// =============================================================================
// Tests
// =============================================================================

test "sha256 empty" {
    const hash = sha256("");
    const expected = [_]u8{
        0xe3, 0xb0, 0xc4, 0x42, 0x98, 0xfc, 0x1c, 0x14,
        0x9a, 0xfb, 0xf4, 0xc8, 0x99, 0x6f, 0xb9, 0x24,
        0x27, 0xae, 0x41, 0xe4, 0x64, 0x9b, 0x93, 0x4c,
        0xa4, 0x95, 0x99, 0x1b, 0x78, 0x52, 0xb8, 0x55,
    };
    try std.testing.expectEqualSlices(u8, &expected, &hash);
}

test "sha256 abc" {
    const hash = sha256("abc");
    const expected = [_]u8{
        0xba, 0x78, 0x16, 0xbf, 0x8f, 0x01, 0xcf, 0xea,
        0x41, 0x41, 0x40, 0xde, 0x5d, 0xae, 0x22, 0x23,
        0xb0, 0x03, 0x61, 0xa3, 0x96, 0x17, 0x7a, 0x9c,
        0xb4, 0x10, 0xff, 0x61, 0xf2, 0x00, 0x15, 0xad,
    };
    try std.testing.expectEqualSlices(u8, &expected, &hash);
}

test "chacha20 basic" {
    const key = [_]u8{0} ** 32;
    const nonce = [_]u8{0} ** 12;
    var chacha = ChaCha20.init(&key, &nonce, 0);

    var out: [64]u8 = undefined;
    const zeros = [_]u8{0} ** 64;
    chacha.xor(&out, &zeros);

    // Verificar que no es todo ceros (el keystream se aplicó)
    try std.testing.expect(!std.mem.eql(u8, &out, &zeros));
}