import OrgSwift
import OrgSwiftUI
import SwiftMath
import SwiftUI
import UIKit

/// One typeset expression: an alpha mask, tinted by the text around it, and
/// how far its baseline sits above the bottom edge.
struct RenderedMath {
    let image: CGImage
    let scale: CGFloat
    let descent: CGFloat
}

/// Typesets TeX with SwiftMath. Input is bounded as the server's converter
/// bounds it (internal/texmath), so a pathological body cannot run the
/// parser deep or long; anything refused shows as its source.
enum MathRenderer {
    static let maxInput = 8 << 10
    static let maxDepth = 64
    /// SwiftMath recurses once per command that takes an argument, so the
    /// count bounds nesting that braces do not show.
    static let maxCommands = 200
    /// The longest side of an image, in points.
    static let maxSide: CGFloat = 4096

    private final class Entry {
        let value: RenderedMath?
        init(_ value: RenderedMath?) { self.value = value }
    }

    private static let cache: NSCache<NSString, Entry> = {
        let cache = NSCache<NSString, Entry>()
        cache.countLimit = 500
        cache.totalCostLimit = 64 << 20
        return cache
    }()

    static func render(_ tex: String, display: Bool, pointSize: CGFloat) -> RenderedMath? {
        let key = "\(display) \(pointSize) \(tex)" as NSString
        if let entry = cache.object(forKey: key) { return entry.value }
        let value = typeset(tex, display: display, pointSize: pointSize)
        cache.setObject(Entry(value), forKey: key, cost: value.map { $0.image.bytesPerRow * $0.image.height } ?? 0)
        return value
    }

    private static func typeset(_ tex: String, display: Bool, pointSize: CGFloat) -> RenderedMath? {
        guard !tex.trimmingCharacters(in: .whitespacesAndNewlines).isEmpty,
              tex.utf8.count <= maxInput, admits(tex) else { return nil }
        let latex = swiftMath(tex)
        let mode: MTMathUILabelMode = display ? .display : .text
        // Measured before it is drawn, so an expression too large to show never
        // allocates its bitmap.
        let label = MTMathUILabel()
        label.labelMode = mode
        label.fontSize = pointSize
        label.latex = latex
        let size = label.intrinsicContentSize
        guard label.error == nil, size.width <= maxSide, size.height <= maxSide else { return nil }
        var math = MathImage(latex: latex, fontSize: pointSize, textColor: .black,
                             labelMode: mode, textAlignment: .left)
        let (error, image, layout) = math.asImage()
        guard error == nil, let image, let cgImage = image.cgImage, let layout,
              image.size.width > 0, image.size.height > 0 else { return nil }
        return RenderedMath(image: cgImage, scale: image.scale, descent: layout.descent)
    }

    /// Whether `tex` nests groups, `\left…\right` and environments no
    /// deeper than `maxDepth`, has at most `maxCommands` commands, and has no
    /// unescaped `$`, `#` or `%`, which the server's converter refuses, and no
    /// character outside printable ASCII, which SwiftMath leaves out without
    /// an error.
    static func admits(_ tex: String) -> Bool {
        guard tex.unicodeScalars.allSatisfy({ (0x20...0x7E).contains($0.value) || $0 == "\n" || $0 == "\t" })
        else { return false }
        var depth = 0
        var commands = 0
        var chars = tex[...]
        while let c = chars.popFirst() {
            switch c {
            case "\\":
                let name = chars.prefix(while: \.isLetter)
                if name.isEmpty {
                    chars = chars.dropFirst()
                    continue
                }
                chars = chars.dropFirst(name.count)
                commands += 1
                if commands > maxCommands { return false }
                if name == "left" || name == "begin" { depth += 1 }
                if name == "right" || name == "end" { depth -= 1 }
            case "{": depth += 1
            case "}": depth -= 1
            case "$", "#", "%": return false
            default: continue
            }
            if depth > maxDepth { return false }
        }
        return true
    }

    private static let commands: [Substring: String] = [
        "dfrac": "frac", "tfrac": "frac", "cfrac": "frac",
        "dbinom": "binom", "tbinom": "binom",
        "operatorname": "mathrm", "boldsymbol": "mathbf", "bm": "mathbf", "mathscr": "mathcal",
        "textrm": "text", "textnormal": "text", "mbox": "text", "middle": "",
    ]

    private static let environments: [Substring: String] = [
        "align": "aligned", "align*": "aligned", "gathered": "gather", "gather*": "gather",
        "smallmatrix": "matrix", "equation": "", "equation*": "", "displaymath": "",
    ]

    /// The web's spellings in SwiftMath's: the `\frac` and `\binom` variants,
    /// font and operator names, primes, `~`, and the environments it names
    /// differently. An environment that only displays its body is dropped.
    static func swiftMath(_ tex: String) -> String {
        var out = ""
        var chars = tex[...]
        while let c = chars.popFirst() {
            switch c {
            case "\\":
                let name = chars.prefix(while: \.isLetter)
                guard !name.isEmpty else {
                    out.append(c)
                    if let next = chars.popFirst() { out.append(next) }
                    continue
                }
                chars = chars.dropFirst(name.count)
                if name == "operatorname", chars.first == "*" { chars = chars.dropFirst() }
                // An operator name is upright with a thin space after it.
                if name == "operatorname", chars.first == "{", let close = chars.firstIndex(of: "}") {
                    out += "\\mathrm" + chars[...close] + "\\,"
                    chars = chars[chars.index(after: close)...]
                    continue
                }
                if name == "begin" || name == "end", chars.first == "{",
                   let close = chars.firstIndex(of: "}") {
                    let env = chars[chars.index(after: chars.startIndex)..<close]
                    if let renamed = environments[env] {
                        chars = chars[chars.index(after: close)...]
                        if !renamed.isEmpty { out += "\\\(name){\(renamed)}" }
                        continue
                    }
                }
                let renamed = commands[name] ?? String(name)
                if !renamed.isEmpty { out += "\\" + renamed }
            case "'":
                var primes = 1
                while chars.first == "'" { primes += 1; chars = chars.dropFirst() }
                out += "^{" + String(repeating: "\\prime", count: primes) + "}"
            case "~":
                out += "\\;"
            default:
                out.append(c)
            }
        }
        return out
    }

    /// The expression as text: the image, set on the baseline and tinted
    /// like the text around it, or the source where it is refused.
    static func text(_ math: MathExpression, pointSize: CGFloat) -> Text {
        drawn(math.tex, display: math.display, pointSize: pointSize) ?? Text(verbatim: math.source)
    }

    static func drawn(_ tex: String, display: Bool, pointSize: CGFloat) -> Text? {
        guard let rendered = render(tex, display: display, pointSize: pointSize) else { return nil }
        return Text(Image(rendered.image, scale: rendered.scale, label: Text(verbatim: tex))
            .renderingMode(.template))
            .baselineOffset(-rendered.descent)
    }

    /// `OrgView`'s math hook: each expression the allowance admits, drawn at `pointSize`.
    static func org(pointSize: CGFloat, allowance: MathAllowance) -> OrgMathRenderer {
        { math in
            allowance.admits(math) ? drawn(math.tex, display: math.display, pointSize: pointSize) : nil
        }
    }
}

/// The web's per-document budget for org, which is drawn view by view: the first
/// expressions asked for are admitted, and keep their place on later passes.
final class MathAllowance {
    private let start: MathBudget
    private var budget: MathBudget
    private var admitted: Set<OrgMath> = []
    private var text: String?

    init(budget: MathBudget = .document) {
        start = budget
        self.budget = budget
    }

    /// Starts over when the document changes.
    func reset(for text: String) {
        guard text != self.text else { return }
        self.text = text
        budget = start
        admitted = []
    }

    func admits(_ math: OrgMath) -> Bool {
        if admitted.contains(math) { return true }
        guard budget.take(math.tex) else { return false }
        admitted.insert(math)
        return true
    }
}

/// A display block: centred, scrolling sideways when wider than the screen,
/// or the source as code where it is refused.
struct MathBlockView: View {
    let math: MathExpression
    let pointSize: CGFloat

    var body: some View {
        if let rendered = MathRenderer.render(math.tex, display: true, pointSize: pointSize) {
            let image = Image(rendered.image, scale: rendered.scale, label: Text(verbatim: math.tex))
                .renderingMode(.template)
                .padding(.vertical, 4)
            ViewThatFits(in: .horizontal) {
                image.frame(maxWidth: .infinity)
                ScrollView(.horizontal, showsIndicators: false) { image }
            }
        } else {
            ScrollView(.horizontal) {
                Text(math.source)
                    .font(.gbMono(.caption))
                    .padding(10)
            }
            .background(Color.gbCodeBackground, in: RoundedRectangle(cornerRadius: 2))
        }
    }
}

extension UIFont {
    /// The point size of `style` at a Dynamic Type size, for typesetting
    /// math beside text set in that style.
    static func pointSize(_ style: UIFont.TextStyle, _ size: DynamicTypeSize) -> CGFloat {
        preferredFont(forTextStyle: style,
                      compatibleWith: UITraitCollection(preferredContentSizeCategory: UIContentSizeCategory(size)))
            .pointSize
    }
}