import Foundation
import Virtualization
import CryptoKit

struct VMDirectory: Prunable {
  enum State: String {
    case Running = "running"
    case Suspended = "suspended"
    case Stopped = "stopped"
  }

  var baseURL: URL

  var configURL: URL {
    baseURL.appendingPathComponent("config.json")
  }
  var diskURL: URL {
    baseURL.appendingPathComponent("disk.img")
  }
  var nvramURL: URL {
    baseURL.appendingPathComponent("nvram.bin")
  }
  var stateURL: URL {
    baseURL.appendingPathComponent("state.vzvmsave")
  }
  var manifestURL: URL {
    baseURL.appendingPathComponent("manifest.json")
  }
  var overlayURL: URL {
    baseURL.appendingPathComponent("overlay.asif")
  }
  var controlSocketURL: URL {
    URL(fileURLWithPath: "control.sock", relativeTo: baseURL)
  }

  var explicitlyPulledMark: URL {
    baseURL.appendingPathComponent(".explicitly-pulled")
  }

  var name: String {
    baseURL.lastPathComponent
  }

  var url: URL {
    baseURL
  }

  func lock() throws -> PIDLock {
    try PIDLock(lockURL: configURL)
  }

  func running() throws -> Bool {
    // The most common reason why PIDLock() instantiation fails is a race with "tart delete" (ENOENT),
    // which is fine to report as "not running".
    //
    // The other reasons are unlikely and the cost of getting a false positive is way less than
    // the cost of crashing with an exception when calling "tart list" on a busy machine, for example.
    guard let lock = try? lock() else {
      return false
    }

    return try lock.pid() != 0
  }

  func state() throws -> State {
    if try running() {
      return State.Running
    } else if FileManager.default.fileExists(atPath: stateURL.path) {
      return State.Suspended
    } else {
      return State.Stopped
    }
  }

  static func temporary() throws -> VMDirectory {
    let tmpDir = try Config().tartTmpDir.appendingPathComponent(UUID().uuidString)
    try FileManager.default.createDirectory(at: tmpDir, withIntermediateDirectories: false)

    return VMDirectory(baseURL: tmpDir)
  }

  //Create tmp directory with hashing
  static func temporaryDeterministic(key: String) throws -> VMDirectory {
    let keyData = Data(key.utf8)
    let hash = Insecure.MD5.hash(data: keyData)
    // Convert hash to string
    let hashString = hash.compactMap { String(format: "%02x", $0) }.joined()
    let tmpDir = try Config().tartTmpDir.appendingPathComponent(hashString)
    try FileManager.default.createDirectory(at: tmpDir, withIntermediateDirectories: true)
    return VMDirectory(baseURL: tmpDir)
  }

  private var hasRequiredMetadata: Bool {
    let fileManager = FileManager.default

    return fileManager.fileExists(atPath: configURL.path) &&
      fileManager.fileExists(atPath: nvramURL.path)
  }

  enum Layout: Equatable {
    /// Existing Tart layout with one independently attachable `disk.img`.
    /// A pulled standalone OCI record may also carry `manifest.json`.
    case standalone

    /// Runnable stacked VM with immutable disk files from `manifest.json` and
    /// a private writable `overlay.asif`.
    case stackedLocal

    /// Pulled OCI record for a stacked image. It intentionally has no writable
    /// overlay and becomes runnable only after `tart clone` creates one.
    case stackedOCIRecord

    var isRunnable: Bool {
      self != .stackedOCIRecord
    }
  }

  var layout: Layout? {
    let fileManager = FileManager.default
    let hasDisk = fileManager.fileExists(atPath: diskURL.path)
    let hasManifest = fileManager.fileExists(atPath: manifestURL.path)
    let hasOverlay = fileManager.fileExists(atPath: overlayURL.path)

    guard hasRequiredMetadata else {
      return nil
    }

    if hasDisk && !hasOverlay {
      return .standalone
    }
    if !hasDisk && hasManifest && hasOverlay {
      return .stackedLocal
    }
    if !hasDisk && hasManifest && !hasOverlay {
      return .stackedOCIRecord
    }

    return nil
  }

  var initialized: Bool {
    layout?.isRunnable == true
  }

  var isStandalone: Bool {
    layout == .standalone
  }

  var isStackedVM: Bool {
    layout == .stackedLocal
  }

  var isStackedCachedImage: Bool {
    layout == .stackedOCIRecord
  }

  /// Shapes that may live in the remote-image cache. A cached stacked image
  /// has no writable overlay and is intentionally not runnable as a local VM.
  var isCachedImage: Bool {
    layout == .standalone || layout == .stackedOCIRecord
  }

  func initialize(overwrite: Bool = false) throws {
    if !overwrite && initialized {
      throw RuntimeError.VMDirectoryAlreadyInitialized("VM directory is already initialized, preventing overwrite")
    }

    try FileManager.default.createDirectory(at: baseURL, withIntermediateDirectories: true, attributes: nil)

    try? FileManager.default.removeItem(at: configURL)
    try? FileManager.default.removeItem(at: diskURL)
    try? FileManager.default.removeItem(at: nvramURL)
    try? FileManager.default.removeItem(at: manifestURL)
    try? FileManager.default.removeItem(at: overlayURL)
    try? FileManager.default.removeItem(at: stateURL)
  }

  func validate(userFriendlyName: String) throws {
    if !FileManager.default.fileExists(atPath: baseURL.path) {
      throw RuntimeError.VMDoesNotExist(name: userFriendlyName)
    }

    if !initialized {
      throw RuntimeError.VMMissingFiles(
        "VM is missing files for a supported layout: "
          + "standalone requires \(configURL.lastPathComponent), \(diskURL.lastPathComponent) and \(nvramURL.lastPathComponent); "
          + "stacked requires \(configURL.lastPathComponent), \(manifestURL.lastPathComponent), "
          + "\(overlayURL.lastPathComponent) and \(nvramURL.lastPathComponent)"
      )
    }
  }

  func validateCachedImage(userFriendlyName: String) throws {
    if !FileManager.default.fileExists(atPath: baseURL.path) {
      throw RuntimeError.VMDoesNotExist(name: userFriendlyName)
    }

    if !isCachedImage {
      throw RuntimeError.VMMissingFiles(
        "cached image is missing files for a supported layout: "
          + "standalone requires \(configURL.lastPathComponent), \(diskURL.lastPathComponent) and \(nvramURL.lastPathComponent); "
          + "stacked requires \(configURL.lastPathComponent), \(manifestURL.lastPathComponent) and \(nvramURL.lastPathComponent)"
      )
    }
  }

  func clone(to: VMDirectory, generateMAC: Bool) throws {
    try FileManager.default.copyItem(at: configURL, to: to.configURL)
    try FileManager.default.copyItem(at: nvramURL, to: to.nvramURL)
    try FileManager.default.copyItem(at: diskURL, to: to.diskURL)
    try? FileManager.default.copyItem(at: stateURL, to: to.stateURL)

    // Re-generate MAC address
    if generateMAC {
      try to.regenerateMACAddress()
    }
  }

  func macAddress() throws -> String {
    try VMConfig(fromURL: configURL).macAddress.string
  }

  func regenerateMACAddress() throws {
    var vmConfig = try VMConfig(fromURL: configURL)

    vmConfig.macAddress = VZMACAddress.randomLocallyAdministered()
    // cleanup state if any
    try? FileManager.default.removeItem(at: stateURL)

    try vmConfig.save(toURL: configURL)
  }

  func resizeDisk(
    _ sizeGB: UInt16,
    format: DiskImageFormat = .raw,
    contentStore: ContentStore? = nil
  ) throws {
    if isStackedVM {
      // Resolve the stack before taking the config.json PID lock. Reading
      // config.json after acquiring an fcntl lock would release that lock
      // when the read file descriptor is closed.
      let stack = try diskImageStack(contentStore: contentStore)
      let lock = try lock()
      guard try lock.trylock() else {
        throw RuntimeError.VMConfigurationError("VM \"\(name)\" must be stopped before resizing its disk")
      }
      defer { try? lock.unlock() }

      // Holding the PID lock proves that the VM is not running. A saved state
      // file is the remaining suspended state that must also reject resize.
      guard !FileManager.default.fileExists(atPath: stateURL.path) else {
        throw RuntimeError.VMConfigurationError("VM \"\(name)\" must be stopped before resizing its disk")
      }

      let desiredSizeBytes = UInt64(sizeGB) * 1000 * 1000 * 1000
      guard desiredSizeBytes.isMultiple(of: stack.blockSize) else {
        throw RuntimeError.InvalidDiskSize("new disk size must align to the stacked disk block size")
      }

      let desiredBlockCount = desiredSizeBytes / stack.blockSize
      try stack.growWritableOverlay(toBlockCount: desiredBlockCount)
      return
    }

    let diskExists = FileManager.default.fileExists(atPath: diskURL.path)

    if diskExists {
      // Existing disk - resize it
      try resizeExistingDisk(sizeGB)
    } else {
      // New disk - create it with the specified format
      try createDisk(sizeGB: sizeGB, format: format)
    }
  }

  private func resizeExistingDisk(_ sizeGB: UInt16) throws {
    // Check if this is an ASIF disk by reading the VM config
    let vmConfig = try VMConfig(fromURL: configURL)

    if vmConfig.diskFormat == .asif {
      try resizeASIFDisk(sizeGB)
    } else {
      try resizeRawDisk(sizeGB)
    }
  }

  private func resizeRawDisk(_ sizeGB: UInt16) throws {
    let diskFileHandle = try FileHandle.init(forWritingTo: diskURL)
    let currentDiskFileLength = try diskFileHandle.seekToEnd()
    let desiredDiskFileLength = UInt64(sizeGB) * 1000 * 1000 * 1000

    if desiredDiskFileLength < currentDiskFileLength {
      let currentLengthHuman = ByteCountFormatter().string(fromByteCount: Int64(currentDiskFileLength))
      let desiredLengthHuman = ByteCountFormatter().string(fromByteCount: Int64(desiredDiskFileLength))
      throw RuntimeError.InvalidDiskSize("new disk size of \(desiredLengthHuman) should be larger " +
        "than the current disk size of \(currentLengthHuman)")
    } else if desiredDiskFileLength > currentDiskFileLength {
      try diskFileHandle.truncate(atOffset: desiredDiskFileLength)
    }
    try diskFileHandle.close()
  }

  private func resizeASIFDisk(_ sizeGB: UInt16) throws {
    do {
      let diskImageInfo = try Diskutil.imageInfo(diskURL)

      let currentSizeBytes = try diskImageInfo.totalBytes()
      let desiredSizeBytes = UInt64(sizeGB) * 1000 * 1000 * 1000

      if desiredSizeBytes < currentSizeBytes {
        let currentLengthHuman = ByteCountFormatter().string(fromByteCount: Int64(currentSizeBytes))
        let desiredLengthHuman = ByteCountFormatter().string(fromByteCount: Int64(desiredSizeBytes))

        throw RuntimeError.InvalidDiskSize("New disk size of \(desiredLengthHuman) should be larger " +
          "than the current disk size of \(currentLengthHuman)")
      } else if desiredSizeBytes > currentSizeBytes {
        // Resize the ASIF disk image using diskutil
        try performASIFResize(sizeGB)
      } else {
        // If sizes are equal, no action needed
      }
    } catch let error as RuntimeError {
      throw error
    } catch {
      throw RuntimeError.FailedToResizeDisk("\(error)")
    }
  }

  private func performASIFResize(_ sizeGB: UInt16) throws {
    guard let diskutilURL = resolveBinaryPath("diskutil") else {
      throw RuntimeError.FailedToResizeDisk("diskutil not found in PATH")
    }

    let process = Process()
    process.executableURL = diskutilURL
    process.arguments = [
      "image", "resize",
      "--size", "\(sizeGB)G",
      diskURL.path
    ]

    let pipe = Pipe()
    process.standardOutput = pipe
    process.standardError = pipe

    do {
      try process.run()
      process.waitUntilExit()

      let data = pipe.fileHandleForReading.readDataToEndOfFile()

      if process.terminationStatus != 0 {
        let output = String(data: data, encoding: .utf8) ?? "Unknown error"
        throw RuntimeError.FailedToResizeDisk("Failed to resize ASIF disk image: \(output)")
      }
    } catch {
      throw RuntimeError.FailedToResizeDisk("Failed to execute diskutil resize: \(error)")
    }
  }

  private func createDisk(sizeGB: UInt16, format: DiskImageFormat) throws {
    switch format {
    case .raw:
      try createRawDisk(sizeGB: sizeGB)
    case .asif:
      try Diskutil.imageCreate(diskURL: diskURL, sizeGB: sizeGB)
    }
  }

  private func createRawDisk(sizeGB: UInt16) throws {
    // Create traditional raw disk image
    FileManager.default.createFile(atPath: diskURL.path, contents: nil, attributes: nil)

    let diskFileHandle = try FileHandle.init(forWritingTo: diskURL)
    let desiredDiskFileLength = UInt64(sizeGB) * 1000 * 1000 * 1000
    try diskFileHandle.truncate(atOffset: desiredDiskFileLength)
    try diskFileHandle.close()
  }


  func delete() throws {
    let lock = try lock()

    if try !lock.trylock() {
      throw RuntimeError.VMIsRunning(name)
    }

    // Standalone local VMs do not reference the shared content store. Delete
    // them directly so a full disk can still be recovered before the content
    // store has ever been initialized.
    if isStandalone {
      try FileManager.default.removeItem(at: baseURL)
    } else {
      try removeFromDisk()
    }

    try lock.unlock()
  }

  /// Removes a VM directory while preserving the content-store reference
  /// protocol for any complete or partially published manifest it contains.
  func removeFromDisk() throws {
    let contentStore = try ContentStore()
    try contentStore.withPruneLock {
      try FileManager.default.removeItem(at: baseURL)
    }
  }

  func accessDate() throws -> Date {
    try baseURL.accessDate()
  }

  func allocatedSizeBytes() throws -> Int {
    try configURL.allocatedSizeBytes() + localDiskStorageAllocatedSizeBytes() + nvramURL.allocatedSizeBytes()
  }

  func allocatedSizeGB() throws -> Int {
    try allocatedSizeBytes() / 1000 / 1000 / 1000
  }

  func deduplicatedSizeBytes() throws -> Int {
    try configURL.deduplicatedSizeBytes() + localDiskStorageDeduplicatedSizeBytes() + nvramURL.deduplicatedSizeBytes()
  }

  func deduplicatedSizeGB() throws -> Int {
    try deduplicatedSizeBytes() / 1000 / 1000 / 1000
  }

  func sizeBytes() throws -> Int {
    try configURL.sizeBytes() + localDiskStorageSizeBytes() + nvramURL.sizeBytes()
  }

  func sizeGB() throws -> Int {
    try sizeBytes() / 1000 / 1000 / 1000
  }

  func diskSizeBytes() throws -> Int {
    if isStackedVM {
      let blockLayout = try DiskImageStack.diskImageBlockLayout(at: overlayURL)
      let product = blockLayout.blockSize.multipliedReportingOverflow(by: blockLayout.blockCount)
      guard !product.overflow, let diskSizeBytes = Int(exactly: product.partialValue) else {
        throw RuntimeError.VMConfigurationError("VM has invalid stacked disk block layout")
      }

      return diskSizeBytes
    }

    if isStackedCachedImage {
      let manifest = try OCIManifest(fromJSON: Data(contentsOf: manifestURL))
      guard let blockSize = manifest.diskBlockSize(),
            let blockCount = manifest.diskBlockCount() else {
        throw RuntimeError.VMConfigurationError("VM has invalid stacked disk block layout")
      }
      let product = blockSize.multipliedReportingOverflow(by: blockCount)
      guard !product.overflow, let diskSizeBytes = Int(exactly: product.partialValue) else {
        throw RuntimeError.VMConfigurationError("VM has invalid stacked disk block layout")
      }

      return diskSizeBytes
    }

    let vmConfig = try VMConfig(fromURL: configURL)

    return switch vmConfig.diskFormat {
    case .raw:
      try sizeBytes()
    case .asif:
      try Diskutil.imageInfo(diskURL).totalBytes()
    }
  }

  func diskSizeGB() throws -> Int {
    try diskSizeBytes() / 1000 / 1000 / 1000
  }

  func markExplicitlyPulled() {
    FileManager.default.createFile(atPath: explicitlyPulledMark.path, contents: nil)
  }

  func isExplicitlyPulled() -> Bool {
    FileManager.default.fileExists(atPath: explicitlyPulledMark.path)
  }

  private var localDiskStorageURL: URL {
    isStackedVM ? overlayURL : diskURL
  }

  // Cached stacked images own no disk file in their VM directory. Their
  // immutable disk content lives in the shared content store and must not be
  // charged to every cached image that references it.
  private func localDiskStorageAllocatedSizeBytes() throws -> Int {
    isStackedCachedImage ? 0 : try localDiskStorageURL.allocatedSizeBytes()
  }

  private func localDiskStorageDeduplicatedSizeBytes() throws -> Int {
    isStackedCachedImage ? 0 : try localDiskStorageURL.deduplicatedSizeBytes()
  }

  private func localDiskStorageSizeBytes() throws -> Int {
    isStackedCachedImage ? 0 : try localDiskStorageURL.sizeBytes()
  }
}
