import Foundation
import FoundationModels

// apple-llm-polish: pipe JSON {id?, mode: "line"|"paragraph", text, prev1?, prev2?, context?, language?} -> {id, text, ok, error?}
// Uses Apple on-device 3B Foundation Model (SystemLanguageModel)
// Much faster than Ollama, ANE-accelerated, offline

struct InMsg: Decodable {
    var id: String?
    var mode: String // "line" | "paragraph" | "check" | "bench"
    var text: String?
    var prev1: String?
    var prev2: String?
    var context: String?
    var prevSource: String?
    var language: String?
    var bench: String?
}

struct OutMsg: Encodable {
    var id: String?
    var ok: Bool
    var text: String
    var error: String?
    var ms: Int?
}

func log(_ s: String) { fputs(s+"\n", stderr) }

@main
struct AppleLLMPolish {
    static func main() async {
        let args = CommandLine.arguments
        if args.contains("--help") || args.contains("-h") {
            fputs("Usage: apple-llm-polish [--check] [--bench \"text\"] [--pipe]\nPipe JSONL in stdin, JSONL out\n", stderr); exit(0)
        }
        if args.contains("--check") {
            await runCheck(); return
        }
        if let idx = args.firstIndex(of: "--bench"), idx+1 < args.count {
            await runBench(args[idx+1]); return
        }
        if args.contains("--bench-text") {
            // bench via stdin text file? not needed
        }
        await runPipe()
    }

    static func runCheck() async {
        let m = SystemLanguageModel.default
        print("{\"available\": \(m.isAvailable), \"availability\": \"\(String(describing: m.availability))\"}")
        if case .unavailable(let reason) = m.availability {
            print("{\"reason\":\"\(String(describing: reason))\"}")
        }
        if m.isAvailable {
            // quick ping
            do {
                let session = LanguageModelSession(model: m, instructions: "You are a concise assistant.")
                let r = try await session.respond(to: "Say ok")
                print("{\"ping\":\"\(r.content)\",\"ok\":true}")
            } catch {
                print("{\"ping_error\":\"\(error.localizedDescription)\",\"ok\":false}")
            }
        }
    }

    static func runBench(_ text: String) async {
        let m = SystemLanguageModel.default
        guard m.isAvailable else { print("{\"ok\":false,\"error\":\"model unavailable \(m.availability)\"}"); return }
        let session = LanguageModelSession(model: m, instructions: instructionsFor(mode: "line", language: nil))
        let t0 = Date()
        do {
            let r = try await session.respond(to: promptFor(mode: "line", text: text, prev1: "", prev2: "", context: nil, prevSource: nil, language: nil))
            let ms = Int(Date().timeIntervalSince(t0)*1000)
            print("{\"ok\":true,\"text\":\"\(r.content.replacingOccurrences(of: "\"", with: "\\\""))\",\"ms\":\(ms)}")
        } catch {
            print("{\"ok\":false,\"error\":\"\(error.localizedDescription)\"}")
        }
    }

    static func runPipe() async {
        let m = SystemLanguageModel.default
        guard m.isAvailable else {
            // emit errors for all input lines so caller doesn't hang
            let reason = "\(m.availability)"
            while let line = readLine() {
                if line.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty { continue }
                var idv: String? = nil
                if let data = line.data(using: .utf8), let dict = try? JSONSerialization.jsonObject(with: data) as? [String:Any] {
                    idv = dict["id"] as? String
                }
                let out = OutMsg(id: idv, ok: false, text: "", error: "model unavailable: \(reason)", ms: nil)
                if let d = try? JSONEncoder().encode(out), let s = String(data: d, encoding: .utf8) { print(s); fflush(stdout) }
            }
            return
        }

        // Create single session reused across requests (keeps KV cache warm)
        let lineSession = LanguageModelSession(model: SystemLanguageModel(useCase: .general, guardrails: .permissiveContentTransformations), instructions: freeflowSystemPrompt(language: nil, mode: "line"))
        let paraSession = LanguageModelSession(model: SystemLanguageModel(useCase: .general, guardrails: .permissiveContentTransformations), instructions: freeflowSystemPrompt(language: nil, mode: "paragraph"))
        // Prewarm
        lineSession.prewarm()
        paraSession.prewarm()

        log("[apple-llm-polish] ready, ANE-backed, model available")

        while let line = readLine() {
            let trimmed = line.trimmingCharacters(in: .whitespacesAndNewlines)
            if trimmed.isEmpty { continue }
            guard let data = line.data(using: .utf8) else { continue }
            guard let req = try? JSONDecoder().decode(InMsg.self, from: data) else {
                let out = OutMsg(id: nil, ok: false, text: "", error: "bad json", ms: nil)
                if let d = try? JSONEncoder().encode(out), let s = String(data: d, encoding: .utf8) { print(s); fflush(stdout) }
                continue
            }
            if req.mode == "check" {
                let out = OutMsg(id: req.id, ok: m.isAvailable, text: "\(m.availability)", error: nil, ms: 0)
                if let d = try? JSONEncoder().encode(out), let s = String(data: d, encoding: .utf8) { print(s); fflush(stdout) }
                continue
            }
            let t0 = Date()
            do {
                let (session, prompt): (LanguageModelSession, String)
                let lang = req.language
                switch req.mode {
                case "paragraph":
                    // For paragraph we want deterministic polish first? We do full LLM here, but freeflow deterministic already did pre-clean.
                    // Use freeflow system prompt
                    session = paraSession
                    prompt = buildParagraphPrompt(context: req.context ?? "", prevSource: req.prevSource ?? "", newText: req.text ?? "")
                default: // line
                    session = lineSession
                    prompt = buildLinePrompt(text: req.text ?? "", prev1: req.prev1 ?? "", prev2: req.prev2 ?? "", language: lang)
                }
                // Generation options - low temp for deterministic polish
                var opts = GenerationOptions()
                opts.temperature = req.mode == "paragraph" ? 0.2 : 0.1
                let resp = try await session.respond(to: prompt, options: opts)
                let ms = Int(Date().timeIntervalSince(t0)*1000)
                let cleaned = resp.content.trimmingCharacters(in: .whitespacesAndNewlines)
                let out = OutMsg(id: req.id, ok: true, text: cleaned.isEmpty ? (req.text ?? "") : cleaned, error: nil, ms: ms)
                if let d = try? JSONEncoder().encode(out), let s = String(data: d, encoding: .utf8) { print(s); fflush(stdout) }
            } catch {
                let ms = Int(Date().timeIntervalSince(t0)*1000)
                let err = error.localizedDescription
                // guardrail or unavailable -> fallback to original text so pipeline doesn't break
                let fallback = req.text ?? ""
                let out = OutMsg(id: req.id, ok: false, text: fallback, error: err, ms: ms)
                if let d = try? JSONEncoder().encode(out), let s = String(data: d, encoding: .utf8) { print(s); fflush(stdout) }
            }
        }
    }

    // MARK: - Prompts (ported from freeflow_polish.py + engine_llm.py freeflow style)

    static func freeflowSystemPrompt(language: String?, mode: String) -> String {
        if mode == "paragraph" {
            return """
            You are a careful live transcript editor working from noisy translated source text.

            Goal:
            Produce the most faithful readable English update.

            Decision rules:
            1. NEW SOURCE TEXT is the primary evidence and should dominate the output.
            2. Use PREVIOUS SOURCE TEXT and PREVIOUS REFINED CONTEXT only when they clearly help resolve or continue the new material.
            3. The final answer must be English only.
            4. Never copy non-English words into the output unless they are proper names.
            5. If a non-English fragment appears as an isolated clause, trailing fragment, or side comment without a clear English anchor, drop it.
            6. Only translate a non-English fragment when it is clearly central to the same idea and its meaning is reasonably obvious from context.
            7. If the new material starts a fresh thought, output only the new material.
            8. Remove exact or near-exact repetition unless the repetition is clearly intentional rhetoric.
            9. Do not add explanations, disclaimers, or meta commentary.
            10. Prefer conservative wording over guessed meaning.
            11. Return only the revised transcript text.
            """
        } else {
            // line
            return """
            You are a real-time caption polisher.

            Task: Polish the current caption line to be readable, preserving meaning exactly.

            Rules:
            - Fix punctuation, casing, obvious STT typos.
            - Expand spoken punctuation like <keep>comma</keep> already expanded upstream.
            - Remove filler words (uh, um, erm) if clearly disfluency.
            - Keep all semantic content; do not rewrite with prior context.
            - If uncertain, return original.
            - Output only one polished line, no explanation.
            """
        }
    }

    static func buildLinePrompt(text: String, prev1: String, prev2: String, language: String?) -> String {
        var p = ""
        if !prev2.isEmpty { p += "Previous line 2: \(prev2)\n" }
        if !prev1.isEmpty { p += "Previous line 1: \(prev1)\n" }
        p += "Current line to polish:\n\(text)\n\nPolished:"
        return p
    }

    static func buildParagraphPrompt(context: String, prevSource: String, newText: String) -> String {
        return """
        Edit this transcript update.
        [PREVIOUS REFINED CONTEXT]
        \(context)

        [PREVIOUS SOURCE TEXT]
        \(prevSource)

        [NEW SOURCE TEXT]
        \(newText)
        """
    }

    static func instructionsFor(mode: String, language: String?) -> String {
        return freeflowSystemPrompt(language: language, mode: mode)
    }

    static func promptFor(mode: String, text: String, prev1: String, prev2: String, context: String?, prevSource: String?, language: String?) -> String {
        if mode == "paragraph" {
            return buildParagraphPrompt(context: context ?? "", prevSource: prevSource ?? "", newText: text)
        } else {
            return buildLinePrompt(text: text, prev1: prev1, prev2: prev2, language: language)
        }
    }
}