// LLMHelper.swift - Interface with macOS native on-device LLM (FoundationModels / Apple Intelligence)
//
// Usage:
//   llm-helper --check
//   llm-helper --prompt "Hello" [--system "System prompt"]

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
struct LLMHelperApp {
    static func main() async {
        let options = parseArguments()

        if options.check {
            let model = SystemLanguageModel.default
            let isAvailable = model.isAvailable
            emit([
                "available": isAvailable,
                "model": "macOS SystemLanguageModel",
                "reason": isAvailable ? "macOS Apple Intelligence model available" : "SystemLanguageModel unavailable on this hardware/OS configuration"
            ])
            exit(0)
        }

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

        do {
            let model = SystemLanguageModel.default
            guard model.isAvailable else {
                fail("macOS SystemLanguageModel is not available on this machine")
            }

            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)")
        }
    }
}