// Third-party libraries
import { ecb } from '@noble/ciphers/aes.js';
import { ed25519, x25519 } from '@noble/curves/ed25519.js';
import { hmac } from '@noble/hashes/hmac.js';
import { sha256, sha512 } from '@noble/hashes/sha2.js';
import { bytesToHex, hexToBytes } from '@noble/hashes/utils.js';

// Local types
import type {
  AckPayload,
  AnonReqPayload,
  AdvertPayload,
  BasePrivateKey,
  BasePublicKey,
  BaseSharedSecret,
  ControlPayload,
  DecryptedGroupMessage,
  Group,
  GroupDataPayload,
  GroupSecretValue,
  GroupTextPayload,
  MultipartPayload,
  NodeHash,
  PathPayload,
  Payload,
  PublicKeyValue,
  RawCustomPayload,
  RequestPayload,
  ResponsePayload,
  TextPayload,
  TracePayload,
} from './meshcore.types';

// Local imports
import { base64ToBytes, BufferReader, constantTimeEqual } from '../util';
import {
  AdvertisementFlags,
  BaseGroupSecret,
  BasePacket,
  PayloadType,
  RouteType,
} from './meshcore.types';

const MAX_PATH_SIZE = 64;

export const isValidRouteType = (value: number): value is typeof RouteType[keyof typeof RouteType] => {
  return Object.values(RouteType).includes(value as any);
}

export const isValidPayloadType = (value: number): value is typeof PayloadType[keyof typeof PayloadType] => {
  return Object.values(PayloadType).includes(value as any);
}

export const hasTransportCodes = (routeType: RouteType): boolean => {
  return routeType === RouteType.TRANSPORT_FLOOD || routeType === RouteType.TRANSPORT_DIRECT;
}

export class Packet extends BasePacket {
  constructor(data?: Uint8Array | string) {
    super();
    if (typeof data !== 'undefined') {
      if (typeof data === 'string') {
        data = base64ToBytes(data);
      }
      this.parse(data);
    }
  }

  public parse(data: Uint8Array) {
    const header = data[0];
    this.routeType = (header >> 0) & 0x03;
    this.payloadType = (header >> 2) & 0x0F;
    this.version = (header >> 6) & 0x03;

    let index = 1;
    if (hasTransportCodes(this.routeType)) {
      const view = new DataView(data.buffer, index, index + 4);
      this.transportCodes = new Uint16Array(2);
      this.transportCodes[0] = view.getUint16(0, true);
      this.transportCodes[1] = view.getUint16(2, true);
      index += 4;
    }

    this.pathLength = data[index];
    index++;
    if (!this.isValidPathLength()) {
      throw new Error(`MeshCore: invalid path length ${this.pathLength}`)
    }

    const pathBytesLength = this.getPathBytesLength();
    this.path = new Uint8Array(pathBytesLength);
    for (let i = 0; i < pathBytesLength; i++) {
      this.path[i] = data[index];
      index++;
    }

    if (index >= data.length) {
      throw new Error('MeshCore: invalid packet: no payload');
    }

    const payloadBytesLength = data.length - index;
    this.payload = new Uint8Array(payloadBytesLength);
    for (let i = 0; i < payloadBytesLength; i++) {
      this.payload[i] = data[index];
      index++;
    }
  }

  public parseBase64(data: string) {
    return this.parse(base64ToBytes(data));
  }

  private isValidPathLength(): boolean {
    const hashCount = this.getPathHashCount();
    const hashSize = this.getPathHashSize();
    if (hashSize === 4) return false; // reserved
    return hashCount * hashSize <= MAX_PATH_SIZE;
  }

  public getPathHashSize(): number {
    return (this.pathLength >> 6) + 1;
  }

  public getPathHashCount(): number {
    return this.pathLength & 63;
  }

  public getPathBytesLength(): number {
    return this.getPathHashCount() * this.getPathHashSize();
  }

  public hash(): Uint8Array {
    let data = new Uint8Array([this.payloadType]);
    if (this.payloadType === PayloadType.TRACE) {
      data = new Uint8Array([...data, ...this.path]);
    }
    data = new Uint8Array([...data, ...this.payload]);
    const hash = sha256.create().update(data).digest();
    return hash.slice(0, 8);
  }

  public decode(): Payload {
    switch (this.payloadType) {
    case PayloadType.REQUEST:
      return this.decodeRequest();
    case PayloadType.RESPONSE:
      return this.decodeResponse();
    case PayloadType.TEXT:
      return this.decodeText();
    case PayloadType.ACK:
      return this.decodeAck();
    case PayloadType.ADVERT:
      return this.decodeAdvert();
    case PayloadType.GROUP_TEXT:
      return this.decodeGroupText();
    case PayloadType.GROUP_DATA:
      return this.decodeGroupData();
    case PayloadType.ANON_REQ:
      return this.decodeAnonReq();
    case PayloadType.PATH:
      return this.decodePath();
    case PayloadType.TRACE:
      return this.decodeTrace();
    case PayloadType.MULTIPART:
      return this.decodeMultipart();
    case PayloadType.CONTROL:
      return this.decodeControl();
    case PayloadType.RAW_CUSTOM:
      return this.decodeRawCustom();
    default:
      throw new Error(`MeshCore: can't decode payload ${this.payloadType}`)
    }
  }

  private decodeEncrypted<T>(kind: PayloadType): T {
    const buffer = new BufferReader(this.payload);
    return {
      payloadType: kind,
      dstHash:     buffer.readUint8().toString(16).padStart(2, '0'),
      srcHash:     buffer.readUint8().toString(16).padStart(2, '0'),
      cipherMAC:   buffer.readBytes(2),
      cipherText:  buffer.readBytes()
    } as T
  }

  private decodeRequest(): RequestPayload {
    return this.decodeEncrypted<RequestPayload>(PayloadType.REQUEST);
  }

  private decodeResponse(): ResponsePayload {
    return this.decodeEncrypted<ResponsePayload>(PayloadType.RESPONSE);
  }

  private decodeText(): TextPayload {
    return this.decodeEncrypted<TextPayload>(PayloadType.TEXT);
  }

  private decodeAck(): AckPayload {
    return this.decodeEncrypted<AckPayload>(PayloadType.ACK);
  }

  private decodeAdvert(): AdvertPayload {
    const buffer = new BufferReader(this.payload);
    const publicKey = buffer.readBytes(32);
    const timestampValue = buffer.readUint32LE();
    const signature = buffer.readBytes(64);
    let payload: AdvertPayload = {
      payloadType: PayloadType.ADVERT,
      publicKey:   publicKey,
      timestamp:   timestampValue === 0 ? undefined : new Date(timestampValue),
      signature:   signature,
      appdata: {
        flags: buffer.readUint8(),
      }
    }
    if (payload.appdata.flags & AdvertisementFlags.HAS_LOCATION) {
      payload.appdata.latitude = buffer.readUint32LE() / 1000000.0
      payload.appdata.longitude = buffer.readUint32LE() / 1000000.0
    }
    if (payload.appdata.flags & AdvertisementFlags.HAS_FEATURE_1) {
      payload.appdata.feature1 = buffer.readUint16LE()
    }
    if (payload.appdata.flags & AdvertisementFlags.HAS_FEATURE_2) {
      payload.appdata.feature2 = buffer.readUint16LE()
    }
    if (payload.appdata.flags & AdvertisementFlags.HAS_NAME) {
      payload.appdata.name = new TextDecoder().decode(buffer.readBytes()).toString()
    }
    return payload;
  }

  private decodeGroupEncrypted<T>(kind: PayloadType): T {
    const buffer = new BufferReader(this.payload);
    return {
      payloadType: kind,
      channelHash: buffer.readUint8().toString(16).padStart(2, '0'),
      cipherMAC:   buffer.readBytes(2),
      cipherText:  buffer.readBytes()
    } as T;
  }

  private decodeGroupText(): GroupTextPayload {
    return this.decodeGroupEncrypted<GroupTextPayload>(PayloadType.GROUP_TEXT);
  }

  private decodeGroupData(): GroupDataPayload {
    return this.decodeGroupEncrypted<GroupDataPayload>(PayloadType.GROUP_DATA);
  }

  private decodeAnonReq(): AnonReqPayload {
    const buffer = new BufferReader(this.payload);
    return {
      payloadType: PayloadType.ANON_REQ,
      dstHash:     buffer.readUint8().toString(16).padStart(2, '0'),
      publicKey:   buffer.readBytes(32),
      cipherMAC:   buffer.readBytes(2),
      cipherText:  buffer.readBytes()
    }
  }

  private decodePath(): PathPayload {
    return this.decodeEncrypted<PathPayload>(PayloadType.PATH);
  }

  private decodeTrace(): TracePayload {
    const buffer = new BufferReader(this.payload);
    return {
      payloadType: PayloadType.TRACE,
      data:        buffer.readBytes()
    }
  }

  private decodeMultipart(): MultipartPayload {
    const buffer = new BufferReader(this.payload);
    return {
      payloadType: PayloadType.MULTIPART,
      data:        buffer.readBytes()
    }
  }

  private decodeControl(): ControlPayload {
    const buffer = new BufferReader(this.payload);
    return {
      payloadType: PayloadType.CONTROL,
      flags:       buffer.readUint8(),
      data:        buffer.readBytes()
    }
  }

  private decodeRawCustom(): RawCustomPayload {
    const buffer = new BufferReader(this.payload);
    return {
      payloadType: PayloadType.RAW_CUSTOM,
      data:        buffer.readBytes()
    }
  }
}

export class PublicKey implements BasePublicKey {
  private pub: Uint8Array;

  constructor(key: PublicKeyValue) {
    if (key instanceof Uint8Array) {
      this.pub = Uint8Array.from(key);
    } else {
      this.pub = hexToBytes(key);
    }
    if (this.pub.length !== 8 && this.pub.length !== 32) {
      throw new Error(`invalid public key, expected 8 or 32 bytes, got ${this.pub.length}`);
    }
  }

  toBytes(): Uint8Array {
    return this.pub.slice();
  }

  toString(): string {
    return bytesToHex(this.pub);
  }

  verify(message: Uint8Array, signature: Uint8Array): boolean {
    return ed25519.verify(signature, message, this.pub);
  }
}

export class PrivateKey implements BasePrivateKey {
  private seed: Uint8Array;
  private pub: Uint8Array = new Uint8Array(32);
  private prefix: Uint8Array = new Uint8Array(32);

  constructor(seed: PublicKeyValue) {
    if (seed instanceof Uint8Array) {
      this.seed = Uint8Array.from(seed);
    } else {
      this.seed = hexToBytes(seed);
    }
    if (this.seed.length !== 32) {
      throw new Error(`invalid public key, expected 32 bytes seed, got ${this.seed.length}`);
    }
    this.precompute();
  }

  private precompute() {
    const hash = sha512.create().update(this.seed).digest();
    const h32 = hash.slice(0, 32);
    const clampedScalar = this.clampScalar(h32);
    const publicKey = ed25519.getPublicKey(clampedScalar);
    this.pub.set(publicKey, 32);
    this.prefix.set(hash.slice(32, 64));
  }

  private clampScalar(scalar: Uint8Array): Uint8Array {
    const clamped = new Uint8Array(scalar);
    clamped[0] &= 248;  // Clear lowest 3 bits
    clamped[31] &= 127; // Clear highest bit
    clamped[31] |= 64;  // Set second highest bit
    return clamped;
  }

  toBytes(): Uint8Array {
    return this.seed.slice();
  }

  toString(): string {
    return bytesToHex(this.seed);
  }

  toPublicKey(): PublicKey {
    return new PublicKey(this.seed.slice(32, 64));
  }

  sign(message: Uint8Array): Uint8Array {
    return ed25519.sign(message, this.seed);
  }

  verify(message: Uint8Array, signature: Uint8Array): boolean {
    return ed25519.verify(signature, message, this.pub);
  }
}

export class SharedSecret implements BaseSharedSecret {
  private bytes: Uint8Array;

  constructor(secret: string | Uint8Array) {
    if (typeof secret === 'string') {
      this.bytes = hexToBytes(secret);
    } else {
      this.bytes = Uint8Array.from(secret);
    }
    if (this.bytes.length !== 32) {
      throw new Error(`invalid shared secret size ${this.bytes.length}`)
    }
  }

  decrypt(cipherText: Uint8Array, cipherMAC: Uint8Array): Uint8Array {
    const ourMAC = hmac(sha256, this.bytes, cipherText).slice(0, 2);
    if (!constantTimeEqual(cipherMAC, ourMAC)) {
      throw new Error(`invalid MAC`)
    }

    const block = ecb(this.bytes);
    const plain = block.decrypt(cipherText);

    return plain;
  }
}

export class StaticSecret {
  private bytes: Uint8Array;

  constructor(secret: string | Uint8Array) {
    if (secret instanceof Uint8Array) {
      this.bytes = Uint8Array.from(secret);
    } else {
      this.bytes = hexToBytes(secret);
    }
    if (this.bytes.length !== 32) {
      throw new Error(`invalid shared secret size ${this.bytes.length}`)
    }
  }

  diffieHellman(other: PublicKey): SharedSecret {
    const bytes = x25519.getSharedSecret(this.bytes, other.toBytes());
    return new SharedSecret(bytes);
  }

  publicKey(): PublicKey {
    const bytes = x25519.getPublicKey(this.bytes);
    return new PublicKey(bytes);
  }
}

export class GroupSecret implements BaseGroupSecret {
  private bytes: Uint8Array;

  constructor(secret: GroupSecretValue) {
    if (secret instanceof Uint8Array) {
      this.bytes = Uint8Array.from(secret);
    } else {
      this.bytes = hexToBytes(secret);
    }

    if (this.bytes.length === 16) {
      const padded = new Uint8Array(32);
      padded.set(this.bytes, 0);
      this.bytes = padded;
    }

    if (this.bytes.length !== 32) {
      throw new Error(`invalid group secret, expected 32 bytes key, got ${this.bytes.length}`);
    }
  }

  toBytes(): Uint8Array {
    return this.bytes.slice();
  }

  toString(): string {
    return bytesToHex(this.bytes);
  }

  toHash(): string {
    return this.bytes[0].toString(16).padStart(2, '0')
  }

  decrypt(cipherText: Uint8Array, cipherMAC: Uint8Array): DecryptedGroupMessage {
    const ourMAC = hmac(sha256, this.bytes, cipherText)
    if (!constantTimeEqual(cipherMAC, ourMAC)) {
      throw new Error('invalid MAC');
    }

    const block = ecb(this.bytes);
    const plain = block.decrypt(cipherText);
    if (plain.length < 5) {
      throw new Error('invalid payload');
    }

    const reader = new BufferReader(plain);
    const timestamp = new Date(reader.readUint32LE());
    const flags = reader.readUint8();
    let message = new TextDecoder('utf-8').decode(reader.readBytes());
    const nullPos = message.indexOf('\0')
    if (nullPos >= 0) {
      message = message.substring(0, nullPos);
    }

    return {
      timestamp,
      flags,
      message
    }
  }

  static fromName(name: string): GroupSecret {
    const hash = sha256.create().update(new TextEncoder().encode(name)).digest();
    return new GroupSecret(hash.slice(0, 32));
  }
}

export interface DecryptedGroupMessageResult {
  success: boolean;
  result?: DecryptedGroupMessage;
  error?: string;
}

export class KeyManager {
  protected groups: Map<string, Group[]> = new Map();

  public addGroup(name: string, secret?: string | Uint8Array) {
    let group: Group;
    if (secret) {
      group = {
        name:   name,
        secret: new GroupSecret(secret)
      }
    } else {
      group = {
        name:   name,
        secret: GroupSecret.fromName(name)
      }
    }
    const hash = group.secret.toHash();
    this.groups.set(hash, [...this.groups.get(hash) || [], group]);
  }

  public decryptGroup(channelHash: NodeHash, cipherText: Uint8Array, cipherMAC: Uint8Array): DecryptedGroupMessage {
    if (!this.groups.has(channelHash)) {
      throw new Error(`no keys found for channel hash ${channelHash}`);
    }

    let error: any;
    for (const group of this.groups.get(channelHash) || []) {
      try {
        return {
          ...group.secret.decrypt(cipherText, cipherMAC),
          group: group.name
        }
      } catch (e) {
        error = e;
      }
    }
    throw error;
  }
}
