import Foundation
@preconcurrency import MediaGenerationKit
import UniformTypeIdentifiers

/// 기기 안 그림 엔진(2026-10-22): Draw Things의 MediaGenerationKit(LGPLv3)으로 그린다.
///
/// 이 모듈은 앱만 링크하는 동적 프레임워크 `MDitImaging`에 든다(LGPLv3 4조: 사용자가 고친 라이브러리로 이 프레임워크를
/// 다시 빌드해 바꿔 끼울 수 있게, docs/notes/img/license.md). 다른 MDit 모듈에 기대지 않고 Foundation 형만 내놓는다:
/// 앱이 켤 때 이것을 Features의 `DeviceImageBackend`에 잇는다.
///
/// 모델 파일은 고른 폴더(`directory`)에만 둔다. 모델 목록과 설정은 엔진에 들어 있는 것과 앱 묶음의 `configs.json`으로 찾아,
/// 설정을 받으러 밖에 나가지 않는다. 받는 곳은 모델 파일뿐이다(static.libnnc.org, 사람이 "받기"를 눌렀을 때만).
public final class LocalImageEngine: Sendable {
  /// 그리는 동안의 걸음.
  public enum Step: Sendable, Equatable {
    case preparing
    case step(Int, of: Int)
    case finishing
  }

  /// 엔진이 아는 모델 하나: MDit의 모델 id(`ImageModels`)와 엔진의 파일, 권장 설정.
  struct Model: Sendable {
    let id: String
    let file: String
    /// 함께 쓰는 파일(글 인코더, VAE). 다른 모델도 쓰면 지울 때 남긴다.
    let shared: [String]
    let steps: Int
    let sampler: SamplerType
    let guidance: Float
    let shift: Float
    /// 참조 그림을 받는가.
    let takesReferences: Bool
    /// 받을 크기(바이트, 받기 전에 남은 양을 어림하는 데만 쓴다: 실제 크기는 받으면서 엔진이 알려 준다).
    let bytes: [String: Int64]
  }

  /// 이 앱이 쓰는 모델들. 설정은 Draw Things의 권장값(docs/notes/img/spike.md): klein은 4걸음, DDIM Trailing, shift 3.
  static let models: [String: Model] = {
    let kleinShared = ["qwen_3_4b_q8p.ckpt", "flux_2_vae_f16.ckpt"]
    let kleinBytes: [String: Int64] = ["qwen_3_4b_q8p.ckpt": 4_535_328_768, "flux_2_vae_f16.ckpt": 168_534_016]
    let list = [
      Model(id: "flux2-klein-4b-8bit", file: "flux_2_klein_4b_q8p.ckpt", shared: kleinShared, steps: 4, sampler: .dDIMTrailing,
            guidance: 1, shift: 3, takesReferences: true, bytes: kleinBytes.merging(["flux_2_klein_4b_q8p.ckpt": 4_039_196_672]) { a, _ in a }),
      Model(id: "flux2-klein-4b-6bit", file: "flux_2_klein_4b_q6p.ckpt", shared: kleinShared, steps: 4, sampler: .dDIMTrailing,
            guidance: 1, shift: 3, takesReferences: true, bytes: kleinBytes.merging(["flux_2_klein_4b_q6p.ckpt": 3_117_088_768]) { a, _ in a }),
      Model(id: "z-image-turbo", file: "z_image_turbo_1.0_q8p.ckpt", shared: ["qwen_3_vl_4b_instruct_q8p.ckpt", "flux_1_vae_f16.ckpt"],
            steps: 8, sampler: .uniPCTrailing, guidance: 1, shift: 3, takesReferences: false, bytes: [:]),
    ]
    return Dictionary(uniqueKeysWithValues: list.map { ($0.id, $0) })
  }()

  /// 이 엔진이 그릴 수 있는 모델 id들.
  public static var modelIDs: [String] { models.keys.sorted() }

  /// 모델 파일을 두는 폴더.
  public let directory: URL
  private let runner = Runner()

  /// - Parameter directory: 모델 폴더(없으면 만들고, 백업에서 뺀다: 다시 받을 수 있는 큰 파일이다).
  public init(directory: URL) throws {
    try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)
    var excluded = directory
    var values = URLResourceValues()
    values.isExcludedFromBackup = true
    try? excluded.setResourceValues(values)
    self.directory = directory
    // 엔진은 모델을 찾는 폴더를 프로세스 전체에 하나 둔다(기본값은 문서 폴더 › Models라 바꾼다).
    MediaGenerationEnvironment.default.externalUrls = [directory]
  }

  // MARK: 모델

  /// 그 모델과 함께 쓰는 파일이 모두 받아져 있는가(받다 만 것이 남아 있으면 아니다).
  public func isInstalled(_ id: String) -> Bool {
    guard let model = Self.models[id] else { return false }
    return ([model.file] + model.shared).allSatisfy(isComplete)
  }

  /// 아직 받을 바이트(어림). 모르면 nil.
  public func remainingBytes(_ id: String) -> Int64? {
    guard let model = Self.models[id], !model.bytes.isEmpty else { return nil }
    return ([model.file] + model.shared).filter { !isComplete($0) }.reduce(0) { $0 + (model.bytes[$1] ?? 0) }
  }

  /// 모델과 함께 쓰는 파일을 받는다(받다 만 파일은 이어 받는다). 진행은 (받은 바이트, 모두)로 알린다.
  public func download(_ id: String, progress: @escaping @Sendable (Int64, Int64) -> Void) async throws {
    guard let model = Self.models[id] else { throw Failure.unknownModel(id) }
    let pending = ([model.file] + model.shared).filter { !isComplete($0) }
    let tally = DownloadTally(estimates: pending.map { model.bytes[$0] ?? 0 })
    _ = try await MediaGenerationEnvironment.default.ensure(model.file, offline: true) { state in
      guard case .downloading(_, let index, _, let written, let total) = state else { return }
      let (done, all) = tally.update(index: index, written: written, total: total)
      progress(done, all)
    }
    guard isInstalled(id) else { throw Failure.incomplete(id) }
  }

  /// 받은 모델을 지운다. 함께 쓰는 파일은 그것을 쓰는 다른 모델이 받아져 있으면 남긴다.
  public func remove(_ id: String) async throws {
    guard let model = Self.models[id] else { throw Failure.unknownModel(id) }
    await runner.release()
    let others = Self.models.values.filter { $0.id != id && isInstalled($0.id) }
    let keep = Set(others.flatMap(\.shared))
    for file in [model.file] + model.shared where !keep.contains(file) {
      for name in [file, file + "-tensordata", file + ".part"] {
        let url = directory.appendingPathComponent(name)
        if FileManager.default.fileExists(atPath: url.path) { try FileManager.default.removeItem(at: url) }
      }
    }
  }

  private func isComplete(_ file: String) -> Bool {
    let url = directory.appendingPathComponent(file)
    return FileManager.default.fileExists(atPath: url.path)
      && !FileManager.default.fileExists(atPath: url.appendingPathExtension("part").path)
  }

  // MARK: 그리기

  /// 그림 한 장을 그려 PNG로 돌려준다. 취소되면 `CancellationError`.
  /// - Parameters:
  ///   - references: 참조 그림(PNG). 받지 않는 모델은 버린다.
  public func draw(model id: String, prompt: String, negative: String, width: Int, height: Int, seed: UInt32,
                   references: [Data], progress: @escaping @Sendable (Step) -> Void) async throws -> Data {
    guard let model = Self.models[id] else { throw Failure.unknownModel(id) }
    guard isInstalled(id) else { throw Failure.notInstalled(id) }
    return try await runner.draw(model: model, directory: directory, prompt: prompt, negative: negative,
                                 width: width, height: height, seed: seed, references: model.takesReferences ? references : [],
                                 progress: progress)
  }

  /// 올려 둔 모델을 내린다(메모리를 돌려준다). 다음 장은 다시 올리느라 조금 더 걸린다.
  public func release() async { await runner.release() }

  public enum Failure: Error, CustomStringConvertible {
    case unknownModel(String)
    case notInstalled(String)
    case incomplete(String)
    case noImage

    public var description: String {
      switch self {
      case .unknownModel(let id): "unknown model \(id)"
      case .notInstalled(let id): "model not installed \(id)"
      case .incomplete(let id): "download incomplete \(id)"
      case .noImage: "no image"
      }
    }
  }

  /// 파이프라인(올린 모델)을 쥐는 곳. 한 번에 한 장만 그린다(부르는 쪽이 차례를 지킨다). 같은 모델로 잇달아 그리면 다시 올리지 않는다.
  private actor Runner {
    private var loaded: (id: String, pipeline: MediaGenerationPipeline)?

    func draw(model: Model, directory: URL, prompt: String, negative: String, width: Int, height: Int, seed: UInt32,
              references: [Data], progress: @escaping @Sendable (Step) -> Void) async throws -> Data {
      progress(.preparing)
      var pipeline: MediaGenerationPipeline
      if let loaded, loaded.id == model.id {
        pipeline = loaded.pipeline
      } else {
        loaded = nil
        pipeline = try await MediaGenerationPipeline.fromPretrained(model.file, backend: .local(directory: directory.path))
      }
      try Task.checkCancellation()
      // 권장 설정을 앱이 직접 준다(설정을 받으러 밖에 나가지 못해도 같은 그림이 나오게).
      pipeline.configuration.width = width
      pipeline.configuration.height = height
      pipeline.configuration.seed = seed
      pipeline.configuration.steps = model.steps
      pipeline.configuration.sampler = model.sampler
      pipeline.configuration.guidanceScale = model.guidance
      pipeline.configuration.shift = model.shift
      pipeline.configuration.resolutionDependentShift = false
      pipeline.configuration.batchCount = 1
      pipeline.configuration.batchSize = 1
      pipeline.configuration.strength = 1
      loaded = (model.id, pipeline)

      let inputs: [MediaGenerationPipeline.Input] = references.map { MediaGenerationPipeline.data($0).moodboard() }
      let results = try await pipeline.generate(prompt: prompt, negativePrompt: negative, inputs: inputs) { state in
        switch state {
        case .generating(let step, let total): progress(.step(step, of: total))
        case .decoding, .postprocessing: progress(.finishing)
        default: break
        }
      }
      guard let first = results.first else { throw Failure.noImage }
      // 결과는 파일로만 적을 수 있어 임시 파일을 거친다.
      let file = FileManager.default.temporaryDirectory.appendingPathComponent("mdit-image-\(UUID().uuidString).png")
      defer { try? FileManager.default.removeItem(at: file) }
      try first.write(to: file, type: .png)
      return try Data(contentsOf: file)
    }

    func release() { loaded = nil }
  }
}

/// 여러 파일을 차례로 받을 때의 모두 합한 진행. 아직 시작하지 않은 파일은 어림한 크기로 센다.
private final class DownloadTally: @unchecked Sendable {
  private let lock = NSLock()
  private var totals: [Int64]
  private var written: [Int64]

  init(estimates: [Int64]) {
    totals = estimates
    written = Array(repeating: 0, count: estimates.count)
  }

  /// - Parameter index: 받는 파일의 차례(1부터).
  func update(index: Int, written bytes: Int64, total: Int64) -> (Int64, Int64) {
    lock.lock(); defer { lock.unlock() }
    let i = index - 1
    if i >= totals.count {
      totals.append(contentsOf: Array(repeating: 0, count: i - totals.count + 1))
      written.append(contentsOf: Array(repeating: 0, count: i - written.count + 1))
    }
    // 앞의 파일은 다 받은 것이다.
    for j in 0..<i where written[j] < totals[j] { written[j] = totals[j] }
    if total > 0 { totals[i] = total }
    written[i] = bytes
    return (written.reduce(0, +), max(totals.reduce(0, +), 1))
  }
}