import OrgSwift
import SwiftUI

/// Draws one TeX expression for ``OrgView``: a `Text` holding it, typically an image set on
/// the baseline, or nil to show the source as written. The style is that of the text the
/// expression sits in, so math in a heading can be set at the heading's size.
public typealias OrgMathRenderer = @MainActor (OrgMath, Font.TextStyle) -> Text?

/// The text style a math run is set in, recorded by ``OrgInlineStyler`` for the renderer.
enum OrgMathStyleAttribute: AttributedStringKey {
    typealias Value = Font.TextStyle
    static let name = "OrgSwiftUI.mathStyle"
}

private struct OrgMathKey: EnvironmentKey {
    static let defaultValue: OrgMathRenderer? = nil
}

extension EnvironmentValues {
    var orgMath: OrgMathRenderer? {
        get { self[OrgMathKey.self] }
        set { self[OrgMathKey.self] = newValue }
    }
}

/// A styled run as `Text`, each expression drawn by `math` where it can be. Runs are taken
/// by their math attribute alone, so an expression a `textTransform` split is still one;
/// two identical expressions side by side share a range, so one holding its source k times
/// draws k.
@MainActor
func orgText(_ string: AttributedString, math: OrgMathRenderer?) -> Text {
    guard let math, string.runs.contains(where: { $0[OrgMathAttribute.self] != nil }) else {
        return Text(string)
    }
    var result = Text(verbatim: "")
    for (expression, range) in string.runs[OrgMathAttribute.self] {
        let piece = AttributedString(string[range])
        let style = piece.runs.first?[OrgMathStyleAttribute.self] ?? .body
        guard let expression, let drawn = math(expression, style) else {
            result = Text("\(result)\(Text(piece))")
            continue
        }
        let count = max(1, piece.characters.count / max(1, expression.source.count))
        for _ in 0..<count { result = Text("\(result)\(drawn)") }
    }
    return result
}

/// Display math on its own: centred, scrolling sideways when wider than the column, or the
/// source as a code block where nothing draws it.
struct OrgMathBlockView: View {
    let math: OrgMath
    @Environment(\.orgMath) private var renderer

    var body: some View {
        if let drawn = renderer?(math, .body) {
            ViewThatFits(in: .horizontal) {
                drawn.frame(maxWidth: .infinity)
                ScrollView(.horizontal, showsIndicators: false) { drawn }
            }
            .padding(.vertical, 4)
        } else {
            ScrollView(.horizontal) {
                Text(math.source)
                    .font(.system(.footnote, design: .monospaced))
                    .textSelection(.enabled)
                    .padding(10)
            }
            .background(Color.secondary.opacity(0.12), in: RoundedRectangle(cornerRadius: 4))
        }
    }
}