import CoreGraphics
import CoreML
import Foundation
import ImageIO
import SuryaCoreMLRuntime

struct Arguments {
    let modelDirectory: URL
    let imageURL: URL?
    let pixelValuesURL: URL?
    let dumpPixelValuesURL: URL?
    let maxTokens: Int
    let useInt8Vision: Bool
}

func parseArguments() throws -> Arguments {
    var modelDirectory: URL?
    var imageURL: URL?
    var pixelValuesURL: URL?
    var dumpPixelValuesURL: URL?
    var maxTokens = 32
    var useInt8Vision = true

    var iterator = CommandLine.arguments.dropFirst().makeIterator()
    while let arg = iterator.next() {
        switch arg {
        case "--model-dir":
            guard let value = iterator.next() else { throw SuryaRuntimeError.malformedConstants("Missing --model-dir value") }
            modelDirectory = URL(fileURLWithPath: value)
        case "--image":
            guard let value = iterator.next() else { throw SuryaRuntimeError.malformedConstants("Missing --image value") }
            imageURL = URL(fileURLWithPath: value)
        case "--pixel-values-fp16":
            guard let value = iterator.next() else { throw SuryaRuntimeError.malformedConstants("Missing --pixel-values-fp16 value") }
            pixelValuesURL = URL(fileURLWithPath: value)
        case "--dump-pixel-values-fp16":
            guard let value = iterator.next() else { throw SuryaRuntimeError.malformedConstants("Missing --dump-pixel-values-fp16 value") }
            dumpPixelValuesURL = URL(fileURLWithPath: value)
        case "--max-tokens":
            guard let value = iterator.next(), let parsed = Int(value) else {
                throw SuryaRuntimeError.malformedConstants("Bad --max-tokens value")
            }
            maxTokens = parsed
        case "--vision":
            guard let value = iterator.next() else { throw SuryaRuntimeError.malformedConstants("Missing --vision value") }
            useInt8Vision = value.lowercased() != "fp16"
        default:
            throw SuryaRuntimeError.malformedConstants("Unknown argument \(arg)")
        }
    }

    guard let modelDirectory else {
        throw SuryaRuntimeError.malformedConstants("Usage: surya-coreml-smoke --model-dir /path/to/repo --image /path/to/image [--max-tokens 32] [--vision int8|fp16]")
    }
    if imageURL == nil && pixelValuesURL == nil {
        throw SuryaRuntimeError.malformedConstants("Pass either --image or --pixel-values-fp16")
    }

    return Arguments(
        modelDirectory: modelDirectory,
        imageURL: imageURL,
        pixelValuesURL: pixelValuesURL,
        dumpPixelValuesURL: dumpPixelValuesURL,
        maxTokens: maxTokens,
        useInt8Vision: useInt8Vision
    )
}

func loadImage(_ url: URL) throws -> CGImage {
    guard let source = CGImageSourceCreateWithURL(url as CFURL, nil),
          let image = CGImageSourceCreateImageAtIndex(source, 0, nil) else {
        throw SuryaRuntimeError.missingFile(url)
    }
    return image
}

func loadPixelValues(_ url: URL) throws -> MLMultiArray {
    let data = try Data(contentsOf: url, options: [.mappedIfSafe])
    let shape = [1024, 1536]
    let expected = shape.reduce(1, *) * MemoryLayout<UInt16>.size
    guard data.count == expected else {
        throw SuryaRuntimeError.invalidShape("Expected \(expected) bytes, got \(data.count)")
    }
    let array = try MLMultiArray(shape: shape.map(NSNumber.init(value:)), dataType: .float32)
    let dst = array.dataPointer.bindMemory(to: Float.self, capacity: shape.reduce(1, *))
    data.withUnsafeBytes { raw in
        let src = raw.bindMemory(to: UInt16.self)
        for i in 0..<src.count {
            dst[i] = Float(Float16(bitPattern: src[i]))
        }
    }
    return array
}

func dumpPixelValues(_ array: MLMultiArray, to url: URL) throws {
    var data = Data(capacity: array.count * MemoryLayout<UInt16>.size)
    for index in 0..<array.count {
        let value = array.dataPointer.bindMemory(to: Float.self, capacity: array.count)[index]
        var bits = Float16(value).bitPattern
        withUnsafeBytes(of: &bits) { data.append(contentsOf: $0) }
    }
    try data.write(to: url)
}

do {
    let args = try parseArguments()
    let runtime = try SuryaCoreMLRuntime(modelDirectory: args.modelDirectory)
    let result: SuryaGenerationResult
    if let pixelValuesURL = args.pixelValuesURL {
        let pixelValues = try loadPixelValues(pixelValuesURL)
        result = try runtime.generate(pixelValues: pixelValues, maxNewTokens: args.maxTokens, useInt8Vision: args.useInt8Vision)
    } else if let imageURL = args.imageURL {
        let image = try loadImage(imageURL)
        let preprocessor = QwenFixedImagePreprocessor()
        let pixelValues = try preprocessor.pixelValues(from: image)
        if let dumpURL = args.dumpPixelValuesURL {
            try dumpPixelValues(pixelValues, to: dumpURL)
        }
        result = try runtime.generate(pixelValues: pixelValues, maxNewTokens: args.maxTokens, useInt8Vision: args.useInt8Vision)
    } else {
        throw SuryaRuntimeError.malformedConstants("Pass either --image or --pixel-values-fp16")
    }
    print(result.tokenIDs.map(String.init).joined(separator: " "))
    if let text = result.text {
        print(text)
    }
} catch {
    fputs("surya-coreml-smoke failed: \(error)\n", stderr)
    exit(1)
}