// Native citation scanner + encoder: a byte-level port of citenet/text.py
// (CITE_RE, normalize, encode). scanparity.swift checks it against Python.
import Foundation

let SEQ = 64
let MAXCHARS = 46

@inline(__always) func isDigit(_ b: UInt8) -> Bool { b >= 0x30 && b <= 0x39 }
@inline(__always) func isUpper(_ b: UInt8) -> Bool { b >= 0x41 && b <= 0x5A }
@inline(__always) func isLower(_ b: UInt8) -> Bool { b >= 0x61 && b <= 0x7A }
@inline(__always) func isAlpha(_ b: UInt8) -> Bool { isUpper(b) || isLower(b) }

struct Text {
    let p: UnsafePointer<UInt8>
    let n: Int
    @inline(__always) subscript(_ i: Int) -> UInt8 { i < n ? p[i] : 0 }

    /// Length in bytes of the whitespace character at i (Python str `\s`), or 0.
    @inline(__always) func ws(_ i: Int) -> Int {
        let b = self[i]
        if b == 0x20 || (b >= 0x09 && b <= 0x0D) || (b >= 0x1C && b <= 0x1F) { return i < n ? 1 : 0 }
        if b < 0x80 { return 0 }
        let b1 = self[i + 1], b2 = self[i + 2]
        if b == 0xC2 && (b1 == 0x85 || b1 == 0xA0) { return 2 }
        if b == 0xE1 && b1 == 0x9A && b2 == 0x80 { return 3 }
        if b == 0xE2 && b1 == 0x80 && ((b2 >= 0x80 && b2 <= 0x8A) || b2 == 0xA8 || b2 == 0xA9 || b2 == 0xAF) { return 3 }
        if b == 0xE2 && b1 == 0x81 && b2 == 0x9F { return 3 }
        if b == 0xE3 && b1 == 0x80 && b2 == 0x80 { return 3 }
        return 0
    }
    @inline(__always) func skipWS(_ i: Int) -> Int { var j = i; while true { let w = ws(j); if w == 0 { return j }; j += w } }
    @inline(__always) func isRSQ(_ i: Int) -> Bool { self[i] == 0xE2 && self[i + 1] == 0x80 && self[i + 2] == 0x99 }
    @inline(__always) func isEnDash(_ i: Int) -> Bool { self[i] == 0xE2 && self[i + 1] == 0x80 && self[i + 2] == 0x93 }

    /// Python `\w` for the character ending just before i (ASCII exact; non-ASCII letters count as word chars).
    func wordBefore(_ i: Int) -> Bool {
        if i <= 0 { return false }
        let b = p[i - 1]
        if b < 0x80 { return isAlpha(b) || isDigit(b) || b == 0x5F }
        var s = i - 1
        while s > 0 && (p[s] & 0xC0) == 0x80 { s -= 1 }
        guard let sc = String(bytes: UnsafeBufferPointer(start: p + s, count: i - s), encoding: .utf8)?.unicodeScalars.first else { return false }
        return sc.properties.isAlphabetic || sc.properties.numericType != nil
    }
    /// Python `\w` for the character starting at i.
    func wordAt(_ i: Int) -> Bool {
        if i >= n { return false }
        let b = p[i]
        if b < 0x80 { return isAlpha(b) || isDigit(b) || b == 0x5F }
        var e = i + 1
        while e < n && (p[e] & 0xC0) == 0x80 { e += 1 }
        return wordBefore(e)
    }

    /// TOKEN = [A-Z][A-Za-z'’&]*\.?  |  \d+(?:d|th|st|nd|rd)\b   -> end index or -1
    @inline(__always) func token(_ i: Int) -> Int {
        let b = self[i]
        if isUpper(b) {
            var q = i + 1
            while q < n {
                let c = p[q]
                if isAlpha(c) || c == 0x27 || c == 0x26 { q += 1 } else if isRSQ(q) { q += 3 } else { break }
            }
            if self[q] == 0x2E { q += 1 }
            return q
        }
        if isDigit(b) {
            var q = i
            while isDigit(self[q]) { q += 1 }
            let c0 = self[q], c1 = self[q + 1]
            var e = -1
            if c0 == 0x64 { e = q + 1 }                                   // d
            else if (c0 == 0x74 && c1 == 0x68) || (c0 == 0x73 && c1 == 0x74) ||
                    (c0 == 0x6E && c1 == 0x64) || (c0 == 0x72 && c1 == 0x64) { e = q + 2 }   // th st nd rd
            return (e >= 0 && !wordAt(e)) ? e : -1
        }
        return -1
    }

    /// (?:at\s+)?\*?\d{1,8}  -> end or -1
    @inline(__always) func starDigits(_ s: Int, _ m: Int) -> Int {
        var t = s
        if self[t] == 0x2A { t += 1 }
        let d0 = t
        while isDigit(self[t]) && t - d0 < m { t += 1 }
        return t > d0 ? t : -1
    }
    @inline(__always) func pageAt(_ i: Int, maxDigits: Int) -> Int {
        if self[i] == 0x61 && self[i + 1] == 0x74 {
            let k = skipWS(i + 2)
            if k > i + 2 { let r = starDigits(k, maxDigits); if r >= 0 { return r } }
        }
        return starDigits(i, maxDigits)
    }

    /// One CITE_RE match starting at i (i must be at a digit). Returns (end, accept) or nil;
    /// accept = the reporter has an uppercase letter and does not open with a STOP word.
    func match(_ i: Int) -> (Int, Bool)? {
        var j = i
        while isDigit(self[j]) { j += 1 }
        if j - i > 4 { return nil }
        let k = skipWS(j)
        if k == j { return nil }
        // greedy maximal tokens, at most 5, separated by at most one whitespace char
        var ends = [Int](); ends.reserveCapacity(5)
        var upper = [Bool](); upper.reserveCapacity(5)
        var t = k
        while ends.count < 5 {
            let e = token(t)
            if e < 0 { break }
            ends.append(e); upper.append(isUpper(self[t]) || (upper.last ?? false))
            t = e + ws(e)                       // \s? : at most one whitespace char
        }
        if ends.isEmpty { return nil }
        for c in stride(from: ends.count - 1, through: 0, by: -1) {
            let e = ends[c]
            let s = skipWS(e)
            var pg = -1
            if self[s] == 0x2C {                                 // \s*,\s*at\s+ page  ("976 F. 2d, at 540")
                let q = skipWS(s + 1)
                if self[q] == 0x61 && self[q + 1] == 0x74 {
                    let r = skipWS(q + 2)
                    if r > q + 2 { pg = starDigits(r, 8) }
                }
            } else if s > e {                                    // \s+(?:at\s+)? page
                pg = pageAt(s, maxDigits: 8)
            }
            if pg < 0 { continue }
            var end = pg
            // pincites: (?:\s*,\s*(?:at\s+)?\*?\d+(?:\s*[-–]\s*\*?\d+)?)*
            while true {
                var q = skipWS(end)
                if self[q] != 0x2C { break }
                q = skipWS(q + 1)
                let d = pageAt(q, maxDigits: Int.max)
                if d < 0 { break }
                var e2 = d
                var r = skipWS(d)
                if self[r] == 0x2D || isEnDash(r) {
                    r = skipWS(r + (self[r] == 0x2D ? 1 : 3))
                    if self[r] == 0x2A { r += 1 }
                    let d0 = r
                    while isDigit(self[r]) { r += 1 }
                    if r > d0 { e2 = r }
                }
                end = e2
            }
            // paren: \s*\([^()]{0,60}?\d{4}\)
            let po = skipWS(end)
            if self[po] == 0x28 {
                var q = po + 1, chars = 0
                while q < n && p[q] != 0x28 && p[q] != 0x29 && chars <= 64 {
                    if (p[q] & 0xC0) != 0x80 { chars += 1 }
                    q += 1
                }
                if q < n && p[q] == 0x29 && chars >= 4 && chars - 4 <= 60 &&
                    isDigit(p[q - 1]) && isDigit(p[q - 2]) && isDigit(p[q - 3]) && isDigit(p[q - 4]) {
                    end = q + 1
                }
            }
            let first = Array(UnsafeBufferPointer(start: p + k, count: ends[0] - k))
            return (end, upper[c] && !STOP.contains(first))
        }
        return nil
    }

    /// finditer over [lo, hi): calls emit(start, end) for each accepted citation.
    func scan(_ lo: Int, _ hi: Int, _ emit: (Int, Int) -> Void) {
        var i = lo
        while i < hi {
            if isDigit(p[i]) && !wordBefore(i) {
                if let (e, up) = match(i) {
                    if up { emit(i, e) }
                    i = e
                    continue
                }
                while i < hi && isDigit(p[i]) { i += 1 }   // \b fails inside the rest of this run
                continue
            }
            i += 1
        }
    }
}

/// citenet/text.py STOP: words that open a match but never a reporter
let STOP: Set<[UInt8]> = Set("""
No. Nos. No Id. Id Ibid. Section Sec. In On At Page Pages Rule Rules Cir. Supp. Tr. Year Years
East West North South And Or The Of To CA Art. Article Para. Count Counts Exhibit Ex. Doc. Dkt. ECF
January February March April May June July August September October November December
Jan. Feb. Mar. Apr. Jun. Jul. Aug. Sep. Sept. Oct. Nov. Dec.
C Civ Civ. Cr Cr. Crim. Crim CV CR
""".split(whereSeparator: { $0 == " " || $0 == "\n" }).map { Array($0.utf8) })

// ---------------------------------------------------------------- normalize + encode
let OTHERB: UInt8 = 0xFF

/// Map the match bytes to one byte per character (normalize() steps 1-3: NFKC-ish,
/// quote/dash/section mapping, whitespace collapse + strip, lowercase).
@inline(__always) func charsOf(_ t: Text, _ s: Int, _ e: Int, into out: inout [UInt8]) {
    out.removeAll(keepingCapacity: true)
    var i = s
    var pendingSpace = false
    while i < e {
        let w = t.ws(i)
        if w > 0 { pendingSpace = !out.isEmpty; i += w; continue }
        var c: UInt8
        let b = t.p[i]
        if b < 0x80 {
            c = isUpper(b) ? b + 32 : (b == 0x60 ? 0x27 : b); i += 1
        } else {
            let b1 = t[i + 1], b2 = t[i + 2]
            if b == 0xE2 && b1 == 0x80 && (b2 == 0x98 || b2 == 0x99) { c = 0x27; i += 3 }
            else if b == 0xE2 && b1 == 0x80 && (b2 == 0x93 || b2 == 0x94) { c = 0x2D; i += 3 }
            else if b == 0xC2 && b1 == 0xA7 { c = 0x73; i += 2 }
            else {
                c = OTHERB
                i += 1
                while i < e && (t.p[i] & 0xC0) == 0x80 { i += 1 }
            }
        }
        if pendingSpace { out.append(0x20); pendingSpace = false }
        out.append(c)
    }
}

@inline(__always) func isWordB(_ c: UInt8) -> Bool { isLower(c) || isDigit(c) || c == 0x5F || c == OTHERB }

let EXPL: [[UInt8]] = ["quoting", "citing", "internal", "alteration", "emphasis", "quotation", "affirmed", "aff'd",
                       "rev'd", "cert", "vacated", "overruled", "abrogated", "superseded", "modified"].map { Array($0.utf8) }

/// _EXPL: \s*\((?:quoting|citing|...|\d+ )  -> cut the string at the leftmost match
@inline(__always) func cutExplanatory(_ s: inout [UInt8]) {
    let n = s.count
    var i = 0
    while i < n {
        if s[i] == 0x28 {
            var hit = false
            for w in EXPL where i + 1 + w.count <= n && s[(i + 1)..<(i + 1 + w.count)].elementsEqual(w) { hit = true; break }
            if !hit {
                var q = i + 1
                while q < n && isDigit(s[q]) { q += 1 }
                if q > i + 1 && q < n && s[q] == 0x20 { hit = true }
            }
            if hit {
                var cut = i
                while cut > 0 && s[cut - 1] == 0x20 { cut -= 1 }
                s.removeSubrange(cut..<n)
                return
            }
        }
        i += 1
    }
}

func normalizeChars(_ s: inout [UInt8], _ tmp: inout [UInt8]) {
    cutExplanatory(&s)
    // _TRAIL: (?<=[\d)])[\s.,;:]+$ | [\s,;:]+$   (leftmost position where either matches)
    var k = s.count
    while k > 0 && (s[k - 1] == 0x20 || s[k - 1] == 0x2E || s[k - 1] == 0x2C || s[k - 1] == 0x3B || s[k - 1] == 0x3A) { k -= 1 }
    if k < s.count {
        var cut = -1
        for pos in k..<s.count {
            let prev = pos > 0 ? s[pos - 1] : 0
            if isDigit(prev) || prev == 0x29 { cut = pos; break }
            var allNoDot = true
            for q in pos..<s.count where s[q] == 0x2E { allNoDot = false; break }
            if allNoDot { cut = pos; break }
        }
        if cut >= 0 { s.removeSubrange(cut..<s.count) }
    }
    // _JOIN: \.\s+(?=(?!at\b)[a-z]|\d+(?:d|th|st|nd|rd)\b) -> "."
    tmp.removeAll(keepingCapacity: true)
    var i = 0
    let n = s.count
    while i < n {
        let c = s[i]
        tmp.append(c)
        if c == 0x2E && i + 1 < n && s[i + 1] == 0x20 {
            let j = i + 2
            var join = false
            if j < n && isLower(s[j]) {
                let isAt = s[j] == 0x61 && j + 1 < n && s[j + 1] == 0x74 && !(j + 2 < n && isWordB(s[j + 2]))
                join = !isAt
            } else if j < n && isDigit(s[j]) {
                var q = j
                while q < n && isDigit(s[q]) { q += 1 }
                var e = -1
                if q < n && s[q] == 0x64 { e = q + 1 }
                else if q + 1 < n && ((s[q] == 0x74 && s[q + 1] == 0x68) || (s[q] == 0x73 && s[q + 1] == 0x74) ||
                                       (s[q] == 0x6E && s[q + 1] == 0x64) || (s[q] == 0x72 && s[q + 1] == 0x64)) { e = q + 2 }
                if e >= 0 && !(e < n && isWordB(s[e])) { join = true }
            }
            if join { i += 2; continue }
        }
        i += 1
    }
    swap(&s, &tmp)
    if s.count > MAXCHARS {
        let head = Array(s[0..<24]), tail = Array(s[(s.count - (MAXCHARS - 24))...])
        s = head + tail
    }
}

let CODE: [UInt8] = {
    var t = [UInt8](repeating: 47, count: 256)
    for d in 0..<10 { t[0x30 + d] = UInt8(3 + d) }
    for l in 0..<26 { t[0x61 + l] = UInt8(13 + l) }
    t[0x20] = 39; t[0x2E] = 40; t[0x2C] = 41; t[0x28] = 42; t[0x29] = 43; t[0x27] = 44; t[0x26] = 45; t[0x2D] = 46
    return t
}()

/// Write BOS + codes + EOS + PAD (SEQ Float16 values) for citation bytes [s, e).
@inline(__always) func encode(_ t: Text, _ s: Int, _ e: Int, _ dst: UnsafeMutablePointer<Float16>,
                              _ buf: inout [UInt8], _ tmp: inout [UInt8]) {
    charsOf(t, s, e, into: &buf)
    normalizeChars(&buf, &tmp)
    dst[0] = 1
    var k = 1
    for c in buf where k < SEQ - 1 { dst[k] = Float16(CODE[Int(c)]); k += 1 }
    dst[k] = 2
    k += 1
    while k < SEQ { dst[k] = 0; k += 1 }
}