//
//  PoCRunner.swift
//  Zyquo Local
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation

/// CLI proof-of-concept mode — the Phase 2 gate:
/// `ZyquoLocal --poc <model-dir> "<prompt>"` loads an MLX model from a local
/// directory and streams generated tokens to stdout, ending with stats.
enum PoCRunner {
    static func run(arguments: [String]) async {
        guard let flagIndex = arguments.firstIndex(of: "--poc"),
            arguments.count >= flagIndex + 3
        else {
            err("usage: ZyquoLocal --poc <model-dir> \"<prompt>\"\n")
            exit(64)
        }
        let directory = URL(fileURLWithPath: arguments[flagIndex + 1])
        let prompt = arguments[flagIndex + 2]

        let size = (try? FileManager.default.allocatedSizeOfDirectory(at: directory)) ?? 0
        let model = LocalModel(
            repoID: directory.pathComponents.suffix(2).joined(separator: "/"),
            directory: directory,
            sizeBytes: size
        )

        let engine = InferenceEngine()
        do {
            err("Loading \(model.repoID) (\(ByteCountFormatter.string(fromByteCount: size, countStyle: .file)))…\n")
            let loadStart = Date()
            try await engine.load(model: model)
            err(String(format: "Loaded in %.2fs. Generating…\n\n", Date().timeIntervalSince(loadStart)))

            let conversation = Conversation(modelID: model.repoID)
            try await engine.startSession(conversation: conversation)

            var stats: GenerationStats?
            let stream = try await engine.generate(prompt: prompt, params: GenerationParams())
            for try await event in stream {
                switch event {
                case .token(let text):
                    out(text)
                case .stats(let s):
                    stats = s
                case .finished(let reason):
                    out("\n")
                    if let s = stats {
                        err(String(
                            format: "\n⚡ %.1f tok/s · %d tokens · %.2fs to first token · prompt %d tokens · peak %@ · stop: %@\n",
                            s.tokensPerSecond,
                            s.generationTokenCount,
                            s.timeToFirstToken,
                            s.promptTokenCount,
                            ByteCountFormatter.string(fromByteCount: Int64(s.peakMemoryBytes), countStyle: .memory),
                            reason.rawValue
                        ))
                    }
                }
            }
            await engine.unload()
            exit(0)
        } catch {
            err("error: \(error.localizedDescription)\n")
            exit(1)
        }
    }

    private static func out(_ text: String) {
        FileHandle.standardOutput.write(Data(text.utf8))
    }

    private static func err(_ text: String) {
        FileHandle.standardError.write(Data(text.utf8))
    }
}

extension FileManager {
    /// Total allocated size of a directory tree in bytes.
    func allocatedSizeOfDirectory(at url: URL) throws -> Int64 {
        var total: Int64 = 0
        let keys: Set<URLResourceKey> = [.totalFileAllocatedSizeKey, .isRegularFileKey]
        guard let enumerator = enumerator(at: url, includingPropertiesForKeys: Array(keys)) else {
            return 0
        }
        for case let file as URL in enumerator {
            let values = try file.resourceValues(forKeys: keys)
            if values.isRegularFile == true {
                total += Int64(values.totalFileAllocatedSize ?? 0)
            }
        }
        return total
    }
}