import CryptoKit
import Foundation
import RastrilloNative
import Security

/// A single Keychain item makes key rotation crash recoverable: retain the old
/// key until every file has been atomically resealed for its original path.
public struct StoreKeyState: Codable, Sendable {
  public var current: Data
  public var previous: Data?
  public var seals: UInt64
  public init(current: Data, previous: Data? = nil, seals: UInt64 = 0) {
    self.current = current; self.previous = previous; self.seals = seals
  }
}

/// Where an account's store key lives. `KeychainKeyProvider` is the real one;
/// tests script every Keychain status with `ScriptedKeyProvider`.
///
/// It answers in raw Keychain statuses on purpose: telling *unavailable*
/// from *lost* is the store's job, and it must see the status to do it.
public protocol KeyProvider: Sendable {
  /// `errSecSuccess` with the key, or the status the Keychain gave.
  func copyKey() -> (OSStatus, SymmetricKey?)
  func addKey(_ key: SymmetricKey) -> OSStatus
  func deleteKey() -> OSStatus
  func sealState() -> (OSStatus, StoreKeyState?)
  func saveSealState(_ state: StoreKeyState) -> OSStatus
}

/// One AES-256 key per account, a generic password under service
/// `<bundle id>.store` and account `store:<account id>|<origin>`, readable
/// only while the device is unlocked, on this device, and never synchronised.
public struct KeychainKeyProvider: KeyProvider {
  public let service: String
  public let account: String

  public init(accountID: String, origin: URL, bundleID: String = Bundle.main.bundleIdentifier ?? "oficina.native") {
    service = bundleID + ".store"
    account = "store:\(accountID)|\(origin.absoluteString)"
  }

  private var query: [String: Any] {
    [kSecClass as String: kSecClassGenericPassword, kSecAttrService as String: service,
     kSecAttrAccount as String: account, kSecAttrSynchronizable as String: false]
  }

  public func copyKey() -> (OSStatus, SymmetricKey?) {
    var q = query
    q[kSecReturnData as String] = true
    q[kSecMatchLimit as String] = kSecMatchLimitOne
    var result: CFTypeRef?
    let status = SecItemCopyMatching(q as CFDictionary, &result)
    guard status == errSecSuccess else { return (status, nil) }
    guard let data = result as? Data else { return (errSecDecode, nil) }
    if data.count == 32 { return (status, SymmetricKey(data: data)) }
    guard let state = try? JSONDecoder().decode(StoreKeyState.self, from: data), state.current.count == 32 else { return (errSecDecode, nil) }
    return (status, SymmetricKey(data: state.current))
  }

  public func addKey(_ key: SymmetricKey) -> OSStatus {
    var q = query
    q[kSecValueData as String] = key.withUnsafeBytes { Data($0) }
    q[kSecAttrAccessible as String] = kSecAttrAccessibleWhenUnlockedThisDeviceOnly
    return SecItemAdd(q as CFDictionary, nil)
  }

  public func sealState() -> (OSStatus, StoreKeyState?) {
    var q = query; q[kSecReturnData as String] = true; q[kSecMatchLimit as String] = kSecMatchLimitOne
    var result: CFTypeRef?
    let status = SecItemCopyMatching(q as CFDictionary, &result)
    guard status == errSecSuccess, let data = result as? Data else { return (status, nil) }
    if data.count == 32 { return (status, StoreKeyState(current: data)) }
    guard let state = try? JSONDecoder().decode(StoreKeyState.self, from: data), state.current.count == 32,
          state.previous == nil || state.previous?.count == 32 else { return (errSecDecode, nil) }
    return (status, state)
  }
  public func saveSealState(_ state: StoreKeyState) -> OSStatus {
    guard let bytes = try? JSONEncoder().encode(state) else { return errSecParam }
    return SecItemUpdate(query as CFDictionary, [kSecValueData as String: bytes] as CFDictionary)
  }

  public func deleteKey() -> OSStatus { SecItemDelete(query as CFDictionary) }
}

/// Atomic authenticated files, bound to their final relative paths. A lock
/// covers key creation and disk mutations; no plaintext is written to disk.
public final class NativeSealedStore: BlobStore, @unchecked Sendable {
  public static let formatVersion = 1
  public let root: URL
  public let accountID: String
  private let keys: KeyProvider
  private let lock = NSLock()
  private let fm = FileManager.default
  // One store owns an account directory. Cached material remains protected on
  // disk by completeFileProtection; cold opens still distinguish locked/lost.
  private var cachedState: StoreKeyState?
  private var cachedKey: SymmetricKey?
  private var reservedUntil: UInt64 = 0
  private var backupExcluded = false
  var rotationLimit: UInt64 = 1 << 31
  public init(root: URL, accountID: String, keys: KeyProvider) { self.root = root; self.accountID = accountID; self.keys = keys }
  static let writeOptions: Data.WritingOptions = [.atomic, .completeFileProtection]
  private func url(_ path: String, allowRoot: Bool = false) throws -> URL {
    guard (allowRoot || !path.isEmpty), !path.hasPrefix("/"), !path.contains("\0"),
          !path.split(separator: "/", omittingEmptySubsequences: false).contains(where: { $0 == ".." || $0 == "." }) else { throw BlobStoreError.unavailable }
    let result = root.appendingPathComponent(path)
    // Reject symlinks too: a relative spelling alone does not contain I/O.
    let resolvedRoot = root.resolvingSymlinksInPath().standardizedFileURL.path
    let resolved = result.resolvingSymlinksInPath().standardizedFileURL.path
    guard resolved == resolvedRoot || resolved.hasPrefix(resolvedRoot + "/") else { throw BlobStoreError.unavailable }
    return result
  }
  private func missing(_ error: Error) -> Bool {
    let error = error as NSError
    return error.domain == NSCocoaErrorDomain && [NSFileReadNoSuchFileError, NSFileNoSuchFileError].contains(error.code)
  }
  private func io<T>(_ body: () throws -> T) throws -> T {
    do { return try body() } catch {
      // A failed write may have committed its reservation. Abandon its entire
      // unused tail before retrying, including interrupted key rotation.
      cachedState = nil; cachedKey = nil; reservedUntil = 0; backupExcluded = false
      throw (error as? BlobStoreError) ?? .unavailable
    }
  }
  private func hasFiles(_ directory: URL) throws -> Bool {
    let children: [URL]
    do { children = try fm.contentsOfDirectory(at: directory, includingPropertiesForKeys: [.isDirectoryKey]) }
    catch { if missing(error) { return false }; throw error }
    for child in children {
      if try child.resourceValues(forKeys: [.isDirectoryKey]).isDirectory == true {
        if try hasFiles(child) { return true }
      } else { return true }
    }
    return false
  }
  private func key(create: Bool) throws -> SymmetricKey? {
    if let cachedKey { return cachedKey }
    let (status, key) = keys.copyKey()
    if status == errSecSuccess, let key { return key }
    guard status == errSecItemNotFound else { throw BlobStoreError.unavailable }
    if try hasFiles(root) {
      try fm.removeItem(at: root)
      throw BlobStoreError.keyLost
    }
    guard create else { return nil }
    let new = SymmetricKey(size: .bits256)
    let added = keys.addKey(new)
    if added == errSecSuccess { return new }
    if added == errSecDuplicateItem {
      let (status, existing) = keys.copyKey()
      if status == errSecSuccess, let existing { return existing }
    }
    throw BlobStoreError.unavailable
  }
  private func aad(_ path: String) -> Data { Data("\(accountID)|\(Self.formatVersion)|\(path)".utf8) }
  private func readKeyState() throws -> StoreKeyState {
    if let cachedState { return cachedState }
    let (status, state) = keys.sealState()
    guard status == errSecSuccess, let state else { throw BlobStoreError.unavailable }
    cachedState = state; cachedKey = SymmetricKey(data: state.current)
    reservedUntil = state.seals
    return state
  }
  private func saveKeyState(_ state: StoreKeyState) throws {
    var durable = state
    durable.seals = max(state.seals, reservedUntil)
    guard keys.saveSealState(durable) == errSecSuccess else { throw BlobStoreError.unavailable }
    cachedState = state; cachedKey = SymmetricKey(data: state.current)
  }
  private func reserveSeal(_ state: inout StoreKeyState) throws {
    guard state.seals < rotationLimit else { throw BlobStoreError.unavailable }
    if state.seals == reservedUntil {
      // Persist a high-water mark before using any slot. Relaunch starts at
      // that mark, sacrificing the unused tail rather than reusing a quota.
      reservedUntil = state.seals + min(1024, rotationLimit - state.seals)
      try saveKeyState(state)
    }
    state.seals += 1
    cachedState = state
  }
  private func open(_ data: Data, path: String, state: StoreKeyState) throws -> Data {
    let box = try AES.GCM.SealedBox(combined: data)
    do { return try AES.GCM.open(box, using: SymmetricKey(data: state.current), authenticating: aad(path)) }
    catch {
      guard let previous = state.previous else { throw error }
      return try AES.GCM.open(box, using: SymmetricKey(data: previous), authenticating: aad(path))
    }
  }
  private func files(_ directory: URL) throws -> [URL] {
    let children: [URL]
    do { children = try fm.contentsOfDirectory(at: directory, includingPropertiesForKeys: [.isDirectoryKey]) }
    catch { if missing(error) { return [] }; throw error }
    return try children.flatMap { child in
      try child.resourceValues(forKeys: [.isDirectoryKey]).isDirectory == true ? files(child) : [child]
    }
  }
  private func rotateIfNeeded(_ state: inout StoreKeyState) throws {
    guard state.previous != nil || state.seals >= rotationLimit else { return }
    if state.previous == nil {
      state = StoreKeyState(current: SymmetricKey(size: .bits256).withUnsafeBytes { Data($0) }, previous: state.current)
      reservedUntil = 0
      try saveKeyState(state)
    }
    for file in try files(root) {
      let path = String(file.resolvingSymlinksInPath().path.dropFirst(root.resolvingSymlinksInPath().path.count + 1))
      let combined = try Data(contentsOf: file)
      if let box = try? AES.GCM.SealedBox(combined: combined),
         (try? AES.GCM.open(box, using: SymmetricKey(data: state.current), authenticating: aad(path))) != nil { continue }
      let plaintext: Data
      do { plaintext = try open(combined, path: path, state: state) }
      catch { try fm.removeItem(at: file); throw BlobStoreError.corrupt(path) }
      try reserveSeal(&state)
      let resealed = try AES.GCM.seal(plaintext, using: SymmetricKey(data: state.current), authenticating: aad(path))
      try resealed.combined!.write(to: file, options: Self.writeOptions)
    }
    state.previous = nil
    try saveKeyState(state)
  }
  public func write(_ data: Data, to path: String) throws {
    try lock.withLock { try io {
      let target = try url(path)
      _ = try key(create: true)
      var state = try readKeyState()
      try rotateIfNeeded(&state)
      try reserveSeal(&state)
      let key = cachedKey!
      try fm.createDirectory(at: target.deletingLastPathComponent(), withIntermediateDirectories: true)
      if !backupExcluded {
        var directory = root; var values = URLResourceValues(); values.isExcludedFromBackup = true
        try directory.setResourceValues(values); backupExcluded = true
      }
      let sealed = try AES.GCM.seal(data, using: key, authenticating: aad(path))
      guard let combined = sealed.combined else { throw BlobStoreError.unavailable }
      try combined.write(to: target, options: Self.writeOptions)
    } }
  }
  public func read(_ path: String) throws -> Data? {
    try lock.withLock { try io {
      let target = try url(path)
      let data: Data
      do { data = try Data(contentsOf: target) } catch { if missing(error) { return nil }; throw error }
      guard try key(create: false) != nil else { throw BlobStoreError.unavailable }
      let state = try readKeyState()
      do { return try open(data, path: path, state: state) }
      catch {
        // Only an authentication/decode failure reaches this catch. I/O or
        // Keychain refusal must never turn into a destructive corruption path.
        try fm.removeItem(at: target)
        throw BlobStoreError.corrupt(path)
      }
    } }
  }
  public func list(_ directory: String) throws -> [String] {
    try lock.withLock { try io {
      let target = try url(directory, allowRoot: true)
      do {
        return try fm.contentsOfDirectory(at: target, includingPropertiesForKeys: [.isRegularFileKey])
          .filter { try $0.resourceValues(forKeys: [.isRegularFileKey]).isRegularFile == true }.map(\.lastPathComponent).sorted()
      } catch { if missing(error) { return [] }; throw error }
    } }
  }
  public func remove(_ path: String) throws {
    try lock.withLock { try io {
      let target = try url(path)
      do { try fm.removeItem(at: target) } catch { if !missing(error) { throw error } }
    } }
  }
  public func wipe() throws {
    try lock.withLock { try io {
      let status = keys.deleteKey()
      guard status == errSecSuccess || status == errSecItemNotFound else { throw BlobStoreError.unavailable }
      cachedState = nil; cachedKey = nil; reservedUntil = 0; backupExcluded = false
      do { try fm.removeItem(at: root) } catch { if !missing(error) { throw error } }
    } }
  }
}