@preconcurrency import CoreML
import Foundation
@preconcurrency import Tokenizers

struct DolphinModelInformation: Sendable, Equatable {
  let loadSeconds: Double
  let maxContextLength: Int
  let maxQueryLength: Int
  let computeUnits: String
  let modelIdentifier: String
}

struct DolphinGenerationProgress: Sendable, Equatable {
  let text: String
  let generatedTokens: Int
}

struct DolphinGenerationResult: Sendable, Equatable {
  let text: String
  let promptTokens: Int
  let generatedTokens: Int
  let timeToFirstTokenSeconds: Double
  let totalSeconds: Double

  var tokensPerSecond: Double {
    guard generatedTokens > 0, totalSeconds > 0 else { return 0 }
    return Double(generatedTokens) / totalSeconds
  }
}

enum DolphinModelRuntimeError: LocalizedError, Sendable, Equatable {
  case modelNotLoaded
  case modelResourceMissing
  case tokenizerResourcesMissing
  case emptyPrompt
  case invalidOutputLimit
  case contextTooLong(actual: Int, limit: Int)
  case generationAlreadyRunning
  case missingLogits
  case unexpectedLogitsShape([Int])
  case nonFiniteLogits
  case unsupportedLogitsType(String)
  case invalidCausalMask

  var errorDescription: String? {
    switch self {
    case .modelNotLoaded:
      "Load the Core ML model before asking Dolphin for generated text."
    case .modelResourceMissing:
      "The compiled Dolphin model is missing from the app bundle. Run prepare-model.sh, regenerate the project, and rebuild."
    case .tokenizerResourcesMissing:
      "The tokenizer resources are missing from the app bundle."
    case .emptyPrompt:
      "The tokenizer produced an empty prompt."
    case .invalidOutputLimit:
      "The output-token limit must be greater than zero."
    case .contextTooLong(let actual, let limit):
      "Prompt plus output requests \(actual) tokens; the state capacity is \(limit)."
    case .generationAlreadyRunning:
      "A generation is already running."
    case .missingLogits:
      "The model prediction did not contain logits."
    case .unexpectedLogitsShape(let shape):
      "Expected rank-3 logits, got \(shape)."
    case .nonFiniteLogits:
      "The model returned non-finite logits."
    case .unsupportedLogitsType(let type):
      "Unsupported logits type: \(type)."
    case .invalidCausalMask:
      "Invalid causal-mask dimensions."
    }
  }
}

/// A fail-closed runtime for deterministic tool plans. Those plans never need
/// tokenization or generation; retaining the protocol boundary lets the same
/// audited AgentLoop execute them while guaranteeing that an accidental model
/// turn cannot silently proceed without the Core ML model.
struct ToolOnlyDolphinRuntime: DolphinModelRuntimeProtocol {
  func information() -> DolphinModelInformation {
    DolphinModelInformation(
      loadSeconds: 0,
      maxContextLength: 1,
      maxQueryLength: 1,
      computeUnits: "None",
      modelIdentifier: "dolphin-tool-only"
    )
  }

  func tokenCount(_ renderedPrompt: String) -> Int {
    renderedPrompt.isEmpty ? 0 : 1
  }

  func generate(
    renderedPrompt: String,
    maxNewTokens: Int,
    progress: @MainActor @Sendable (DolphinGenerationProgress) -> Void
  ) async throws -> DolphinGenerationResult {
    throw DolphinModelRuntimeError.modelNotLoaded
  }
}

protocol DolphinModelRuntimeProtocol: Sendable {
  func information() async -> DolphinModelInformation
  func tokenCount(_ renderedPrompt: String) async -> Int
  func generate(
    renderedPrompt: String,
    maxNewTokens: Int,
    progress: @MainActor @Sendable (DolphinGenerationProgress) -> Void
  ) async throws -> DolphinGenerationResult
}

actor DolphinModelRuntime: DolphinModelRuntimeProtocol {
  static let modelName = "Dolphin3.0-Llama3.2-3B-stateful-int4"
  static let modelIdentifier = "ales27pm/Dolphin3.0-CoreML@v2.0.0"
  static let stopTokenIDs: Set<Int> = [128256, 128001, 128008, 128009]

  private let model: MLModel
  private let tokenizer: any Tokenizer
  private let maxContextLength: Int
  private let maxQueryLength: Int
  private let loadSeconds: Double
  private var isGenerating = false

  init(bundle: Bundle = .main) async throws {
    guard
      let modelURL = bundle.url(
        forResource: Self.modelName,
        withExtension: "mlmodelc"
      )
    else {
      throw DolphinModelRuntimeError.modelResourceMissing
    }
    guard let tokenizerFolder = bundle.resourceURL else {
      throw DolphinModelRuntimeError.tokenizerResourcesMissing
    }

    let loadStart = ContinuousClock.now
    let configuration = MLModelConfiguration()
    configuration.computeUnits = .cpuAndGPU
    model = try await MLModel.load(
      contentsOf: modelURL,
      configuration: configuration
    )
    tokenizer = try await AutoTokenizer.from(modelFolder: tokenizerFolder)
    loadSeconds = Self.seconds(since: loadStart)

    let metadata =
      model.modelDescription.metadata[
        MLModelMetadataKey.creatorDefinedKey
      ] as? [String: String] ?? [:]
    maxContextLength =
      Int(metadata["com.ales27pm.dolphin.max_context_length"] ?? "2048") ?? 2048
    maxQueryLength =
      Int(metadata["com.ales27pm.dolphin.max_query_length"] ?? "512") ?? 512
  }

  func information() -> DolphinModelInformation {
    DolphinModelInformation(
      loadSeconds: loadSeconds,
      maxContextLength: maxContextLength,
      maxQueryLength: maxQueryLength,
      computeUnits: "CPU + GPU",
      modelIdentifier: Self.modelIdentifier
    )
  }

  func tokenCount(_ renderedPrompt: String) -> Int {
    tokenizer.encode(text: renderedPrompt, addSpecialTokens: false).count
  }

  func generate(
    renderedPrompt: String,
    maxNewTokens: Int,
    progress: @MainActor @Sendable (DolphinGenerationProgress) -> Void
  ) async throws -> DolphinGenerationResult {
    guard !isGenerating else {
      throw DolphinModelRuntimeError.generationAlreadyRunning
    }
    guard maxNewTokens > 0 else {
      throw DolphinModelRuntimeError.invalidOutputLimit
    }
    isGenerating = true
    defer { isGenerating = false }

    let promptTokens = tokenizer.encode(
      text: renderedPrompt,
      addSpecialTokens: false
    )
    guard !promptTokens.isEmpty else {
      throw DolphinModelRuntimeError.emptyPrompt
    }
    let requestedContext = promptTokens.count + maxNewTokens
    guard requestedContext <= maxContextLength else {
      throw DolphinModelRuntimeError.contextTooLong(
        actual: requestedContext,
        limit: maxContextLength
      )
    }

    let state = model.makeState()
    let generationStart = ContinuousClock.now
    var endStep = 0
    var finalPrefillLogits: MLMultiArray?

    for chunkStart in stride(
      from: 0,
      to: promptTokens.count,
      by: maxQueryLength
    ) {
      try Task.checkCancellation()
      let chunkEnd = min(chunkStart + maxQueryLength, promptTokens.count)
      let chunk = Array(promptTokens[chunkStart..<chunkEnd])
      endStep += chunk.count
      let output = try await prediction(
        tokens: chunk,
        endStep: endStep,
        state: state
      )
      try Task.checkCancellation()
      if chunkEnd == promptTokens.count {
        finalPrefillLogits = output
      }
    }

    guard let finalPrefillLogits else {
      throw DolphinModelRuntimeError.missingLogits
    }

    var generated: [Int] = []
    var logits = finalPrefillLogits
    var timeToFirstTokenSeconds = 0.0

    while generated.count < maxNewTokens {
      try Task.checkCancellation()
      let token = try greedyToken(from: logits)
      if generated.isEmpty {
        timeToFirstTokenSeconds = Self.seconds(since: generationStart)
      }
      if Self.stopTokenIDs.contains(token) || token == tokenizer.eosTokenId {
        break
      }

      generated.append(token)
      await progress(
        DolphinGenerationProgress(
          text: tokenizer.decode(tokens: generated),
          generatedTokens: generated.count
        )
      )

      guard generated.count < maxNewTokens else { break }
      endStep += 1
      logits = try await prediction(
        tokens: [token],
        endStep: endStep,
        state: state
      )
      try Task.checkCancellation()
    }

    return DolphinGenerationResult(
      text: tokenizer.decode(tokens: generated),
      promptTokens: promptTokens.count,
      generatedTokens: generated.count,
      timeToFirstTokenSeconds: timeToFirstTokenSeconds,
      totalSeconds: Self.seconds(since: generationStart)
    )
  }

  private func prediction(
    tokens: [Int],
    endStep: Int,
    state: MLState
  ) async throws -> MLMultiArray {
    let inputs = try MLDictionaryFeatureProvider(dictionary: [
      "inputIds": MLFeatureValue(multiArray: try inputIDs(tokens)),
      "causalMask": MLFeatureValue(
        multiArray: try causalMask(
          queryLength: tokens.count,
          endStep: endStep
        )
      ),
    ])
    let output = try await model.prediction(from: inputs, using: state)
    guard
      let logits = output.featureValue(for: "logits")?.multiArrayValue
    else {
      throw DolphinModelRuntimeError.missingLogits
    }
    return logits
  }

  private func inputIDs(_ tokens: [Int]) throws -> MLMultiArray {
    let result = try MLMultiArray(
      shape: [1, NSNumber(value: tokens.count)],
      dataType: .int32
    )
    let strides = result.strides.map(\.intValue)
    let values = result.dataPointer.bindMemory(
      to: Int32.self,
      capacity: result.count
    )
    for (index, token) in tokens.enumerated() {
      values[index * strides[1]] = Int32(token)
    }
    return result
  }

  private func causalMask(queryLength: Int, endStep: Int) throws -> MLMultiArray {
    guard queryLength > 0, endStep >= queryLength else {
      throw DolphinModelRuntimeError.invalidCausalMask
    }
    let result = try MLMultiArray(
      shape: [
        1,
        1,
        NSNumber(value: queryLength),
        NSNumber(value: endStep),
      ],
      dataType: .float16
    )
    let pastLength = endStep - queryLength
    let strides = result.strides.map(\.intValue)
    let values = result.dataPointer.bindMemory(
      to: UInt16.self,
      capacity: result.count
    )
    for row in 0..<queryLength {
      for column in 0..<endStep {
        let offset = row * strides[2] + column * strides[3]
        values[offset] = column <= pastLength + row ? 0x0000 : 0xFBFF
      }
    }
    return result
  }

  private func greedyToken(from logits: MLMultiArray) throws -> Int {
    guard logits.shape.count == 3 else {
      throw DolphinModelRuntimeError.unexpectedLogitsShape(
        logits.shape.map(\.intValue)
      )
    }
    let shape = logits.shape.map(\.intValue)
    let strides = logits.strides.map(\.intValue)
    let rowOffset = (shape[1] - 1) * strides[1]

    var bestToken = 0
    var bestScore = -Float.infinity
    switch logits.dataType {
    case .float16:
      let values = logits.dataPointer.bindMemory(
        to: UInt16.self,
        capacity: logits.count
      )
      for token in 0..<shape[2] {
        let score = float32(
          fromFloat16Bits: values[rowOffset + token * strides[2]]
        )
        guard score.isFinite else {
          throw DolphinModelRuntimeError.nonFiniteLogits
        }
        if score > bestScore {
          bestScore = score
          bestToken = token
        }
      }
    case .float32:
      let values = logits.dataPointer.bindMemory(
        to: Float.self,
        capacity: logits.count
      )
      for token in 0..<shape[2] {
        let score = values[rowOffset + token * strides[2]]
        guard score.isFinite else {
          throw DolphinModelRuntimeError.nonFiniteLogits
        }
        if score > bestScore {
          bestScore = score
          bestToken = token
        }
      }
    default:
      throw DolphinModelRuntimeError.unsupportedLogitsType(
        String(describing: logits.dataType)
      )
    }
    return bestToken
  }

  private func float32(fromFloat16Bits bits: UInt16) -> Float {
    let sign = UInt32(bits & 0x8000) << 16
    let exponent = bits & 0x7C00
    var significand = bits & 0x03FF

    let floatExponent: UInt32
    let floatSignificand: UInt32
    if exponent == 0 {
      if significand == 0 {
        return Float(bitPattern: sign)
      }
      var shift = 0
      while significand & 0x0400 == 0 {
        significand <<= 1
        shift += 1
      }
      significand &= 0x03FF
      floatExponent = UInt32(127 - 14 - shift) << 23
      floatSignificand = UInt32(significand) << 13
    } else if exponent == 0x7C00 {
      floatExponent = 0xFF << 23
      floatSignificand = UInt32(significand) << 13
    } else {
      floatExponent = UInt32((exponent >> 10) + (127 - 15)) << 23
      floatSignificand = UInt32(significand) << 13
    }
    return Float(bitPattern: sign | floatExponent | floatSignificand)
  }

  private static func seconds(
    since instant: ContinuousClock.Instant
  ) -> Double {
    let duration = instant.duration(to: .now)
    return Double(duration.components.seconds)
      + Double(duration.components.attoseconds) / 1_000_000_000_000_000_000
  }
}