import { x25519 } from "@noble/curves/ed25519";
import { PrivateKey, PublicKey, SharedSecret, StaticSecret } from "./crypto";
import { IPublicKey } from "./crypto.types";
import { NodeHash, IIdentity, ILocalIdentity } from "./identity.types";
import { DecryptedGroupData, DecryptedGroupText } from "./packet.types";
import { equalBytes, hexToBytes, BufferReader, BufferWriter } from "./parser";

export const parseNodeHash = (hash: NodeHash | string | Uint8Array): NodeHash => {
  if (hash instanceof Uint8Array) {
    return hash[0] as NodeHash;
  }
  if (typeof hash === "number") {
    if (hash < 0 || hash > 255) {
      throw new Error("NodeHash number must be between 0x00 and 0xFF");
    }
    return hash as NodeHash;
  } else if (typeof hash === "string") {
    const parsed = hexToBytes(hash);
    if (parsed.length !== 1) {
      throw new Error("NodeHash string must represent a single byte");
    }
    return parsed[0] as NodeHash;
  }
  throw new Error("Invalid NodeHash type");
}

const toPublicKeyBytes = (key: Identity | PublicKey | Uint8Array | string): Uint8Array => {
  if (key instanceof Identity) {
    return key.publicKey.toBytes();
  } else if (key instanceof PublicKey) {
    return key.toBytes();
  } else if (key instanceof Uint8Array) {
    return key;
  } else if (typeof key === 'string') {
    return hexToBytes(key);
  } else {
    throw new Error('Invalid type for toPublicKeyBytes');
  }
}

export class Identity implements IIdentity {
  public publicKey: PublicKey;

  constructor(publicKey: PublicKey | Uint8Array | string) {
    if (publicKey instanceof PublicKey) {
      this.publicKey = publicKey;
    } else if (publicKey instanceof Uint8Array || typeof publicKey === 'string') {
      this.publicKey = new PublicKey(publicKey);
    } else {
      throw new Error('Invalid type for Identity constructor');
    }
  }

  hash(): NodeHash {
    return this.publicKey.toHash();
  }

  toString(): string {
    return this.publicKey.toString();
  }

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

  matches(other: Identity | PublicKey | Uint8Array | string): boolean {
    return this.publicKey.equals(toPublicKeyBytes(other));
  }
}

export class LocalIdentity extends Identity implements ILocalIdentity {
  private privateKey: PrivateKey;

  constructor(privateKey: PrivateKey | Uint8Array | string, publicKey: PublicKey | Uint8Array | string) {
    if (publicKey instanceof PublicKey) {
      super(publicKey.toBytes());
    } else {
      super(publicKey);
    }

    if (privateKey instanceof PrivateKey) {
      this.privateKey = privateKey;
    } else {
      this.privateKey = new PrivateKey(privateKey);
    }
  }

  sign(message: Uint8Array): Uint8Array {
    return this.privateKey.sign(message);
  }

  calculateSharedSecret(other: IIdentity | IPublicKey): SharedSecret {
    let otherPublicKey: PublicKey;
    if (other instanceof Identity) {
      otherPublicKey = other.publicKey;
    } else if ('toBytes' in other) {
      otherPublicKey = new PublicKey(other.toBytes());
    } else if ('publicKey' in other && other.publicKey instanceof Uint8Array) {
      otherPublicKey = new PublicKey(other.publicKey);
    } else if ('publicKey' in other && other.publicKey instanceof PublicKey) {
      otherPublicKey = other.publicKey;
    } else if ('publicKey' in other && typeof other.publicKey === 'function') {
      otherPublicKey = new PublicKey(other.publicKey().toBytes());
    } else {
      throw new Error('Invalid type for calculateSharedSecret comparison');
    }
    return new SharedSecret(this.privateKey.calculateSharedSecret(otherPublicKey));
  }
}

export class Contact {
  public name: string = "";
  public identity: Identity;

  constructor(name: string, identity: Identity | PublicKey | Uint8Array | string) {
    this.name = name;
    if (identity instanceof Identity) {
      this.identity = identity;
    } else if (identity instanceof PublicKey) {
      this.identity = new Identity(identity);
    } else if (identity instanceof Uint8Array || typeof identity === 'string') {
      this.identity = new Identity(identity);
    } else {
      throw new Error('Invalid type for Contact constructor');
    }
  }

  public matches(hash: Uint8Array | PublicKey): boolean {
    return this.identity.publicKey.equals(hash);
  }

  public publicKey(): PublicKey {
    return this.identity.publicKey;
  }

  public calculateSharedSecret(me: LocalIdentity): SharedSecret {
    return me.calculateSharedSecret(this.identity);
  }
}

export class Group {
  public name: string;
  private secret: SharedSecret;

  constructor(name: string, secret?: SharedSecret) {
    this.name = name;
    if (secret) {
      this.secret = secret;
    } else {
      this.secret = SharedSecret.fromName(name);
    }
  }

  public hash(): NodeHash {
    return this.secret.toHash();
  }

  public decryptText(hmac: Uint8Array, ciphertext: Uint8Array): DecryptedGroupText {
    const data = this.secret.decrypt(hmac, ciphertext);
    if (data.length < 5) {
      throw new Error("Invalid ciphertext");
    }

    const reader = new BufferReader(data);
    const timestamp = reader.readTimestamp();
    const flags = reader.readByte();
    const textType = (flags >> 2) & 0x3F;
    const attempt = flags & 0x03;
    const message = new TextDecoder('utf-8').decode(reader.readBytes());
    return {
      timestamp,
      textType,
      attempt,
      message
    }
  }

  public encryptText(plain: DecryptedGroupText): { hmac: Uint8Array, ciphertext: Uint8Array } {
    const writer = new BufferWriter();
    writer.writeTimestamp(plain.timestamp);
    const flags = ((plain.textType & 0x3F) << 2) | (plain.attempt & 0x03);
    writer.writeByte(flags);
    writer.writeBytes(new TextEncoder().encode(plain.message));
    const data = writer.toBytes();
    return this.secret.encrypt(data);
  }

  public decryptData(hmac: Uint8Array, ciphertext: Uint8Array): DecryptedGroupData {
    const data = this.secret.decrypt(hmac, ciphertext);
    if (data.length < 4) {
      throw new Error("Invalid ciphertext");
    }

    const reader = new BufferReader(data);
    return {
      timestamp: reader.readTimestamp(),
      data:      reader.readBytes(reader.remainingBytes())
    }
  }

  public encryptData(plain: DecryptedGroupData): { hmac: Uint8Array, ciphertext: Uint8Array } {
    const writer = new BufferWriter();
    writer.writeTimestamp(plain.timestamp);
    writer.writeBytes(plain.data);
    const data = writer.toBytes();
    return this.secret.encrypt(data);
  }
}

interface CachedLocalIdentity {
  identity:      LocalIdentity;
  sharedSecrets: Record<string, SharedSecret>;
}

export class Contacts {
  private localIdentities: CachedLocalIdentity[] = [];
  private contacts: Record<number, Contact[]> = {};
  private groups: Record<number, Group[]> = {};

  public addLocalIdentity(identity: LocalIdentity) {
    this.localIdentities.push({ identity, sharedSecrets: {} });
  }

  public addContact(contact: Contact) {
    const hash = parseNodeHash(contact.identity.hash()) as number;
    if (!this.contacts[hash]) {
      this.contacts[hash] = [];
    }
    this.contacts[hash].push(contact);
  }

  public decrypt(src: NodeHash | PublicKey, dst: NodeHash, hmac: Uint8Array, ciphertext: Uint8Array): {
    localIdentity: LocalIdentity,
    contact: Contact,
    decrypted: Uint8Array,
  } {
    // Find the public key associated with the source hash.
    let contacts: Contact[] = [];
    if (src instanceof PublicKey) {
      // Check if we have a contact with this exact public key (for direct messages).
      const srcHash = parseNodeHash(src.toHash()) as number;
      for (const contact of this.contacts[srcHash] || []) {
        if (contact.identity.matches(src)) {
          contacts.push(contact);
        }
      }

      // If no contact matches the public key, add a temporary contact with the hash and no name.
      if (contacts.length === 0) {
        contacts.push(new Contact("", new Identity(src.toBytes())));
      }
    } else {
      const srcHash = parseNodeHash(src) as number;
      contacts = this.contacts[srcHash] || [];
    }
    if (contacts.length === 0) {
      throw new Error("Unknown source hash");
    }

    // Find the local identity associated with the destination hash.
    const dstHash = parseNodeHash(dst) as number;
    const localIdentities = this.localIdentities.filter(li => li.identity.publicKey.key[0] === dstHash);
    if (localIdentities.length === 0) {
      throw new Error("Unknown destination hash");
    }

    // Try to decrypt with each combination of local identity and their public key.
    for (const localIdentity of localIdentities) {
      for (const contact of contacts) {
        const sharedSecret: SharedSecret = this.calculateSharedSecret(localIdentity, contact);
        try {
          const decrypted = sharedSecret.decrypt(hmac, ciphertext);
          return { localIdentity: localIdentity.identity, contact, decrypted };
        } catch {
          // Ignore decryption errors and try the next combination.
        }
      }
    }

    throw new Error("Decryption failed with all known identities and contacts");
  }

  // Caches the calculated shared secret for a given local identity and contact to avoid redundant calculations.
  private calculateSharedSecret(localIdentity: CachedLocalIdentity, contact: Contact): SharedSecret {
    const cacheKey = contact.identity.toString();
    if (localIdentity.sharedSecrets[cacheKey]) {
      return localIdentity.sharedSecrets[cacheKey];
    }
    const sharedSecret = localIdentity.identity.calculateSharedSecret(contact.identity);
    localIdentity.sharedSecrets[cacheKey] = sharedSecret;
    return sharedSecret;
  }

  public addGroup(group: Group) {
    const hash = parseNodeHash(group.hash()) as number;
    if (!this.groups[hash]) {
      this.groups[hash] = [];
    }
    this.groups[hash].push(group);
  }

  public decryptGroupText(channelHash: NodeHash, hmac: Uint8Array, ciphertext: Uint8Array): {
    decrypted: DecryptedGroupText,
    group: Group
  } {
    const hash = parseNodeHash(channelHash) as number;
    const groups = this.groups[hash] || [];
    if (groups.length === 0) {
      throw new Error("Unknown group hash");
    }
    for (const group of groups) {
      try {
        const decrypted = group.decryptText(hmac, ciphertext);
        return { decrypted, group };
      } catch {
        // Ignore decryption errors and try the next group.
      }
    }
    throw new Error("Decryption failed with all known groups");
  }

  public decryptGroupData(channelHash: NodeHash, hmac: Uint8Array, ciphertext: Uint8Array): {
    decrypted: DecryptedGroupData,
    group: Group
  } {
    const hash = parseNodeHash(channelHash) as number;
    const groups = this.groups[hash] || [];
    if (groups.length === 0) {
      throw new Error("Unknown group hash");
    }
    for (const group of groups) {
      try {
        const decrypted = group.decryptData(hmac, ciphertext);
        return { decrypted, group };
      } catch {
        // Ignore decryption errors and try the next group.
      }
    }
    throw new Error("Decryption failed with all known groups");
  }
}