import crypto from 'crypto';
import fs from 'fs';
import path from 'path';
import zlib from 'zlib';
import * as shamir from 'shamir';
import config from './config.json';

export class Encryption {
    private algorithm: string;
    private keyLength: number;
    private ivLength: number;
    private rsaKeySize: number;
    private keyFolder: string;
    private dataFile: string;
    private dataDecryptedFile: string;

    constructor(algorithm: string = config.defaultAlgorithm, mode: string = 'cbc', keyLength: number = 128, dataFile: string, dataDecryptedFile: string) {
        this.algorithm = `${algorithm}-${mode}`;
        this.keyLength = keyLength / 8; 
        this.ivLength = 16; 
        this.rsaKeySize = config.rsaKeySize;
        this.keyFolder = __dirname; 
        this.dataFile = dataFile;
        this.dataDecryptedFile = dataDecryptedFile;
    }

    setKeyFolder(folder: string): void {
        this.keyFolder = folder;
    }

    setDataFile(file: string): void {
        this.dataFile = file;
    }

    setDataDecryptedFile(file: string): void {
        this.dataDecryptedFile = file;
    }

    generateKey(): Buffer {
        return crypto.randomBytes(this.keyLength);
    }

    generateIv(): Buffer {
        return crypto.randomBytes(this.ivLength);
    }

    generateRsaKeyPair(): { publicKey: string, privateKey: string } {
        const { publicKey, privateKey } = crypto.generateKeyPairSync('rsa', {
            modulusLength: this.rsaKeySize,
            publicKeyEncoding: {
                type: 'pkcs1',
                format: 'pem'
            },
            privateKeyEncoding: {
                type: 'pkcs1',
                format: 'pem'
            }
        });
        return { publicKey, privateKey };
    }

    deriveKeyFromPassword(password: string, salt: Buffer): Buffer {
        return crypto.pbkdf2Sync(password, salt, 100000, this.keyLength, 'sha256');
    }

    saveKey(key: Buffer | string, filename: string): void {
        // Store that boi
        fs.writeFileSync(path.join(this.keyFolder, filename), key);
    }

    loadKey(filename: string): Buffer | string {
        // Securly load that boi
        return fs.readFileSync(path.join(this.keyFolder, filename));
    }

    hashData(data: string): string {
        return crypto.createHash('sha256').update(data).digest('hex');
    }

    compressData(data: string): Buffer {
        return zlib.gzipSync(data);
    }

    decompressData(data: Buffer): string {
        return zlib.gunzipSync(data).toString();
    }

    logOperation(operation: string, details: string): void {
        const logMessage = `${new Date().toISOString()} - ${operation}: ${details}\n`;
        fs.appendFileSync('encryption.log', logMessage);
    }

    
    encrypt(text: string, key: Buffer, iv: Buffer): string {
        try {
            const compressed = this.compressData(text);
            const cipher = crypto.createCipheriv(this.algorithm, key, iv); 
            let encrypted = cipher.update(compressed);
            encrypted = Buffer.concat([encrypted, cipher.final()]);
            fs.writeFileSync(this.dataFile, encrypted); // Save that shii to a file
            this.logOperation('Encrypt', `Data encrypted and saved to ${this.dataFile}`);
            return encrypted.toString('hex');
        } catch (error) {
            this.logOperation('Encrypt', `Encryption failed: ${error.message}`);
            throw new Error('Encryption failed: ' + error.message);
        }
    }


    decrypt(encryptedText: string, key: Buffer, iv: Buffer): string {
        try {
            const encryptedBuffer = Buffer.from(encryptedText, 'hex');
            const decipher = crypto.createDecipheriv(this.algorithm, key, iv); 
            let decrypted = decipher.update(encryptedBuffer);
            decrypted = Buffer.concat([decrypted, decipher.final()]);
            const decompressed = this.decompressData(decrypted);
            fs.writeFileSync(this.dataDecryptedFile, decompressed); // Save that decrypted shii to a file
            this.logOperation('Decrypt', `Data decrypted and saved to ${this.dataDecryptedFile}`);
            return decompressed;
        } catch (error) {
            this.logOperation('Decrypt', `Decryption failed: ${error.message}`);
            throw new Error('Decryption failed: ' + error.message);
        }
    }

    encryptWithPublicKey(text: string, publicKey: string): string {
        try {
            const buffer = Buffer.from(text, 'utf8');
            const encrypted = crypto.publicEncrypt(publicKey, buffer);
            fs.writeFileSync(this.dataFile, encrypted.toString('hex')); // Save that encrypted data shii to a file
            this.logOperation('EncryptWithPublicKey', `Data encrypted with public key and saved to ${this.dataFile}`);
            return encrypted.toString('hex');
        } catch (error) {
            this.logOperation('EncryptWithPublicKey', `Encryption with public key failed: ${error.message}`);
            throw new Error('Encryption with public key failed: ' + error.message);
        }
    }

    decryptWithPrivateKey(encryptedText: string, privateKey: string): string {
        try {
            const buffer = Buffer.from(encryptedText, 'hex');
            const decrypted = crypto.privateDecrypt(privateKey, buffer);
            fs.writeFileSync(this.dataDecryptedFile, decrypted.toString('utf8')); // Save that decrypted data shii to a file
            this.logOperation('DecryptWithPrivateKey', `Data decrypted with private key and saved to ${this.dataDecryptedFile}`);
            return decrypted.toString('utf8');
        } catch (error) {
            this.logOperation('DecryptWithPrivateKey', `Decryption with private key failed: ${error.message}`);
            throw new Error('Decryption with private key failed: ' + error.message);
        }
    }

    // Key Rotation 
    async rotateKey(): Promise<void> {
        const newKey = this.generateKey();
        this.saveKey(newKey, 'key.bin');
        this.logOperation('RotateKey', 'Encryption key rotated');
    }

    // *** New Features  as of version 1.01 :D ***


    hybridEncrypt(text: string, publicKey: string): string {
        const symmetricKey = this.generateKey();
        const iv = this.generateIv();
        const encryptedText = this.encrypt(text, symmetricKey, iv);
        const encryptedKey = crypto.publicEncrypt(publicKey, symmetricKey);
        return `${encryptedKey.toString('hex')}:${encryptedText}`;
    }

    hybridDecrypt(encryptedData: string, privateKey: string): string {
        const [encryptedKey, encryptedText] = encryptedData.split(':');
        const symmetricKey = crypto.privateDecrypt(privateKey, Buffer.from(encryptedKey, 'hex'));
        const iv = this.generateIv();
        return this.decrypt(encryptedText, symmetricKey, iv);
    }

   
    signData(data: string, privateKey: string): string {
        const sign = crypto.createSign('SHA256');
        sign.update(data);
        return sign.sign(privateKey, 'hex');
    }

    verifySignature(data: string, signature: string, publicKey: string): boolean {
        const verify = crypto.createVerify('SHA256');
        verify.update(data);
        return verify.verify(publicKey, signature, 'hex');
    }

  
    generateDhKeys(): { publicKey: string, privateKey: string } {
        const dh = crypto.createDiffieHellman(2048);
        const privateKey = dh.generateKeys('hex');
        const publicKey = dh.getPublicKey('hex');
        return { publicKey, privateKey };
    }

    computeSecret(publicKey: string, privateKey: string): Buffer {
        const dh = crypto.createDiffieHellman(Buffer.from(privateKey, 'hex'));
        dh.setPublicKey(Buffer.from(publicKey, 'hex'));
        return dh.computeSecret(Buffer.from(publicKey, 'hex'));
    }

    
    generateOtp(secret: string): string {
        const otp = crypto.createHmac('sha1', secret).update(String(Date.now())).digest('hex');
        return otp;
    }

    validateOtp(secret: string, otp: string): boolean {
        const generatedOtp = this.generateOtp(secret);
        return otp === generatedOtp;
    }


    createEncryptedBackup(data: string, key: Buffer, iv: Buffer, backupFile: string): void {
        const encrypted = this.encrypt(data, key, iv);
        fs.writeFileSync(backupFile, encrypted);
    }

    recoverEncryptedBackup(backupFile: string, key: Buffer, iv: Buffer): string {
        const encryptedData = fs.readFileSync(backupFile, 'utf8');
        return this.decrypt(encryptedData, key, iv);
    }


    addMetadata(filename: string, metadata: object): void {
        const filePath = path.join(this.keyFolder, `${filename}.meta`);
        fs.writeFileSync(filePath, JSON.stringify(metadata));
    }

    getMetadata(filename: string): object {
        const filePath = path.join(this.keyFolder, `${filename}.meta`);
        return JSON.parse(fs.readFileSync(filePath, 'utf8'));
    }


    addKeyExpiration(key: Buffer, expiresIn: number): void {
        const expirationDate = new Date(Date.now() + expiresIn * 1000).toISOString();
        fs.writeFileSync(path.join(this.keyFolder, 'key-expiration'), expirationDate);
    }

    checkKeyExpiration(): boolean {
        const expirationDate = fs.readFileSync(path.join(this.keyFolder, 'key-expiration'), 'utf8');
        return new Date(expirationDate) > new Date();
    }
    
    encryptTwofish(text: string, key: Buffer, iv: Buffer): string {
        const cipher = crypto.createCipheriv('twofish', key, iv);
        let encrypted = cipher.update(text, 'utf8', 'hex');
        encrypted += cipher.final('hex');
        return encrypted;
    }

    
    decryptTwofish(encryptedText: string, key: Buffer, iv: Buffer): string {
        const decipher = crypto.createDecipheriv('twofish', key, iv);
        let decrypted = decipher.update(encryptedText, 'hex', 'utf8');
        decrypted += decipher.final('utf8');
        return decrypted;
    }

    // *** sum additional Key Derivation Functions cuz why not yk ***

    deriveKeyFromPasswordArgon2(password: string, salt: Buffer): Buffer {
        const key = crypto.scryptSync(password, salt, this.keyLength);
        return key;
    }

    deriveKeyFromPasswordBcrypt(password: string, salt: Buffer): string {
        const bcrypt = require('bcrypt');
        const hash = bcrypt.hashSync(password, salt.toString('hex'));
        return hash;
    }
   
    encryptChaCha20(text: string, key: Buffer, nonce: Buffer): string {
        const cipher = crypto.createCipheriv('chacha20', key, nonce);
        let encrypted = cipher.update(text, 'utf8', 'hex');
        encrypted += cipher.final('hex');
        return encrypted;
    }

    
    decryptChaCha20(encryptedText: string, key: Buffer, nonce: Buffer): string {
        const decipher = crypto.createDecipheriv('chacha20', key, nonce);
        let decrypted = decipher.update(encryptedText, 'hex', 'utf8');
        decrypted += decipher.final('utf8');
        return decrypted;
    }
}


class KeyManagementSystem {
    private keyRegistry: Map<string, { key: Buffer; revoked: boolean }>;

    constructor() {
        this.keyRegistry = new Map();
    }

    // Save a key in the registry with a unique identifier.
    saveKey(keyId: string, key: Buffer): void {
        if (this.keyRegistry.has(keyId)) {
            throw new Error(`Key ID '${keyId}' already exists.`);
        }
        this.keyRegistry.set(keyId, { key, revoked: false });
        console.log(`Key '${keyId}' saved successfully.`);
    }

    // Revoke a key, marking it as unusable for further operations.
    revokeKey(keyId: string): void {
        if (!this.keyRegistry.has(keyId)) {
            throw new Error(`Key ID '${keyId}' does not exist.`);
        }

        const keyEntry = this.keyRegistry.get(keyId);
        if (keyEntry) {
            keyEntry.revoked = true;
        }

        console.log(`Key '${keyId}' has been revoked.`);
    }

    // Check if a key is active (not revoked).
    isKeyActive(keyId: string): boolean {
        const keyEntry = this.keyRegistry.get(keyId);
        if (!keyEntry) {
            throw new Error(`Key ID '${keyId}' does not exist.`);
        }
        return !keyEntry.revoked;
    }

    // Get a key from the registry if it is active.
    getKey(keyId: string): Buffer {
        if (!this.isKeyActive(keyId)) {
            throw new Error(`Key '${keyId}' is revoked and cannot be used.`);
        }

        const keyEntry = this.keyRegistry.get(keyId);
        if (keyEntry) {
            return keyEntry.key;
        }

        throw new Error(`Key '${keyId}' does not exist in the registry.`);
    }
}

export class KeyManagementSystemHandler {
    private keyManager: KeyManagementSystem;

    constructor() {
        this.keyManager = new KeyManagementSystem();
    }

    // Delegate saveKey to the keyManager
    saveKey(keyId: string, key: Buffer): void {
        this.keyManager.saveKey(keyId, key);
    }

    // Delegate getKey to the keyManager
    getKey(keyId: string): Buffer {
        return this.keyManager.getKey(keyId);
    }

    // You can add more methods to handle key management as needed
}

// Example usage
const keyManager = new KeyManagementSystemHandler();
const secret = Buffer.from('mySecret');
const numShares = 5;
const threshold = 3;

try {
    // Generate shares
    const shares = shamir.split(secret, numShares, threshold);
    console.log(`Generated shares: ${shares}`);

    // Save shares in the key management system
    shares.forEach((share, index) => {
        const keyId = `share-${index + 1}`;
        keyManager.saveKey(keyId, share);
    });

    // Retrieve shares from the key management system
    const retrievedShares: Buffer[] = [];
    for (let i = 1; i <= threshold; i++) {
        const keyId = `share-${i}`;
        retrievedShares.push(keyManager.getKey(keyId));
    }

    // Reconstruct the secret
    const reconstructedSecret = shamir.combine(retrievedShares);
    console.log(`Reconstructed secret: ${reconstructedSecret.toString()}`);
} catch (error) {
    console.error('An error occurred:', error);
}

export class ShamirSecretSharing {}




export class FileSharder {
    private chunkSize: number; // Size of each chunk in bytes
    private algorithm: string; // Encryption algorithm
    private key: Buffer; // Encryption key

    constructor(chunkSize: number = 1024 * 1024, algorithm: string = 'aes-256-cbc', key: string = 'defaultencryptionkeydefault') {
        this.chunkSize = chunkSize; // Default: 1 MB chunks
        this.algorithm = algorithm;
        this.key = Buffer.from(key.padEnd(32, '0')); // Ensure key is 32 bytes for aes-256-cbc
    }

    shardAndEncrypt(filePath: string, outputDir: string): void {
        if (!fs.existsSync(filePath)) {
            throw new Error(`File not found: ${filePath}`);
        }

        if (!fs.existsSync(outputDir)) {
            fs.mkdirSync(outputDir); // Create output directory if it doesn't exist
            console.log(`Created output directory: ${outputDir}`);
        }

        const fileStream = fs.createReadStream(filePath, { highWaterMark: this.chunkSize });
        let chunkIndex = 0;

        fileStream.on('data', (chunk) => {
            if (Buffer.isBuffer(chunk)) {
                const encryptedChunk = this.encrypt(chunk);
                const chunkFilePath = `${outputDir}/chunk-${chunkIndex}.enc`;
        
                fs.writeFileSync(chunkFilePath, encryptedChunk); // Save encrypted chunk to file
                console.log(`Saved encrypted chunk ${chunkIndex} to ${chunkFilePath}`);
        
                chunkIndex++;
            } else {
                throw new Error('Unexpected string data encountered in file stream.');
            }
        });

        fileStream.on('end', () => {
            console.log('File sharding and encryption complete.');
        });

        fileStream.on('error', (err) => {
            console.error('An error occurred during file processing:', err);
        });
    }

    reassembleAndDecrypt(chunkDir: string, outputFile: string): void {
        if (!fs.existsSync(chunkDir)) {
            throw new Error(`Chunk directory not found: ${chunkDir}`);
        }

        const chunkFiles = fs.readdirSync(chunkDir).filter(file => file.endsWith('.enc'));
        chunkFiles.sort(); // Ensure chunks are processed in the correct order

        const writeStream = fs.createWriteStream(outputFile);

        chunkFiles.forEach(chunkFile => {
            const chunkFilePath = `${chunkDir}/${chunkFile}`;
            const encryptedChunk = fs.readFileSync(chunkFilePath);

            const decryptedChunk = this.decrypt(encryptedChunk);
            writeStream.write(decryptedChunk);

            console.log(`Processed chunk: ${chunkFilePath}`);
        });

        writeStream.end();
        console.log(`Reassembled and decrypted file saved to: ${outputFile}`);
    }

    private encrypt(data: Buffer): Buffer {
        const iv = crypto.randomBytes(16); // Initialization vector
        const cipher = crypto.createCipheriv(this.algorithm, this.key, iv);
        const encrypted = Buffer.concat([iv, cipher.update(data), cipher.final()]); // Prepend IV
        return encrypted;
    }

    private decrypt(encryptedData: Buffer): Buffer {
        const iv = encryptedData.slice(0, 16); // Extract IV
        const encryptedContent = encryptedData.slice(16); // Extract content
        const decipher = crypto.createDecipheriv(this.algorithm, this.key, iv);
        const decrypted = Buffer.concat([decipher.update(encryptedContent), decipher.final()]);
        return decrypted;
    }
}