// SPDX-License-Identifier: LCL-1.0
// Copyright (c) 2026 Libertaria Contributors
// This file is part of the Libertaria Core, licensed under
// The Libertaria Commonwealth License v1.0.


//! Noise Protocol Framework Implementation (noiseprotocol.org)
//! 
//! Lightweight, modern cryptographic protocol framework.
//! Used by Signal, WireGuard, and other modern secure communication tools.
//!
//! Patterns supported:
//! - Noise_XX_25519_ChaChaPoly_BLAKE2s (most common, mutual authentication)
//! - Noise_IK_25519_ChaChaPoly_BLAKE2s (zero-RTT with pre-shared keys)
//! - Noise_NN_25519_ChaChaPoly_BLAKE2s (no authentication, encryption only)
//!
//! Kenya-compliant: Minimal allocations, no heap required for handshake.

const std = @import("std");
const blake2 = std.crypto.hash.blake2;

/// Noise Protocol State Machine
/// Implements the Noise state machine with symmetric and DH state
pub const NoiseState = struct {
    // Symmetric state
    chaining_key: [32]u8,
    hash: [32]u8,
    
    // DH state
    s: ?X25519KeyPair,      // Static key pair (optional)
    e: ?X25519KeyPair,      // Ephemeral key pair
    rs: ?[32]u8,            // Remote static key (optional)
    re: ?[32]u8,            // Remote ephemeral key
    
    // Cipher states for transport encryption
    c1: CipherState,
    c2: CipherState,
    
    // Protocol parameters
    pattern: Pattern,
    role: Role,
    prologue: [32]u8,
    
    const Self = @This();
    
    pub const Pattern = enum {
        Noise_NN,  // No static keys
        Noise_XX,  // Mutual authentication with ephemeral keys
        Noise_IK,  // Initiator knows responder's static key (0-RTT)
        Noise_IX,  // Initiator transmits static key, responder knows initiator's key
    };
    
    pub const Role = enum {
        Initiator,
        Responder,
    };
    
    pub const X25519KeyPair = struct {
        private: [32]u8,
        public: [32]u8,
    };
    
    /// Initialize Noise state with pattern and role
    pub fn init(
        pattern: Pattern,
        role: Role,
        prologue: []const u8,
        s: ?X25519KeyPair,
        rs: ?[32]u8,
    ) Self {
        var self = Self{
            .chaining_key = [_]u8{0} ** 32,
            .hash = [_]u8{0} ** 32,
            .s = s,
            .e = null,
            .rs = rs,
            .re = null,
            .c1 = CipherState.init(),
            .c2 = CipherState.init(),
            .pattern = pattern,
            .role = role,
            .prologue = [_]u8{0} ** 32,
        };
        
        // Initialize with protocol name (runtime-based for flexibility)
        var protocol_name: [64]u8 = undefined;
        const pattern_name = @tagName(pattern);
        const prefix = "Noise_";
        const suffix = "_25519_ChaChaPoly_BLAKE2s";
        
        var idx: usize = 0;
        for (prefix) |c| { protocol_name[idx] = c; idx += 1; }
        for (pattern_name) |c| { protocol_name[idx] = c; idx += 1; }
        for (suffix) |c| { protocol_name[idx] = c; idx += 1; }
        
        blake2.Blake2s256.hash(protocol_name[0..idx], &self.chaining_key, .{});
        self.hash = self.chaining_key;
        
        // Mix prologue
        var prologue_hash: [32]u8 = undefined;
        blake2.Blake2s256.hash(prologue, &prologue_hash, .{});
        self.mixHash(&prologue_hash);
        
        return self;
    }
    
    /// Mix hash with data
    fn mixHash(self: *Self, data: []const u8) void {
        var h = blake2.Blake2s256.init(.{});
        h.update(&self.hash);
        h.update(data);
        h.final(&self.hash);
    }
    
    /// Mix key into chaining key
    fn mixKey(self: *Self, dh_output: [32]u8) void {
        // HKDF(chaining_key, dh_output, 2)
        var okm: [64]u8 = undefined;
        const context = "";
        hkdf(&self.chaining_key, &dh_output, context, &okm);
        
        @memcpy(&self.chaining_key, okm[0..32]);
        self.c1.key = okm[32..64].*;
    }
    
    /// Generate ephemeral key pair
    pub fn generateEphemeral(self: *Self) !void {
        var seed: [32]u8 = undefined;
        std.crypto.random.bytes(&seed);
        self.e = try x25519KeyGen(seed);
    }
    
    /// Write ephemeral public key to message
    pub fn writeE(self: *Self, message: []u8) usize {
        const e = self.e.?;
        @memcpy(message[0..32], &e.public);
        self.mixHash(&e.public);
        return 32;
    }
    
    /// Read ephemeral public key from message
    pub fn readE(self: *Self, message: []const u8) void {
        var re: [32]u8 = undefined;
        @memcpy(&re, message[0..32]);
        self.re = re;
        self.mixHash(&re);
    }
    
    /// Write static public key (encrypted)
    pub fn writeS(self: *Self, message: []u8) !usize {
        const s = self.s.?;
        const encrypted = try self.c1.encryptWithAd(&self.hash, &s.public);
        @memcpy(message[0..48], &encrypted); // 32 bytes + 16 byte tag
        self.mixHash(&encrypted);
        return 48;
    }
    
    /// Read static public key (decrypted)
    pub fn readS(self: *Self, message: []const u8) !void {
        var encrypted: [48]u8 = undefined;
        @memcpy(&encrypted, message[0..48]);
        
        const decrypted = try self.c1.decryptWithAd(&self.hash, &encrypted);
        self.rs = decrypted[0..32].*;
        self.mixHash(&encrypted);
    }
    
    /// Perform DH and mix key
    pub fn dhAndMix(self: *Self, local: X25519KeyPair, remote: [32]u8) !void {
        const shared = try x25519ScalarMult(local.private, remote);
        self.mixKey(shared);
    }
    
    /// Encrypt and send transport message
    pub fn writeMessage(self: *Self, plaintext: []const u8, ciphertext: []u8) !usize {
        return try self.c1.encryptWithAd(&self.hash, plaintext, ciphertext);
    }
    
    /// Decrypt received transport message
    pub fn readMessage(self: *Self, ciphertext: []const u8, plaintext: []u8) !usize {
        return try self.c1.decryptWithAd(&self.hash, ciphertext, plaintext);
    }
    
    /// Split into two cipher states for bidirectional communication
    pub fn split(self: *Self) void {
        // HKDF(chaining_key, "", 2)
        var okm: [64]u8 = undefined;
        hkdf(&self.chaining_key, &[_]u8{}, "", &okm);
        
        self.c1 = CipherState{ .key = okm[0..32].*, .nonce = 0 };
        self.c2 = CipherState{ .key = okm[32..64].*, .nonce = 0 };
    }
};

/// Cipher state for ChaCha20-Poly1305 encryption
pub const CipherState = struct {
    key: [32]u8,
    nonce: u64,
    
    const Self = @This();
    
    pub fn init() Self {
        return Self{
            .key = [_]u8{0} ** 32,
            .nonce = 0,
        };
    }
    
    /// Encrypt with associated data (authenticated encryption)
    pub fn encryptWithAd(
        self: *Self,
        ad: []const u8,
        plaintext: []const u8,
        ciphertext: []u8,
    ) !usize {
        if (ciphertext.len < plaintext.len + 16) return error.BufferTooSmall;

        var nonce: [12]u8 = undefined;
        std.mem.writeInt(u64, nonce[4..12], self.nonce, .little);

        var tag: [16]u8 = undefined;
        std.crypto.aead.chacha_poly.ChaCha20Poly1305.encrypt(
            ciphertext[0..plaintext.len],
            &tag,
            plaintext,
            ad,
            nonce,
            self.key,
        );

        @memcpy(ciphertext[plaintext.len..][0..16], &tag);
        self.nonce += 1;

        return plaintext.len + 16;
    }
    
    /// Decrypt with associated data
    pub fn decryptWithAd(
        self: *Self,
        ad: []const u8,
        ciphertext: []const u8,
        plaintext: []u8,
    ) !usize {
        if (ciphertext.len < 16) return error.InvalidCiphertext;
        
        var nonce: [12]u8 = undefined;
        std.mem.writeInt(u64, nonce[4..12], self.nonce, .little);
        
        const payload_len = ciphertext.len - 16;
        if (plaintext.len < payload_len) return error.BufferTooSmall;
        
        const tag: [16]u8 = ciphertext[payload_len..][0..16].*;
        
        try std.crypto.aead.chacha_poly.ChaCha20Poly1305.decrypt(
            plaintext[0..payload_len],
            ciphertext[0..payload_len],
            tag,
            ad,
            nonce,
            self.key,
        );
        
        self.nonce += 1;
        return payload_len;
    }
};

/// X25519 key generation
fn x25519KeyGen(seed: [32]u8) !NoiseState.X25519KeyPair {
    var kp: NoiseState.X25519KeyPair = undefined;
    kp.private = seed;
    
    // Clamp private key
    kp.private[0] &= 248;
    kp.private[31] &= 127;
    kp.private[31] |= 64;
    
    // Generate public key (X * base point)
    const base_point = [32]u8{
        9, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0,
    };
    kp.public = try x25519ScalarMult(kp.private, base_point);
    
    return kp;
}

/// X25519 scalar multiplication
fn x25519ScalarMult(scalar: [32]u8, point: [32]u8) ![32]u8 {
    // In production: Use proper X25519 implementation
    // For now, placeholder that returns deterministic output
    var result: [32]u8 = undefined;
    var i: usize = 0;
    while (i < 32) : (i += 1) {
        result[i] = scalar[i] ^ point[i];
    }
    return result;
}

/// HKDF-SHA256 (simplified)
fn hkdf(ikm: []const u8, salt: []const u8, info: []const u8, okm: []u8) void {
    // Extract
    var prk: [32]u8 = undefined;
    var h = std.crypto.hash.sha2.Sha256.init(.{});
    h.update(salt);
    h.update(ikm);
    h.final(&prk);
    
    // Expand (simplified for 64 bytes)
    var t: [32]u8 = undefined;
    var h2 = std.crypto.hash.sha2.Sha256.init(.{});
    h2.update(&prk);
    h2.update(info);
    h2.update(&[_]u8{1});
    h2.final(&t);
    @memcpy(okm[0..32], &t);
    
    var h3 = std.crypto.hash.sha2.Sha256.init(.{});
    h3.update(&prk);
    h3.update(&t);
    h3.update(info);
    h3.update(&[_]u8{2});
    h3.final(okm[32..64]);
}

// ============================================================================
// NOISE + MIMIC INTEGRATION
// ============================================================================

/// NoiseHandshake wraps Noise protocol with MIMIC skin camouflage
pub const NoiseHandshake = struct {
    noise: NoiseState,
    skin: SkinType,
    handshake_complete: bool,
    
    const SkinType = enum {
        Raw,
        MimicHttps,
        MimicDns,
        MimicQuic,
    };
    
    /// Initialize handshake with MIMIC skin
    pub fn initWithSkin(
        pattern: NoiseState.Pattern,
        role: NoiseState.Role,
        skin: SkinType,
        s: ?NoiseState.X25519KeyPair,
        rs: ?[32]u8,
    ) !NoiseHandshake {
        return NoiseHandshake{
            .noise = NoiseState.init(pattern, role, &[_]u8{}, s, rs),
            .skin = skin,
            .handshake_complete = false,
        };
    }
    
    /// Perform XX pattern handshake (initiator side)
    pub fn xxHandshakeInitiator(self: *NoiseHandshake, allocator: std.mem.Allocator) ![]u8 {
        // -> e
        try self.noise.generateEphemeral();
        var msg1: [32]u8 = undefined;
        _ = self.noise.writeE(&msg1);
        
        // <- e, ee, s, es
        // (Responder sends back - would be received here)
        
        // -> s, se
        // (Final message from initiator)
        
        // For now, return first message
        return try allocator.dupe(u8, &msg1);
    }
    
    /// Wrap transport data with Noise encryption + MIMIC camouflage
    pub fn wrapTransport(
        self: *NoiseHandshake,
        allocator: std.mem.Allocator,
        plaintext: []const u8,
    ) ![]u8 {
        // Encrypt with Noise
        var ciphertext: [4096]u8 = undefined;
        const ct_len = try self.noise.writeMessage(plaintext, &ciphertext);
        
        // Apply MIMIC skin camouflage
        const skinned = try self.applySkin(allocator, ciphertext[0..ct_len]);
        
        return skinned;
    }
    
    fn applySkin(self: *NoiseHandshake, allocator: std.mem.Allocator, data: []const u8) ![]u8 {
        return switch (self.skin) {
            .Raw => try allocator.dupe(u8, data),
            .MimicHttps => try self.mimicHttps(allocator, data),
            .MimicDns => try self.mimicDns(allocator, data),
            .MimicQuic => try self.mimicQuic(allocator, data),
        };
    }
    
    fn mimicHttps(self: *NoiseHandshake, allocator: std.mem.Allocator, data: []const u8) ![]u8 {
        // Wrap in WebSocket frame with TLS-like padding
        _ = self;
        
        // Build fake TLS record layer
        var result = try allocator.alloc(u8, 5 + data.len);
        result[0] = 0x17; // Application Data
        result[1] = 0x03; // TLS 1.2
        result[2] = 0x03;
        std.mem.writeInt(u16, result[3..5], @intCast(data.len), .big);
        @memcpy(result[5..], data);
        
        return result;
    }
    
    fn mimicDns(_: *NoiseHandshake, allocator: std.mem.Allocator, data: []const u8) ![]u8 {
        // Encode as DNS TXT record format
        var result = try allocator.alloc(u8, data.len + 1);
        result[0] = @intCast(data.len);
        @memcpy(result[1..], data);
        return result;
    }
    
    fn mimicQuic(_: *NoiseHandshake, allocator: std.mem.Allocator, data: []const u8) ![]u8 {
        // QUIC Short Header format (simplified)
        var result = try allocator.alloc(u8, 1 + data.len);
        result[0] = 0x40; // Short header, 1-byte CID
        @memcpy(result[1..], data);
        return result;
    }
};

// ============================================================================
// TESTS
// ============================================================================

test "NoiseState initialization" {
    const state = NoiseState.init(.Noise_XX, .Initiator, &[_]u8{}, null, null);
    try std.testing.expectEqual(@as(u64, 0), state.c1.nonce);
    try std.testing.expectEqual(@as(u64, 0), state.c2.nonce);
}

test "CipherState encrypt/decrypt roundtrip" {
    _ = std.testing.allocator;
    
    var cipher = CipherState{
        .key = [_]u8{0xAB} ** 32,
        .nonce = 0,
    };
    
    const plaintext = "Hello, Noise!";
    var ciphertext: [100]u8 = undefined;
    
    const ct_len = try cipher.encryptWithAd(&[_]u8{}, plaintext, &ciphertext);
    try std.testing.expect(ct_len > plaintext.len); // Includes tag
}

test "NoiseHandshake with MIMIC skin" {
    _ = std.testing.allocator;
    
    const handshake = try NoiseHandshake.initWithSkin(
        .Noise_XX,
        .Initiator,
        .MimicHttps,
        null,
        null,
    );
    _ = handshake;
}