import ExpoModulesCore
import LibSignalClient

// MARK: - Record Types for Type-Safe JS Bridge Conversion

struct IdentityKeyPairRecord: Record {
  @Field var publicKey: Data = Data()
  @Field var privateKey: Data = Data()
}

struct PreKeyDataRecord: Record {
  @Field var id: Int = 0
  @Field var publicKey: Data = Data()
  @Field var privateKey: Data = Data()
}

struct SignedPreKeyDataRecord: Record {
  @Field var id: Int = 0
  @Field var publicKey: Data = Data()
  @Field var privateKey: Data = Data()
  @Field var signature: Data = Data()
  @Field var timestamp: UInt64 = 0
}

struct AddressRecord: Record {
  @Field var name: String = ""
  @Field var deviceId: Int = 0
}

struct CiphertextRecord: Record {
  @Field var type: String = ""
  @Field var body: Data = Data()
}

struct PreKeyBundleRecord: Record {
  @Field var registrationId: Int = 0
  @Field var deviceId: Int = 0
  @Field var preKeyId: Int = 0
  @Field var preKeyPublic: Data = Data()
  @Field var signedPreKeyId: Int = 0
  @Field var signedPreKeyPublic: Data = Data()
  @Field var signedPreKeySignature: Data = Data()
  @Field var identityKey: Data = Data()
  @Field var kyberPreKeyId: Int? = nil
  @Field var kyberPreKey: Data? = nil
  @Field var kyberPreKeySignature: Data? = nil
}

public class SignalExpoModule: Module {
  // Shared stores - persist across function calls
  private var identityKeyPair: IdentityKeyPair?
  private var localRegistrationId: UInt32 = 0

  // Protocol stores
  private lazy var sessionStore = InMemorySessionStore()
  private lazy var identityStore = InMemoryIdentityStore()
  private lazy var preKeyStore = InMemoryPreKeyStore()
  private lazy var signedPreKeyStore = InMemorySignedPreKeyStore()
  private lazy var kyberPreKeyStore = InMemoryKyberPreKeyStore()

  public func definition() -> ModuleDefinition {
    Name("SignalExpo")

    // MARK: - Key Generation Functions

    Function("generateIdentityKeyPair") { () -> [String: Any] in
      let keyPair = IdentityKeyPair.generate()
      return [
        "publicKey": Data(keyPair.publicKey.serialize()),
        "privateKey": Data(keyPair.privateKey.serialize())
      ]
    }

    Function("generateRegistrationId") { () -> UInt32 in
      return UInt32.random(in: 1...0x3FFF) // 14-bit value per Signal spec
    }

    Function("generatePreKeys") { (start: Int, count: Int) -> [[String: Any]] in
      var preKeys: [[String: Any]] = []

      for i in 0..<count {
        let preKeyId = UInt32(start + i)
        let preKeyRecord = try PreKeyRecord(id: preKeyId, privateKey: PrivateKey.generate())

        preKeys.append([
          "id": preKeyId,
          "publicKey": Data(try preKeyRecord.publicKey().serialize()),
          "privateKey": Data(try preKeyRecord.privateKey().serialize())
        ])
      }

      return preKeys
    }

    Function("generateSignedPreKey") { (identityPrivateKey: Data, signedPreKeyId: Int) -> [String: Any] in
      let privateKey = try PrivateKey(identityPrivateKey)
      let signedPreKeyPrivate = PrivateKey.generate()
      let signedPreKeyPublic = signedPreKeyPrivate.publicKey

      let signature = privateKey.generateSignature(message: signedPreKeyPublic.serialize())
      let timestamp = UInt64(Date().timeIntervalSince1970 * 1000)

      return [
        "id": signedPreKeyId,
        "publicKey": Data(signedPreKeyPublic.serialize()),
        "privateKey": Data(signedPreKeyPrivate.serialize()),
        "signature": Data(signature),
        "timestamp": timestamp
      ]
    }

    // MARK: - Kyber (Post-Quantum) Key Generation

    Function("generateKyberPreKeys") { (start: Int, count: Int, identityPrivateKey: Data) -> [[String: Any]] in
      let privateKey = try PrivateKey(identityPrivateKey)
      var kyberPreKeys: [[String: Any]] = []

      for i in 0..<count {
        let kyberPreKeyId = UInt32(start + i)
        let kyberKeyPair = KEMKeyPair.generate()

        // Sign the Kyber public key with identity key
        let signature = privateKey.generateSignature(message: kyberKeyPair.publicKey.serialize())
        let timestamp = UInt64(Date().timeIntervalSince1970 * 1000)

        let record = try KyberPreKeyRecord(
          id: kyberPreKeyId,
          timestamp: timestamp,
          keyPair: kyberKeyPair,
          signature: signature
        )

        // Return serialized record for persistence (includes all data needed for restoration)
        kyberPreKeys.append([
          "id": kyberPreKeyId,
          "publicKey": Data(kyberKeyPair.publicKey.serialize()),
          "signature": Data(signature),
          "timestamp": timestamp,
          "serialized": Data(record.serialize())
        ])

        // Store the record
        try self.kyberPreKeyStore.storeKyberPreKey(record, id: kyberPreKeyId, context: NullContext())
      }

      return kyberPreKeys
    }

    // MARK: - Initialization

    AsyncFunction("initialize") { (
      identityPrivateKey: Data,
      identityPublicKey: Data,
      registrationId: UInt32,
      preKeyIds: [Int],
      preKeyPrivateKeys: [Data],
      signedPreKeyId: Int,
      signedPreKeyPrivateKey: Data,
      signedPreKeySignature: Data,
      signedPreKeyTimestamp: Double
    ) in
      // Reconstruct identity key pair from direct params
      let privateKey = try PrivateKey(identityPrivateKey)
      let publicKey = try PublicKey(identityPublicKey)
      self.identityKeyPair = IdentityKeyPair(publicKey: publicKey, privateKey: privateKey)
      self.localRegistrationId = registrationId

      // Store identity in identity store
      self.identityStore.identityKeyPair = self.identityKeyPair
      self.identityStore.localRegistrationId = registrationId

      // Store pre-keys from parallel arrays
      for i in 0..<preKeyIds.count {
        let preKeyPrivate = try PrivateKey(preKeyPrivateKeys[i])
        let preKeyRecord = try LibSignalClient.PreKeyRecord(id: UInt32(preKeyIds[i]), privateKey: preKeyPrivate)
        try self.preKeyStore.storePreKey(preKeyRecord, id: UInt32(preKeyIds[i]), context: NullContext())
      }

      // Store signed pre-key from direct params
      let signedPrivateKey = try PrivateKey(signedPreKeyPrivateKey)
      let signedPreKeyRecord = try LibSignalClient.SignedPreKeyRecord(
        id: UInt32(signedPreKeyId),
        timestamp: UInt64(signedPreKeyTimestamp),
        privateKey: signedPrivateKey,
        signature: [UInt8](signedPreKeySignature)
      )
      try self.signedPreKeyStore.storeSignedPreKey(signedPreKeyRecord, id: UInt32(signedPreKeyId), context: NullContext())
    }

    // MARK: - Kyber Pre-Key Restoration (for initializeFromStore)

    AsyncFunction("storeKyberPreKeys") { (
      kyberPreKeyIds: [Int],
      kyberPreKeySerializedRecords: [Data]
    ) in
      guard kyberPreKeyIds.count == kyberPreKeySerializedRecords.count else {
        throw SignalExpoError.invalidKeyData
      }

      for i in 0..<kyberPreKeyIds.count {
        let kyberPreKeyId = UInt32(kyberPreKeyIds[i])
        // Deserialize the full record from stored bytes
        let record = try KyberPreKeyRecord(bytes: kyberPreKeySerializedRecords[i])
        try self.kyberPreKeyStore.storeKyberPreKey(record, id: kyberPreKeyId, context: NullContext())
      }
    }

    // MARK: - Session Management Functions

    AsyncFunction("createSession") { (
      addressName: String,
      intParams: [Int],  // [addressDeviceId, registrationId, deviceId, preKeyId, signedPreKeyId, kyberPreKeyId]
      preKeyPublic: Data,
      signedPreKeyPublic: Data,
      signedPreKeySignature: Data,
      identityKey: Data,
      kyberPreKey: Data,
      kyberPreKeySignature: Data
    ) in
      // Extract int params (Data must be separate params - arrays don't convert Uint8Array properly)
      guard intParams.count == 6 else {
        throw SignalExpoError.invalidKeyData
      }
      let addressDeviceId = intParams[0]
      let registrationId = intParams[1]
      let deviceId = intParams[2]
      let preKeyId = intParams[3]
      let signedPreKeyId = intParams[4]
      let kyberPreKeyId = intParams[5]

      let address = try ProtocolAddress(name: addressName, deviceId: UInt32(addressDeviceId))
      let identity = try IdentityKey(publicKey: PublicKey(identityKey))
      let preKey = try PublicKey(preKeyPublic)
      let signedPreKey = try PublicKey(signedPreKeyPublic)
      let kyberKey = try KEMPublicKey(kyberPreKey)

      let bundle = try LibSignalClient.PreKeyBundle(
        registrationId: UInt32(registrationId),
        deviceId: UInt32(deviceId),
        prekeyId: UInt32(preKeyId),
        prekey: preKey,
        signedPrekeyId: UInt32(signedPreKeyId),
        signedPrekey: signedPreKey,
        signedPrekeySignature: [UInt8](signedPreKeySignature),
        identity: identity,
        kyberPrekeyId: UInt32(kyberPreKeyId),
        kyberPrekey: kyberKey,
        kyberPrekeySignature: [UInt8](kyberPreKeySignature)
      )

      // Process the pre-key bundle to establish a session
      try processPreKeyBundle(
        bundle,
        for: address,
        sessionStore: self.sessionStore,
        identityStore: self.identityStore,
        context: NullContext()
      )
    }

    AsyncFunction("hasSession") { (addressName: String, addressDeviceId: Int) -> Bool in
      let address = try ProtocolAddress(name: addressName, deviceId: UInt32(addressDeviceId))
      let session = try self.sessionStore.loadSession(for: address, context: NullContext())
      return session != nil
    }

    AsyncFunction("deleteSession") { (addressName: String, addressDeviceId: Int) in
      let address = try ProtocolAddress(name: addressName, deviceId: UInt32(addressDeviceId))
      self.sessionStore.removeSession(for: address)
    }

    // MARK: - Encryption/Decryption Functions

    AsyncFunction("encrypt") { (addressName: String, addressDeviceId: Int, plaintext: Data) -> [String: Any] in
      let address = try ProtocolAddress(name: addressName, deviceId: UInt32(addressDeviceId))

      let ciphertext = try signalEncrypt(
        message: [UInt8](plaintext),
        for: address,
        sessionStore: self.sessionStore,
        identityStore: self.identityStore,
        context: NullContext()
      )

      let messageType: String
      switch ciphertext.messageType {
      case .preKey:
        messageType = "preKey"
      case .whisper:
        messageType = "whisper"
      case .senderKey:
        messageType = "senderKey"
      case .plaintext:
        messageType = "plaintext"
      default:
        messageType = "unknown"
      }

      return [
        "type": messageType,
        "body": Data(ciphertext.serialize())
      ]
    }

    AsyncFunction("decrypt") { (addressName: String, addressDeviceId: Int, ciphertextType: String, ciphertextBody: Data) -> [String: Any] in
      let address = try ProtocolAddress(name: addressName, deviceId: UInt32(addressDeviceId))

      let plaintext: [UInt8]

      if ciphertextType == "preKey" {
        let preKeyMessage = try PreKeySignalMessage(bytes: [UInt8](ciphertextBody))
        let decryptedData = try signalDecryptPreKey(
          message: preKeyMessage,
          from: address,
          sessionStore: self.sessionStore,
          identityStore: self.identityStore,
          preKeyStore: self.preKeyStore,
          signedPreKeyStore: self.signedPreKeyStore,
          kyberPreKeyStore: self.kyberPreKeyStore,
          context: NullContext()
        )
        plaintext = [UInt8](decryptedData)
      } else {
        let signalMessage = try SignalMessage(bytes: [UInt8](ciphertextBody))
        let decryptedData = try signalDecrypt(
          message: signalMessage,
          from: address,
          sessionStore: self.sessionStore,
          identityStore: self.identityStore,
          context: NullContext()
        )
        plaintext = [UInt8](decryptedData)
      }

      return [
        "plaintext": Data(plaintext)
      ]
    }

    // MARK: - Session Persistence

    AsyncFunction("exportSession") { (addressName: String, addressDeviceId: Int) -> Data? in
      let address = try ProtocolAddress(name: addressName, deviceId: UInt32(addressDeviceId))
      guard let session = try self.sessionStore.loadSession(for: address, context: NullContext()) else {
        return nil
      }
      return Data(session.serialize())
    }

    AsyncFunction("importSession") { (addressName: String, addressDeviceId: Int, serializedSession: Data) in
      let address = try ProtocolAddress(name: addressName, deviceId: UInt32(addressDeviceId))
      let session = try SessionRecord(bytes: [UInt8](serializedSession))
      try self.sessionStore.storeSession(session, for: address, context: NullContext())
    }

    AsyncFunction("listSessions") { () -> [[Any]] in
      return self.sessionStore.getAllAddresses().map { [$0.name, $0.deviceId] }
    }

    // MARK: - Storage Accessors

    AsyncFunction("getIdentityPublicKey") { () -> Data in
      guard let keyPair = self.identityKeyPair else {
        throw SignalExpoError.notInitialized
      }
      return Data(keyPair.publicKey.serialize())
    }

    AsyncFunction("getLocalRegistrationId") { () -> UInt32 in
      guard self.identityKeyPair != nil else {
        throw SignalExpoError.notInitialized
      }
      return self.localRegistrationId
    }

    // MARK: - Clear all stores (for user switching)

    AsyncFunction("clear") { () in
      // Clear identity
      self.identityKeyPair = nil
      self.localRegistrationId = 0
      self.identityStore.identityKeyPair = nil
      self.identityStore.localRegistrationId = 0

      // Reset all in-memory stores
      self.sessionStore.clear()
      self.preKeyStore.clear()
      self.signedPreKeyStore.clear()
      self.kyberPreKeyStore.clear()
      self.identityStore.clear()
    }
  }

  // MARK: - Helper Methods

  private func parsePreKeyBundle(_ bundle: PreKeyBundleRecord) throws -> LibSignalClient.PreKeyBundle {
    // Kyber keys are required in LibSignalClient v0.86.10+
    guard let kyberPreKeyId = bundle.kyberPreKeyId,
          let kyberPreKey = bundle.kyberPreKey,
          let kyberPreKeySignature = bundle.kyberPreKeySignature else {
      throw SignalExpoError.kyberKeysRequired
    }

    return try LibSignalClient.PreKeyBundle(
      registrationId: UInt32(bundle.registrationId),
      deviceId: UInt32(bundle.deviceId),
      prekeyId: UInt32(bundle.preKeyId),
      prekey: PublicKey(bundle.preKeyPublic),
      signedPrekeyId: UInt32(bundle.signedPreKeyId),
      signedPrekey: PublicKey(bundle.signedPreKeyPublic),
      signedPrekeySignature: [UInt8](bundle.signedPreKeySignature),
      identity: IdentityKey(publicKey: PublicKey(bundle.identityKey)),
      kyberPrekeyId: UInt32(kyberPreKeyId),
      kyberPrekey: try KEMPublicKey(kyberPreKey),
      kyberPrekeySignature: [UInt8](kyberPreKeySignature)
    )
  }
}

// MARK: - Error Types

enum SignalExpoError: Error {
  case invalidKeyData
  case invalidAddressData
  case invalidBundleData
  case invalidCiphertextData
  case notInitialized
  case noSession
  case kyberKeysRequired
}

// MARK: - In-Memory Stores

class InMemorySessionStore: SessionStore {
  private var sessions: [ProtocolAddress: SessionRecord] = [:]

  func loadSession(for address: ProtocolAddress, context: StoreContext) throws -> SessionRecord? {
    return sessions[address]
  }

  func loadExistingSessions(for addresses: [ProtocolAddress], context: StoreContext) throws -> [SessionRecord] {
    return addresses.compactMap { sessions[$0] }
  }

  func storeSession(_ record: SessionRecord, for address: ProtocolAddress, context: StoreContext) throws {
    sessions[address] = record
  }

  func removeSession(for address: ProtocolAddress) {
    sessions.removeValue(forKey: address)
  }

  func clear() {
    sessions.removeAll()
  }

  func getAllAddresses() -> [ProtocolAddress] {
    return Array(sessions.keys)
  }
}

class InMemoryIdentityStore: IdentityKeyStore {
  var identityKeyPair: IdentityKeyPair?
  var localRegistrationId: UInt32 = 0
  private var identities: [ProtocolAddress: IdentityKey] = [:]

  func identityKeyPair(context: StoreContext) throws -> IdentityKeyPair {
    guard let keyPair = identityKeyPair else {
      throw SignalExpoError.notInitialized
    }
    return keyPair
  }

  func localRegistrationId(context: StoreContext) throws -> UInt32 {
    return localRegistrationId
  }

  func saveIdentity(_ identity: IdentityKey, for address: ProtocolAddress, context: StoreContext) throws -> IdentityChange {
    let existing = identities[address]
    identities[address] = identity
    if let existing = existing, existing != identity {
      return .replacedExisting
    } else {
      return .newOrUnchanged
    }
  }

  func isTrustedIdentity(_ identity: IdentityKey, for address: ProtocolAddress, direction: Direction, context: StoreContext) throws -> Bool {
    guard let existingIdentity = identities[address] else {
      return true // Trust on first use
    }
    return existingIdentity == identity
  }

  func identity(for address: ProtocolAddress, context: StoreContext) throws -> IdentityKey? {
    return identities[address]
  }

  func clear() {
    identities.removeAll()
  }
}

class InMemoryPreKeyStore: PreKeyStore {
  private var preKeys: [UInt32: PreKeyRecord] = [:]

  func loadPreKey(id: UInt32, context: StoreContext) throws -> PreKeyRecord {
    guard let record = preKeys[id] else {
      throw SignalError.invalidKeyIdentifier("PreKey not found: \(id)")
    }
    return record
  }

  func storePreKey(_ record: PreKeyRecord, id: UInt32, context: StoreContext) throws {
    preKeys[id] = record
  }

  func removePreKey(id: UInt32, context: StoreContext) throws {
    preKeys.removeValue(forKey: id)
  }

  func clear() {
    preKeys.removeAll()
  }
}

class InMemorySignedPreKeyStore: SignedPreKeyStore {
  private var signedPreKeys: [UInt32: SignedPreKeyRecord] = [:]

  func loadSignedPreKey(id: UInt32, context: StoreContext) throws -> SignedPreKeyRecord {
    guard let record = signedPreKeys[id] else {
      throw SignalError.invalidKeyIdentifier("SignedPreKey not found: \(id)")
    }
    return record
  }

  func storeSignedPreKey(_ record: SignedPreKeyRecord, id: UInt32, context: StoreContext) throws {
    signedPreKeys[id] = record
  }

  func clear() {
    signedPreKeys.removeAll()
  }
}

class InMemoryKyberPreKeyStore: KyberPreKeyStore {
  private var kyberPreKeys: [UInt32: KyberPreKeyRecord] = [:]

  func loadKyberPreKey(id: UInt32, context: StoreContext) throws -> KyberPreKeyRecord {
    guard let record = kyberPreKeys[id] else {
      throw SignalError.invalidKeyIdentifier("KyberPreKey not found: \(id)")
    }
    return record
  }

  func storeKyberPreKey(_ record: KyberPreKeyRecord, id: UInt32, context: StoreContext) throws {
    kyberPreKeys[id] = record
  }

  func markKyberPreKeyUsed(id: UInt32, signedPreKeyId: UInt32, baseKey: PublicKey, context: StoreContext) throws {
    // Mark as used - could remove or flag for rotation
    // The additional parameters help with key rotation tracking
  }

  func clear() {
    kyberPreKeys.removeAll()
  }
}