import Foundation

/// One TeX expression in a body: what is typeset, whether it is display
/// math, and the source with its delimiters, shown where it is refused.
nonisolated struct MathExpression: Equatable, Sendable {
    let tex: String
    let display: Bool
    let source: String
}

nonisolated enum MathPiece: Equatable, Sendable {
    case text(String)
    case math(MathExpression)
}

/// The web's per-document limits: past them, delimiters stay literal text.
nonisolated struct MathBudget: Sendable {
    var expressions: Int
    var bytes: Int

    static let document = MathBudget(expressions: 1000, bytes: 256 << 10)

    mutating func take(_ tex: String) -> Bool {
        let size = tex.utf8.count
        guard expressions > 0, size <= bytes else { return false }
        expressions -= 1
        bytes -= size
        return true
    }

    /// `text` with every expression past the budget escaped, so a later
    /// split leaves it as prose and the markdown renderer shows its source.
    mutating func limit(_ text: String) -> String {
        let pieces = splitMarkdownMath(text)
        guard pieces.contains(where: { if case .math = $0 { true } else { false } }) else { return text }
        var out = ""
        for piece in pieces {
            switch piece {
            case .text(let text):
                out += text
            case .math(let math) where take(math.tex):
                out += math.source
            case .math(let math):
                out += math.source.replacingOccurrences(of: "$", with: #"\$"#)
            }
        }
        return out
    }
}

/// Splits inline markdown into prose and math by the web's rules
/// (internal/httpd/math.go): `$$…$$` is display math; `$…$` needs a
/// non-space after the opener, and a non-space before the closer with no
/// digit after it, so "$5 and $10" is prose. A backslash escapes the next
/// character, and code spans and link destinations are never math.
nonisolated func splitMarkdownMath(_ text: String) -> [MathPiece] {
    let chars = Array(text)
    var pieces: [MathPiece] = []
    var prose = ""
    var i = 0

    func flush() {
        if !prose.isEmpty { pieces.append(.text(prose)) }
        prose = ""
    }
    func isSpace(_ c: Character) -> Bool { c == " " || c == "\t" || c == "\n" || c == "\r" }

    // Every `$` that can close inline math, found in one pass so each opener
    // looks its closer up rather than rescanning the rest of the text.
    var closers: [Int] = []
    var k = 1
    while k < chars.count {
        if chars[k] == "\\" {
            k += 2
            continue
        }
        if chars[k] == "$", !isSpace(chars[k - 1]),
           !(k + 1 < chars.count && chars[k + 1].isASCII && chars[k + 1].isNumber) {
            closers.append(k)
        }
        k += 1
    }
    func firstCloser(atOrAfter offset: Int) -> Int? {
        var low = 0, high = closers.count
        while low < high {
            let mid = (low + high) / 2
            if closers[mid] < offset { low = mid + 1 } else { high = mid }
        }
        return low < closers.count ? closers[low] : nil
    }

    while i < chars.count {
        let c = chars[i]
        switch c {
        case "\\":
            prose += String(chars[i..<min(i + 2, chars.count)])
            i += 2
        case "`":
            var run = 1
            while i + run < chars.count, chars[i + run] == "`" { run += 1 }
            var j = i + run
            var close: Int?
            while j < chars.count {
                guard chars[j] == "`" else { j += 1; continue }
                var k = 1
                while j + k < chars.count, chars[j + k] == "`" { k += 1 }
                if k == run { close = j + k; break }
                j += k
            }
            let end = close ?? i + run
            prose += String(chars[i..<end])
            i = end
        case "]" where i + 1 < chars.count && chars[i + 1] == "(":
            let end = chars[(i + 2)...].firstIndex(of: ")").map { $0 + 1 } ?? chars.count
            prose += String(chars[i..<end])
            i = end
        case "$" where i + 1 < chars.count && chars[i + 1] == "$":
            var j = i + 2
            var close: Int?
            while j + 1 < chars.count {
                if chars[j] == "\\" { j += 2; continue }
                if chars[j] == "$", chars[j + 1] == "$" { close = j; break }
                j += 1
            }
            if let close, close > i + 2 {
                flush()
                let tex = String(chars[(i + 2)..<close])
                pieces.append(.math(MathExpression(tex: tex, display: true, source: "$$" + tex + "$$")))
                i = close + 2
            } else {
                prose += "$$"
                i += 2
            }
        case "$" where i + 1 < chars.count && !isSpace(chars[i + 1]):
            if let close = firstCloser(atOrAfter: i + 2) {
                flush()
                let tex = String(chars[(i + 1)..<close])
                pieces.append(.math(MathExpression(tex: tex, display: false, source: "$" + tex + "$")))
                i = close + 1
            } else {
                prose.append(c)
                i += 1
            }
        default:
            prose.append(c)
            i += 1
        }
    }
    flush()
    return pieces
}

/// A display block at the head of `lines`, by the web's rules: a line that
/// is `$$` alone, closed by a line ending in `$$` before a blank line, or a
/// line that is `$$…$$` whole. Returns the expression and the lines it
/// spans.
nonisolated func markdownMathBlock(_ lines: ArraySlice<String>) -> (MathExpression, lines: Int)? {
    guard let first = lines.first else { return nil }
    let opener = first.trimmingCharacters(in: .whitespaces)
    guard opener.hasPrefix("$$") else { return nil }
    let rest = String(opener.dropFirst(2))

    if rest.isEmpty {
        var tex = ""
        var source = "$$\n"
        for (offset, line) in lines.dropFirst().enumerated() {
            if line.trimmingCharacters(in: .whitespaces).isEmpty { return nil }
            let trimmed = String(line.reversed().drop(while: { $0 == " " || $0 == "\t" }).reversed())
            if trimmed.hasSuffix("$$") {
                tex += trimmed.dropLast(2)
                source += trimmed
                return (MathExpression(tex: tex, display: true, source: source), offset + 2)
            }
            tex += line + "\n"
            source += line + "\n"
        }
        return nil
    }
    if rest.count > 2, rest.hasSuffix("$$") {
        return (MathExpression(tex: String(rest.dropLast(2)), display: true, source: opener), 1)
    }
    return nil
}