//
//  CLI.swift
//  Zyquo MLX
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation
import MLXLMCommon

/// Command-line proof-of-concept mode (Phase 2 gate): run inference on local
/// model directories without the UI.
///
///     ZyquoMLX --infer <model-dir> [--prompt "…"] [--max-tokens N] [--image <path>]
///     ZyquoMLX --embed <model-dir> --text "…" [--text "…"]…
enum CLI {

    static var shouldRun: Bool {
        let args = CommandLine.arguments
        return args.contains("--infer") || args.contains("--embed")
    }

    static func run() async -> Int32 {
        do {
            let args = CommandLine.arguments
            if let dir = value(after: "--infer", in: args) {
                try await infer(
                    directory: dir,
                    prompt: value(after: "--prompt", in: args)
                        ?? "Explain in one short sentence what MLX is.",
                    maxTokens: value(after: "--max-tokens", in: args).flatMap(Int.init) ?? 256,
                    imagePath: value(after: "--image", in: args)
                )
                return 0
            }
            if let dir = value(after: "--embed", in: args) {
                let texts = values(after: "--text", in: args)
                try await embed(
                    directory: dir,
                    texts: texts.isEmpty
                        ? ["The quick brown fox", "A fast auburn fox", "Quarterly revenue grew 4%"]
                        : texts)
                return 0
            }
            FileHandle.standardError.write(Data("usage: ZyquoMLX --infer <dir> | --embed <dir>\n".utf8))
            return 2
        } catch {
            FileHandle.standardError.write(Data("error: \(error.localizedDescription)\n".utf8))
            return 1
        }
    }

    // MARK: - Subcommands

    private static func infer(directory: String, prompt: String, maxTokens: Int, imagePath: String?) async throws {
        let url = URL(fileURLWithPath: (directory as NSString).expandingTildeInPath)
        let model = try await ModelStore.shared.describe(directory: url)
        print("model:   \(model.name) [\(model.type.displayName)\(model.quantization.map { ", \($0.label)" } ?? "")]")
        print("verdict: \(MemoryAdvisor.inferenceVerdict(for: model).displayName)")

        let engine = InferenceEngine.shared
        let loadStart = Date()
        try await engine.load(model: model)
        print("loaded in \(String(format: "%.2f", Date().timeIntervalSince(loadStart)))s\n")

        var message = Chat.Message.user(prompt)
        if let imagePath {
            let imageURL = URL(fileURLWithPath: (imagePath as NSString).expandingTildeInPath)
            message = Chat.Message.user(prompt, images: [.url(imageURL)])
        }

        var params = GenerationParams()
        params.maxTokens = maxTokens

        let stream = try await engine.generate(messages: [message], params: params)
        var stats: InferenceStats?
        for try await event in stream {
            switch event {
            case .chunk(let text):
                print(text, terminator: "")
                fflush(stdout)
            case .finished(let s):
                stats = s
            }
        }
        print("\n")
        if let stats {
            print("── stats ──────────────────────────────")
            print("prompt tokens:    \(stats.promptTokens)")
            print("generated tokens: \(stats.generatedTokens)")
            print(String(format: "ttft:             %.2fs", stats.ttft))
            print(String(format: "speed:            %.1f tok/s", stats.tokensPerSecond))
            print("stop reason:      \(stats.stopReason)")
        }
        let freed = try await engine.unload()
        print("unloaded (freed \(ByteCountFormatter.string(fromByteCount: freed, countStyle: .memory)))")
    }

    private static func embed(directory: String, texts: [String]) async throws {
        let url = URL(fileURLWithPath: (directory as NSString).expandingTildeInPath)
        let model = try await ModelStore.shared.describe(directory: url)
        print("model: \(model.name) [\(model.type.displayName)]")

        let engine = InferenceEngine.shared
        try await engine.load(model: model)

        let start = Date()
        let vectors = try await engine.embed(texts: texts)
        let elapsed = Date().timeIntervalSince(start)

        for (text, vector) in zip(texts, vectors) {
            let preview = vector.prefix(4).map { String(format: "%+.4f", $0) }.joined(separator: ", ")
            print("dim=\(vector.count)  [\(preview), …]  \"\(text)\"")
        }
        if vectors.count >= 2 {
            print("\n── cosine similarity ──────────────────")
            for i in 0..<vectors.count {
                for j in (i + 1)..<vectors.count {
                    let sim = InferenceEngine.cosineSimilarity(vectors[i], vectors[j])
                    print(String(format: "%.4f  \"%@\" ↔ \"%@\"", sim, texts[i], texts[j]))
                }
            }
        }
        print(String(format: "\nembedded %d texts in %.2fs", texts.count, elapsed))
        try await engine.unload()
    }

    // MARK: - Arg parsing (shared with CLIFoundry)

    static func value(after flag: String, in args: [String]) -> String? {
        guard let index = args.firstIndex(of: flag), index + 1 < args.count else { return nil }
        return args[index + 1]
    }

    private static func values(after flag: String, in args: [String]) -> [String] {
        var out: [String] = []
        var i = 0
        while i < args.count {
            if args[i] == flag, i + 1 < args.count {
                out.append(args[i + 1])
                i += 2
            } else {
                i += 1
            }
        }
        return out
    }
}