import crypto from 'crypto';
import fs from 'fs/promises';
import path from 'path';

const JSON_FORMAT_VERSION = 2;
const BINARY_FORMAT_VERSION = 1;
const ALGORITHM = 'aes-256-gcm';
const KEY_LENGTH = 32;
const IV_LENGTH = 12;
const AUTH_TAG_LENGTH = 16;

function getKey() {
  const keyEnv = process.env.DATA_ENCRYPTION_KEY;
  if (!keyEnv) throw new Error('DATA_ENCRYPTION_KEY environment variable not set');
  return crypto.createHash('sha256').update(keyEnv).digest().subarray(0, KEY_LENGTH);
}

function normalizeAad(aad = '') {
  if (Buffer.isBuffer(aad)) return aad;
  return Buffer.from(String(aad || ''), 'utf8');
}

export function encryptBuffer(buffer, aad = '') {
  const key = getKey();
  const iv = crypto.randomBytes(IV_LENGTH);
  const cipher = crypto.createCipheriv(ALGORITHM, key, iv);
  const aadBuffer = normalizeAad(aad);
  if (aadBuffer.length) cipher.setAAD(aadBuffer);
  const encrypted = Buffer.concat([cipher.update(buffer), cipher.final()]);
  const authTag = cipher.getAuthTag();
  return {
    version: BINARY_FORMAT_VERSION,
    iv,
    authTag,
    encrypted,
  };
}

export function decryptBuffer(payload, aad = '') {
  const key = getKey();
  const iv = Buffer.isBuffer(payload?.iv) ? payload.iv : Buffer.from(payload?.iv || '', 'hex');
  const authTag = Buffer.isBuffer(payload?.authTag)
    ? payload.authTag
    : Buffer.from(payload?.authTag || '', 'hex');
  const encrypted = Buffer.isBuffer(payload?.encrypted)
    ? payload.encrypted
    : Buffer.from(payload?.encrypted || '', 'hex');

  const decipher = crypto.createDecipheriv(ALGORITHM, key, iv);
  const aadBuffer = normalizeAad(aad);
  if (aadBuffer.length) decipher.setAAD(aadBuffer);
  decipher.setAuthTag(authTag);
  return Buffer.concat([decipher.update(encrypted), decipher.final()]);
}

export function packEncryptedBuffer(payload) {
  const header = Buffer.allocUnsafe(1 + 1 + payload.iv.length + payload.authTag.length);
  header.writeUInt8(payload.version || BINARY_FORMAT_VERSION, 0);
  header.writeUInt8(payload.iv.length, 1);
  payload.iv.copy(header, 2);
  payload.authTag.copy(header, 2 + payload.iv.length);
  return Buffer.concat([header, payload.encrypted]);
}

export function unpackEncryptedBuffer(buffer) {
  const version = buffer.readUInt8(0);
  if (version !== BINARY_FORMAT_VERSION) {
    throw new Error(`Unsupported encrypted buffer version: ${version}`);
  }
  const ivLength = buffer.readUInt8(1);
  const ivStart = 2;
  const ivEnd = ivStart + ivLength;
  const tagEnd = ivEnd + AUTH_TAG_LENGTH;
  return {
    version,
    iv: buffer.subarray(ivStart, ivEnd),
    authTag: buffer.subarray(ivEnd, tagEnd),
    encrypted: buffer.subarray(tagEnd),
  };
}

export async function writeEncryptedFile(filePath, buffer, aad = '') {
  await fs.mkdir(path.dirname(filePath), { recursive: true });
  const payload = encryptBuffer(buffer, aad);
  await fs.writeFile(filePath, packEncryptedBuffer(payload));
}

export async function readEncryptedFile(filePath, aad = '') {
  const packed = await fs.readFile(filePath);
  return decryptBuffer(unpackEncryptedBuffer(packed), aad);
}

function legacyDecryptJson(encryptedData) {
  const key = getKey();
  const decipher = crypto.createDecipher(ALGORITHM, key);
  decipher.setAuthTag(Buffer.from(encryptedData.authTag, 'hex'));
  let decrypted = decipher.update(encryptedData.encrypted, 'hex', 'utf8');
  decrypted += decipher.final('utf8');
  return JSON.parse(decrypted);
}

export function encryptJson(data, aad = '') {
  const payload = encryptBuffer(Buffer.from(JSON.stringify(data), 'utf8'), aad);
  return {
    version: JSON_FORMAT_VERSION,
    iv: payload.iv.toString('hex'),
    authTag: payload.authTag.toString('hex'),
    encrypted: payload.encrypted.toString('hex'),
  };
}

export function decryptJson(encryptedData, aad = '') {
  if (!encryptedData) return null;
  if ((encryptedData.version || 0) >= JSON_FORMAT_VERSION) {
    const decrypted = decryptBuffer(encryptedData, aad);
    return JSON.parse(decrypted.toString('utf8'));
  }
  return legacyDecryptJson(encryptedData);
}

export async function saveEncryptedJson(filePath, data, aad = '') {
  const encrypted = encryptJson(data, aad);
  await fs.mkdir(path.dirname(filePath), { recursive: true });
  await fs.writeFile(filePath, JSON.stringify(encrypted, null, 2), 'utf8');
}

export async function loadEncryptedJson(filePath, aad = '') {
  try {
    const content = await fs.readFile(filePath, 'utf8');
    const encrypted = JSON.parse(content);
    return decryptJson(encrypted, aad);
  } catch (err) {
    if (err.code === 'ENOENT') return null;
    throw err;
  }
}