import { Encryption } from './index';
import { ShamirSecretSharing } from "./index";
import { KeyManagementSystemHandler } from './index';
import { FileSharder } from './index';
import * as readline from 'readline';
import fs from 'fs';
import yargs from 'yargs';
import { hideBin } from 'yargs/helpers';

const sss = new ShamirSecretSharing();

const argv = yargs(hideBin(process.argv))
    .option('encrypt', {
        alias: 'e',
        type: 'string',
        description: 'Text or file to encrypt'
    })
    .option('decrypt', {
        alias: 'd',
        type: 'string',
        description: 'File to decrypt'
    })
    .option('file', {
        alias: 'f',
        type: 'boolean',
        description: 'Specify if the input is a file'
    })
    .option('keyfolder', {
        alias: 'k',
        type: 'string',
        description: 'Folder to save the encrypted key'
    })
    .option('dataefile', {
        alias: 'ef',
        type: 'string',
        description: 'File to save the encrypted data'
    })
    .option('datadfile', {
        alias: 'df',
        type: 'string',
        description: 'File to save the decrypted data'
    })
    .option('algorithm', {
        alias: 'a',
        type: 'string',
        description: 'Encryption algorithm to use (aes-128-cbc, rsa)' 
    })
    .option('mode', {
        alias: 'm',
        type: 'string',
        description: 'Mode of operation to use (ecb, cbc, cfb, ofb, gcm)', 
        choices: ['ecb', 'cbc', 'cfb', 'ofb', 'gcm'],
        default: 'cbc'
    })
    .option('keysize', {
        alias: 'ks',
        type: 'number',
        description: 'Key size to use (128)', 
        choices: [128],
        default: 128
    })
    .option('rotatekey', {
        alias: 'r',
        type: 'boolean',
        description: 'Rotate the encryption key'
    })
    .option('hybridEncrypt', {
        alias: 'he',
        type: 'string',
        description: 'Text to hybrid encrypt with public key'
    })
    .option('hybridDecrypt', {
        alias: 'hd',
        type: 'string',
        description: 'Text to hybrid decrypt with private key'
    })
    .option('sign', {
        alias: 's',
        type: 'string',
        description: 'Text to sign with private key'
    })
    .option('verify', {
        alias: 'v',
        type: 'string',
        description: 'Text to verify signature with public key'
    })
    .option('generateDhKeys', {
        alias: 'gk',
        type: 'boolean',
        description: 'Generate Diffie-Hellman keys'
    })
    .option('computeSecret', {
        alias: 'cs',
        type: 'string',
        description: 'Compute shared secret with public key'
    })
    .option('generateOtp', {
        alias: 'otp',
        type: 'string',
        description: 'Generate one-time password with a secret'
    })
    .option('validateOtp', {
        alias: 'votp',
        type: 'string',
        description: 'Validate one-time password with a secret and otp'
    })
    .option('createBackup', {
        alias: 'cb',
        type: 'string',
        description: 'Create encrypted backup of data'
    })
    .option('recoverBackup', {
        alias: 'rb',
        type: 'string',
        description: 'Recover encrypted backup of data'
    })
    .option('addMetadata', {
        alias: 'am',
        type: 'string',
        description: 'Add metadata to a file'
    })
    .option('getMetadata', {
        alias: 'gm',
        type: 'string',
        description: 'Get metadata from a file'
    })
    .option('addKeyExpiration', {
        alias: 'ake',
        type: 'number',
        description: 'Add expiration to a key in seconds'
    })
    .option('checkKeyExpiration', {
        alias: 'cke',
        type: 'boolean',
        description: 'Check if the key has expired'
    })
    .option('encryptTwofish', {
        alias: 'etf',
        type: 'string',
        description: 'Text to encrypt with Twofish'
    })
    .option('decryptTwofish', {
        alias: 'dtf',
        type: 'string',
        description: 'Text to decrypt with Twofish'
    })
    .option('encryptChaCha20', {
        alias: 'ecc',
        type: 'string',
        description: 'Text to encrypt with ChaCha20'
    })
    .option('decryptChaCha20', {
        alias: 'dcc',
        type: 'string',
        description: 'Text to decrypt with ChaCha20'
    })
    .option("generateShares", {
        type: "boolean",
        description: "Generate shares from a secret",
      })
      .option("reconstructSecret", {
        type: "boolean",
        description: "Reconstruct the secret from shares",
      })
    .help()
    .argv;

// Create da instance
const she = new Encryption(argv.algorithm, argv.mode, argv.keysize, argv.dataefile, argv.datadfile);

if (argv.keyfolder) {
    she.setKeyFolder(argv.keyfolder);
}

// Save that boi
const saveKey = (key: Buffer | string, filename: string) => {
    she.saveKey(key, filename);
};

// Load that boi up
const loadKey = (filename: string): Buffer | string => {
    return she.loadKey(filename);
};

// Encryption and Decryption :D
const main = async () => {
    if (argv.rotatekey) {
        await she.rotateKey();
        console.log('Encryption key rotated');
    } else if (argv.encrypt) {
        try {
            const salt = she.generateIv(); 
            const key = she.deriveKeyFromPassword('', salt); 
            const iv = she.generateIv();
            
            saveKey(key, 'key.bin');
            saveKey(iv, 'iv.bin');
            saveKey(salt, 'salt.bin');
            
            let plaintext;
            if (argv.file) {
                plaintext = fs.readFileSync(argv.encrypt, 'utf8');
            } else {
                plaintext = argv.encrypt;
            }
            
            const encryptedText = she.encrypt(plaintext, key, iv); 
            console.log(`Encrypted Text: ${encryptedText}`);
        } catch (error) {
            console.error('An error occurred:', error);
        }
    } else if (argv.decrypt) {
        try {
            const encryptedText = fs.readFileSync(argv.decrypt, 'utf8');
            const loadedIv = loadKey('iv.bin') as Buffer;
            const loadedSalt = loadKey('salt.bin') as Buffer;
            
            const derivedKey = she.deriveKeyFromPassword('', loadedSalt); 
            const decryptedText = she.decrypt(encryptedText, derivedKey, loadedIv); 
            console.log(`Decrypted Text: ${decryptedText}`);
        } catch (error) {
            console.error('An error occurred:', error);
        }
    } else if (argv.hybridEncrypt) {
        try {
            const publicKey = fs.readFileSync('publicKey.pem', 'utf8');
            const hybridEncryptedText = she.hybridEncrypt(argv.hybridEncrypt, publicKey);
            console.log(`Hybrid Encrypted Text: ${hybridEncryptedText}`);
        } catch (error) {
            console.error('An error occurred during hybrid encryption:', error);
        }
    } else if (argv.hybridDecrypt) {
        try {
            const privateKey = fs.readFileSync('privateKey.pem', 'utf8');
            const hybridDecryptedText = she.hybridDecrypt(argv.hybridDecrypt, privateKey);
            console.log(`Hybrid Decrypted Text: ${hybridDecryptedText}`);
        } catch (error) {
            console.error('An error occurred during hybrid decryption:', error);
        }
    } else if (argv.sign) {
        try {
            const privateKey = fs.readFileSync('privateKey.pem', 'utf8');
            const signature = she.signData(argv.sign, privateKey);
            console.log(`Signature: ${signature}`);
        } catch (error) {
            console.error('An error occurred during signing:', error);
        }
    } else if (argv.verify) {
        try {
            const publicKey = fs.readFileSync('publicKey.pem', 'utf8');
            const isValid = she.verifySignature(argv.verify, argv._[1], publicKey);
            console.log(`Signature is valid: ${isValid}`);
        } catch (error) {
            console.error('An error occurred during signature verification:', error);
        }
    } else if (argv.generateDhKeys) {
        try {
            const { publicKey, privateKey } = she.generateDhKeys();
            console.log(`Diffie-Hellman Public Key: ${publicKey}`);
            console.log(`Diffie-Hellman Private Key: ${privateKey}`);
        } catch (error) {
            console.error('An error occurred during Diffie-Hellman key generation:', error);
        }
    } else if (argv.computeSecret) {
        try {
            const privateKey = fs.readFileSync('dhPrivateKey.pem', 'utf8');
            const secret = she.computeSecret(argv.computeSecret, privateKey);
            console.log(`Computed Shared Secret: ${secret.toString('hex')}`);
        } catch (error) {
            console.error('An error occurred during computing the shared secret:', error);
        }
    } else if (argv.generateOtp) {
        try {
            const otp = she.generateOtp(argv.generateOtp);
            console.log(`Generated OTP: ${otp}`);
        } catch (error) {
            console.error('An error occurred during OTP generation:', error);
        }
    } else if (argv.validateOtp) {
        try {
            const isValid = she.validateOtp(argv.validateOtp, argv._[1]);
            console.log(`OTP is valid: ${isValid}`);
        } catch (error) {
            console.error('An error occurred during OTP validation:', error);
        }
    } else if (argv.createBackup) {
        try {
            const key = loadKey('backupKey.bin') as Buffer;
            const iv = loadKey('backupIv.bin') as Buffer;
            she.createEncryptedBackup(argv.createBackup, key, iv, 'backup.enc');
            console.log('Encrypted backup created successfully.');
        } catch (error) {
            console.error('An error occurred during creating an encrypted backup:', error);
        }
    } else if (argv.recoverBackup) {
        try {
            const key = loadKey('backupKey.bin') as Buffer;
            const iv = loadKey('backupIv.bin') as Buffer;
            const recoveredData = she.recoverEncryptedBackup('backup.enc', key, iv);
            console.log(`Recovered Data: ${recoveredData}`);
        } catch (error) {
            console.error('An error occurred during recovering the encrypted backup:', error);
        }
    } else if (argv.addMetadata) {
        try {
            const metadata = JSON.parse(argv._[1]);
            she.addMetadata(argv.addMetadata, metadata);
            console.log('Metadata added successfully.');
        } catch (error) {
            console.error('An error occurred during adding metadata:', error);
        }
    } else if (argv.getMetadata) {
        try {
            const metadata = she.getMetadata(argv.getMetadata);
            console.log(`Metadata: ${JSON.stringify(metadata, null, 2)}`);
        } catch (error) {
            console.error('An error occurred during fetching metadata:', error);
        }
    } else if (argv.addKeyExpiration) {
        try {
            const key = loadKey('key.bin') as Buffer;
            she.addKeyExpiration(key, argv.addKeyExpiration);
            console.log('Key expiration added successfully.');
        } catch (error) {
            console.error('An error occurred during adding key expiration:', error);
        }
    } else if (argv.checkKeyExpiration) {
        try {
            const isExpired = she.checkKeyExpiration();
            console.log(`Key has expired: ${isExpired}`);
        } catch (error) {
            console.error('An error occurred during checking key expiration:', error);
        }
    } else if (argv.encryptTwofish) {
        try {
            const key = she.generateKey();
            const iv = she.generateIv();
            const encryptedText = she.encryptTwofish(argv.encryptTwofish, key, iv);
            saveKey(key, 'twofishKey.bin');
            saveKey(iv, 'twofishIv.bin');
            console.log(`Twofish Encrypted Text: ${encryptedText}`);
        } catch (error) {
            console.error('An error occurred during Twofish encryption:', error);
        }
    } else if (argv.decryptTwofish) {
        try {
            const key = loadKey('twofishKey.bin') as Buffer;
            const iv = loadKey('twofishIv.bin') as Buffer;
            const decryptedText = she.decryptTwofish(argv.decryptTwofish, key, iv);
            console.log(`Twofish Decrypted Text: ${decryptedText}`);
        } catch (error) {
            console.error('An error occurred during Twofish decryption:', error);
        }
    } else if (argv.encryptChaCha20) {
        try {
            const key = she.generateKey();
            const nonce = she.generateIv(); // ChaCha20
            const encryptedText = she.encryptChaCha20(argv.encryptChaCha20, key, nonce);
            saveKey(key, 'chaCha20Key.bin');
            saveKey(nonce, 'chaCha20Nonce.bin');
            console.log(`ChaCha20 Encrypted Text: ${encryptedText}`);
        } catch (error) {
            console.error('An error occurred during ChaCha20 encryption:', error);
        }
    } else if (argv.decryptChaCha20) {
        try {
            const key = loadKey('chaCha20Key.bin') as Buffer;
            const nonce = loadKey('chaCha20Nonce.bin') as Buffer;
            const decryptedText = she.decryptChaCha20(argv.decryptChaCha20, key, nonce);
            console.log(`ChaCha20 Decrypted Text: ${decryptedText}`);
        } catch (error) {
            console.error('An error occurred during ChaCha20 decryption:', error);
        }
    } else {
        console.log('Please specify an action: --encrypt, --decrypt, --hybridEncrypt, --hybridDecrypt, --sign, --verify, --generateDhKeys, --computeSecret, --generateOtp, --validateOtp, --createBackup, --recoverBackup, --addMetadata, --getMetadata, --addKeyExpiration, --checkKeyExpiration, --encryptTwofish, --decryptTwofish, --encryptChaCha20, or --decryptChaCha20');
    }
};

const shamir = require('shamir');

// Generate shares
function generateShares(secret, numShares, threshold) {
    try {
        const shares = shamir.split(secret, numShares, threshold);
        console.log(`Generated shares: ${shares}`);
        return shares;
    } catch (error) {
        console.error('An error occurred during generating shares:', error);
    }
}

// Reconstruct secret
function reconstructSecrets(shares) {
    try {
        const secret = shamir.combine(shares);
        console.log(`Reconstructed secret: ${secret}`);
        return secret;
    } catch (error) {
        console.error('An error occurred during reconstructing secret:', error);
    }
}


const keyManager = new KeyManagementSystemHandler();

// Example usage
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);
}

const prompt = async (question: string): Promise<string> => {
    const rl = readline.createInterface({
        input: process.stdin,
        output: process.stdout,
    });
    return new Promise((resolve) => rl.question(question, (answer) => {
        rl.close();
        resolve(answer);
    }));
};

(async () => {
    try {
        const fileSharder = new FileSharder();

        // Prompt user for input file
        const inputFilePath = await prompt("Enter the path to the input file: ");
        if (!fs.existsSync(inputFilePath)) {
            throw new Error(`File not found: ${inputFilePath}`);
        }

        // Prompt user for directory to save shards
        let outputDir = await prompt("Enter the directory to save shards (or leave blank for default './shardedFiles'): ");
        outputDir = outputDir || './shardedFiles'; // Default directory
        if (!fs.existsSync(outputDir)) {
            fs.mkdirSync(outputDir); // Automatically create the directory if it doesn't exist
            console.log(`Created directory: ${outputDir}`);
        }

        // Prompt user for output file to save reassembled data
        let reassembledFilePath = await prompt("Enter the path for the reassembled file (or leave blank for default 'reconstructedFile.txt'): ");
        reassembledFilePath = reassembledFilePath || 'reconstructedFile.txt'; // Default file

        // Step 1: Split and encrypt the file
        console.log('Starting file sharding and encryption...');
        fileSharder.shardAndEncrypt(inputFilePath, outputDir);

        // Step 2: Reassemble and decrypt the file
        console.log('Starting reassembly and decryption...');
        fileSharder.reassembleAndDecrypt(outputDir, reassembledFilePath);

        console.log(`File processing complete! Check shards in ${outputDir} and the reassembled file at ${reassembledFilePath}.`);
    } catch (error) {
        console.error('An error occurred:', error.message);
    }
})();

main();