import Foundation
import CoreAI

struct Fixture: Decodable {
    let input_ids: [Int32]
    let attention_mask: [Int32]
    let position_ids: [Int32]
    let expected: [Float]
    let output: String
}

@main struct VerifyDecision {
    static func main() async throws {
        let root = URL(fileURLWithPath: CommandLine.arguments[1], isDirectory: true)
        let fixture = try JSONDecoder().decode(Fixture.self, from: Data(contentsOf: root.appendingPathComponent("swift_fixture.json")))
        let model = try await AIModel(contentsOf: root.appendingPathComponent("decision.aimodel"))
        guard let descriptor = model.functionDescriptor(for: "main"), descriptor.stateNames.isEmpty,
              let function = try model.loadFunction(named: "main") else { fatalError("Missing stateless main function") }
        var inputs: [String: NDArray] = [:]
        for (name, values) in [("input_ids", fixture.input_ids), ("attention_mask", fixture.attention_mask), ("position_ids", fixture.position_ids)] {
            guard case .ndArray(let d) = descriptor.inputDescriptor(of: name) else { fatalError("Missing tensor") }
            var array = NDArray(descriptor: d)
            var view = array.mutableView(as: Int32.self)
            view.copyElements(fromContentsOf: values)
            inputs[name] = array
        }
        var result = try await function.run(inputs: inputs)
        guard let output = result.remove(fixture.output)?.ndArray else { fatalError("Missing output") }
        let logits: [Float] = output.view(as: Float.self).withUnsafePointer { ptr, _, _ in
            Array(UnsafeBufferPointer(start: ptr, count: fixture.expected.count))
        }
        let maximum = logits.max()!
        let exps = logits.map { exp($0 - maximum) }
        let total = exps.reduce(0, +)
        let probabilities = exps.map { $0 / total }
        let error = zip(probabilities, fixture.expected).map { abs($0 - $1) }.max()!
        guard error < 0.02, probabilities.allSatisfy({ $0.isFinite }) else { fatalError("Swift/Python parity failed: \(error)") }
        let report: [String: Any] = ["swift_native_coreai": true, "probabilities": probabilities, "maximum_probability_error_vs_python": error, "stateless": true]
        print(String(data: try JSONSerialization.data(withJSONObject: report, options: [.prettyPrinted, .sortedKeys]), encoding: .utf8)!)
    }
}