// LLMHelper.swift - Interface with macOS native Apple Intelligence FoundationModels
import Foundation
import FoundationModels

struct Options {
    var prompt: String = ""
    var systemPrompt: String? = nil
    var check: Bool = false
}

func parseArguments() -> Options {
    var options = Options()
    var arguments = Array(CommandLine.arguments.dropFirst())
    while let argument = arguments.first {
        arguments.removeFirst()
        switch argument {
        case "--prompt":
            if !arguments.isEmpty { options.prompt = arguments.removeFirst() }
        case "--system":
            if !arguments.isEmpty { options.systemPrompt = arguments.removeFirst() }
        case "--check":
            options.check = true
        default:
            if options.prompt.isEmpty && !argument.hasPrefix("--") {
                options.prompt = argument
            }
        }
    }
    return options
}

func emit(_ payload: [String: Any]) {
    if let data = try? JSONSerialization.data(withJSONObject: payload, options: [.sortedKeys]),
       let str = String(data: data, encoding: .utf8) {
        print(str)
        fflush(stdout)
    }
}

func fail(_ message: String) -> Never {
    emit(["error": message])
    exit(1)
}

@main
@available(macOS 26.0, *)
struct LLMHelperApp {
    static func main() async {
        let options = parseArguments()

        let model = SystemLanguageModel.default
        let isAvailable = model.isAvailable

        if options.check {
            var attemptError: String = ""
            do {
                let session = LanguageModelSession(model: model)
                _ = try await session.respond(to: "test")
            } catch {
                attemptError = error.localizedDescription
            }

            emit([
                "available": isAvailable,
                "model": "macOS SystemLanguageModel",
                "attemptError": attemptError,
                "reason": isAvailable ? "macOS Apple Intelligence model available" : "SystemLanguageModel isAvailable returned false (\(attemptError))"
            ])
            exit(0)
        }

        guard !options.prompt.isEmpty else {
            fail("no prompt provided")
        }

        do {
            let instructions = options.systemPrompt ?? ""
            let session: LanguageModelSession
            if !instructions.isEmpty {
                session = LanguageModelSession(model: model, instructions: instructions)
            } else {
                session = LanguageModelSession(model: model)
            }

            let response = try await session.respond(to: options.prompt)
            let resultText = String(describing: response)
            emit(["done": true, "text": resultText])
        } catch {
            fail("LanguageModel error: \(error.localizedDescription)")
        }
    }
}