import {
  generateSalt,
  deriveWrappingKey,
  encryptKeyBundle,
  decryptKeyBundle,
  bufToBase64url,
  base64urlToBuf,
} from './crypto.js';
import { getAllKeys, storeDocKey } from './local-store.js';

export interface KeyBundle {
  [docId: string]: string;
}

export interface WrappedKeyStore {
  v: 2;
  salt: string;
  data: string;
}

const STORAGE_KEY = 'atmos-keys';

let _wrappingKey: CryptoKey | null = null;
let _salt: Uint8Array | null = null;

export function hasWrappingKey(): boolean {
  return _wrappingKey !== null;
}

export async function initWrappingKey(passphrase: string, salt?: Uint8Array): Promise<void> {
  const existingSalt = getStoredSalt();
  _salt = salt ?? existingSalt ?? generateSalt();
  _wrappingKey = await deriveWrappingKey(passphrase, _salt);
}

export function clearWrappingKey(): void {
  _wrappingKey = null;
  _salt = null;
}

function getStoredSalt(): Uint8Array | null {
  try {
    const raw = localStorage.getItem(STORAGE_KEY);
    if (!raw) return null;
    const parsed = JSON.parse(raw);
    if (parsed?.v === 2 && parsed.salt) {
      return base64urlToBuf(parsed.salt);
    }
  } catch { /* ignore */ }
  return null;
}

export function isWrappedFormat(): boolean {
  try {
    const raw = localStorage.getItem(STORAGE_KEY);
    if (!raw) return false;
    const parsed = JSON.parse(raw);
    return parsed?.v === 2;
  } catch {
    return false;
  }
}

export async function unwrapToPlaintext(): Promise<void> {
  if (!_wrappingKey) throw new Error('No wrapping key');
  const keys = await getLocalKeys();
  if (Object.keys(keys).length === 0) return;
  const savedKey = _wrappingKey;
  const savedSalt = _salt;
  _wrappingKey = null;
  _salt = null;
  await setLocalKeys(keys);
  _wrappingKey = savedKey;
  _salt = savedSalt;
}

export function isLegacyFormat(): boolean {
  try {
    const raw = localStorage.getItem(STORAGE_KEY);
    if (!raw) return false;
    const parsed = JSON.parse(raw);
    return parsed?.v !== 2;
  } catch {
    return false;
  }
}

export function getLegacyKeys(): KeyBundle {
  try {
    const raw = localStorage.getItem(STORAGE_KEY);
    if (!raw) return {};
    const parsed = JSON.parse(raw);
    if (parsed?.v === 2) return {};
    return parsed;
  } catch {
    return {};
  }
}

export async function getLocalKeys(): Promise<KeyBundle> {
  try {
    const raw = localStorage.getItem(STORAGE_KEY);
    if (!raw) return {};
    const parsed = JSON.parse(raw);

    if (parsed?.v === 2) {
      if (!_wrappingKey) return {};
      return await decryptKeyBundle(parsed.data, _wrappingKey);
    }

    return parsed;
  } catch {
    return {};
  }
}

export async function setLocalKeys(keys: KeyBundle): Promise<void> {
  if (!_wrappingKey || !_salt) {
    localStorage.setItem(STORAGE_KEY, JSON.stringify(keys));
    return;
  }

  const data = await encryptKeyBundle(keys, _wrappingKey);
  const store: WrappedKeyStore = {
    v: 2,
    salt: bufToBase64url(_salt),
    data,
  };
  localStorage.setItem(STORAGE_KEY, JSON.stringify(store));
}

export async function storeKey(docId: string, keyStr: string): Promise<void> {
  const keys = await getLocalKeys();
  keys[docId] = keyStr;
  await setLocalKeys(keys);
}

export function mergeKeys(local: KeyBundle, other: KeyBundle): KeyBundle {
  return { ...other, ...local };
}

export async function migrateLegacyKeys(): Promise<void> {
  const legacy = getLegacyKeys();
  if (Object.keys(legacy).length === 0) return;
  await setLocalKeys(legacy);
}

export async function syncKeys(): Promise<KeyBundle> {
  return getLocalKeys();
}

export async function exportKeyBundle(): Promise<Blob> {
  const idbKeys = await getAllKeys();
  const bundle: KeyBundle = {};

  for (const [docId, key] of idbKeys) {
    const raw = await crypto.subtle.exportKey('raw', key);
    const bytes = new Uint8Array(raw);
    let b64 = '';
    for (let i = 0; i < bytes.length; i++) {
      b64 += String.fromCharCode(bytes[i]);
    }
    bundle[docId] = btoa(b64);
  }

  const lsKeys = await getLocalKeys();
  for (const [docId, keyStr] of Object.entries(lsKeys)) {
    if (!bundle[docId]) {
      bundle[docId] = keyStr;
    }
  }

  return new Blob([JSON.stringify(bundle)], { type: 'application/json' });
}

export async function importKeyBundle(json: string): Promise<number> {
  let bundle: KeyBundle;
  try {
    bundle = JSON.parse(json);
  } catch {
    return 0;
  }

  if (typeof bundle !== 'object' || bundle === null) return 0;

  let imported = 0;
  const lsKeys = await getLocalKeys();

  for (const [docId, keyStr] of Object.entries(bundle)) {
    if (typeof keyStr !== 'string') continue;

    try {
      const binary = atob(keyStr);
      const bytes = new Uint8Array(binary.length);
      for (let i = 0; i < binary.length; i++) {
        bytes[i] = binary.charCodeAt(i);
      }
      const cryptoKey = await crypto.subtle.importKey(
        'raw',
        bytes.buffer,
        { name: 'AES-GCM', length: 256 },
        false,
        ['encrypt', 'decrypt'],
      );
      await storeDocKey(docId, cryptoKey);
      lsKeys[docId] = keyStr;
      imported++;
    } catch {
      // skip invalid key entries
    }
  }

  await setLocalKeys(lsKeys);
  return imported;
}