ShadowsocksX-NG/ShadowsocksX-NG2/Domain/Rules/Custom/CustomRuleStore.swift

import Foundation

enum CustomRuleStoreError: Error, Equatable, Sendable {
  case missing
  case corrupt(detail: String)
  case schemaVersionMismatch(found: Int, expected: Int)
  case ioFailure(detail: String)
}

/// 自定义规则持久化：独立 JSON 文档，原子写盘。
/// Saving intent is independent of subsequent ACL application and runtime recovery.
struct CustomRuleStore {
  static let currentSchemaVersion = 3

  let fileURL: URL

  init(fileURL: URL = CustomRuleStore.defaultFileURL()) {
    self.fileURL = fileURL
  }

  static func defaultFileURL() -> URL {
    ApplicationSupportPaths.current.customRulesFileURL
  }

  /// 读取规则集合；文件不存在视为无自定义规则。
  func load() throws -> [CustomRule] { try loadDocument().rules }

  func loadDocument() throws -> CustomRuleDocument {
    let data: Data
    do {
      data = try Data(contentsOf: fileURL)
    } catch let error as CocoaError where error.code == .fileReadNoSuchFile {
      return CustomRuleDocument(rules: [])
    } catch {
      throw CustomRuleStoreError.ioFailure(detail: String(describing: error))
    }
    let document: CustomRuleDocument
    do {
      document = try Self.decode(data)
    } catch let error as CustomRuleStoreError {
      throw error
    } catch {
      throw CustomRuleStoreError.corrupt(detail: String(describing: error))
    }
    return document
  }

  /// 原子写入完整规则集合。
  func save(_ rules: [CustomRule]) throws {
    let old = try loadDocument()
    try saveDocument(CustomRuleDocument(rules: rules, disabledIdentities: old.disabledIdentities))
  }

  func saveDocument(_ document: CustomRuleDocument) throws {
    let data = try Self.encode(document)
    _ = try Self.decode(data)
    do {
      try AtomicFileWriter.write(data, to: fileURL)
    } catch {
      throw CustomRuleStoreError.ioFailure(detail: String(describing: error))
    }
  }

  /// 安全摘要：数量 + 内容版本。
  func summary() throws -> CustomRuleSummary {
    CustomRuleSummary.summarizing(try loadDocument())
  }

  static func decode(_ data: Data) throws -> CustomRuleDocument {
    try JSONDecoder().decode(StoredCustomRuleDocument.self, from: data).domain()
  }

  static func encode(_ document: CustomRuleDocument) throws -> Data {
    let encoder = JSONEncoder()
    encoder.outputFormatting = [.prettyPrinted, .sortedKeys]
    return try encoder.encode(StoredCustomRuleDocument(document))
  }
}

/// One saved-document owner per app process. Only successful atomic saves update
/// this fact; runtime recovery never restores an older rule document.
@MainActor
final class RuleDocumentSession {
  private let store: CustomRuleStore
  private(set) var current: CustomRuleDocument?
  private(set) var revision = 0

  init(store: CustomRuleStore) { self.store = store }

  func load() throws -> CustomRuleDocument {
    if let current { return current }
    let loaded = try store.loadDocument()
    current = loaded
    return loaded
  }

  func save(_ document: CustomRuleDocument) throws {
    try store.saveDocument(document)
    if current != document { revision += 1 }
    current = document
  }
}

private struct StoredCustomRule: Codable {
  let id: UUID
  let rule: StoredProxyRule

  init(_ custom: CustomRule) {
    id = custom.id
    rule = StoredProxyRule(custom.rule)
  }

  private enum CodingKeys: CodingKey { case id }

  init(from decoder: Decoder) throws {
    id = try decoder.container(keyedBy: CodingKeys.self).decode(UUID.self, forKey: .id)
    rule = try StoredProxyRule(from: decoder)
  }

  func encode(to encoder: Encoder) throws {
    var values = encoder.container(keyedBy: CodingKeys.self)
    try values.encode(id, forKey: .id)
    try rule.encode(to: encoder)
  }

  func domain() throws -> CustomRule {
    let value = try rule.domain()
    guard !value.match.isFixedLocalSuffix else {
      throw CustomRuleStoreError.corrupt(detail: "Fixed local suffix is not a custom rule")
    }
    return CustomRule(id: id, action: value.action, match: value.match)
  }
}

private struct StoredCustomRuleDocument: Codable {
  let schemaVersion: Int
  let rules: [StoredCustomRule]
  let disabledIdentities: [StoredProxyRule]

  init(_ document: CustomRuleDocument) {
    schemaVersion = CustomRuleStore.currentSchemaVersion
    rules = document.rules.map(StoredCustomRule.init)
    disabledIdentities = document.disabledIdentities.sorted { $0.contentToken < $1.contentToken }
      .map(StoredProxyRule.init)
  }

  private enum CodingKeys: CodingKey { case schemaVersion, rules, disabledIdentities }

  init(from decoder: Decoder) throws {
    let values = try decoder.container(keyedBy: CodingKeys.self)
    schemaVersion = try values.decode(Int.self, forKey: .schemaVersion)
    guard [1, 2, CustomRuleStore.currentSchemaVersion].contains(schemaVersion) else {
      throw CustomRuleStoreError.schemaVersionMismatch(
        found: schemaVersion, expected: CustomRuleStore.currentSchemaVersion)
    }
    rules = try values.decode([StoredCustomRule].self, forKey: .rules)
    disabledIdentities =
      schemaVersion == 1
      ? [] : try values.decode([StoredProxyRule].self, forKey: .disabledIdentities)
  }

  func domain() throws -> CustomRuleDocument {
    let custom = try rules.map { try $0.domain() }
    let disabled = try disabledIdentities.map { try $0.domain() }
    guard Set(custom.map(\.id)).count == custom.count,
      Set(custom.map(\.rule)).count == custom.count,
      Set(disabled).count == disabled.count
    else { throw CustomRuleStoreError.corrupt(detail: "Duplicate custom rule records") }
    return CustomRuleDocument(rules: custom, disabledIdentities: Set(disabled))
  }
}