ten31-transcripts/Ten31Transcripts/Visual/TimelineBuilder.swift

import Foundation

/// Turns noisy per-frame `SpeakerObservation`s into clean
/// `(start, end, name, confidence)` segments.
///
/// - Hysteresis: open a segment after `openFrames` consecutive speaking frames,
///   close after `closeFrames` quiet frames — rides out UI-cue lag/flicker.
/// - Overlaps allowed: each name is tracked independently (crosstalk).
/// - mic-VAD "self" spans are merged in as high-confidence segments.
/// - OCR name variants are normalized via an alias table.
///
/// Pure logic, no UI/capture deps — fully unit-testable offline.
final class TimelineBuilder {
    private let openFrames: Int
    private let closeFrames: Int
    private var aliases: [String: String] = [:]      // normalized variant -> canonical
    private var states: [String: NameState] = [:]
    private var lastFrameT: Double = 0
    private(set) var segments: [VisualTimeline.Segment] = []

    init(openFrames: Int = 2, closeFrames: Int = 2) {
        self.openFrames = max(1, openFrames)
        self.closeFrames = max(1, closeFrames)
    }

    /// Register that `variant` (e.g. "Sarah J") should map to `canonical`
    /// (e.g. "Sarah Jones").
    func addAlias(_ variant: String, canonical: String) {
        aliases[Self.normalize(variant)] = canonical
    }

    /// Ingest one frame's observations (all sharing time `t`). Names not present
    /// (or present but not speaking) count as a quiet frame for any open segment.
    func ingest(_ observations: [SpeakerObservation], at t: TimeInterval) {
        lastFrameT = t

        // Best confidence per canonical name that is speaking this frame.
        var speaking: [String: Double] = [:]
        for obs in observations where obs.speaking && !obs.name.isEmpty {
            let name = canonical(obs.name)
            speaking[name] = max(speaking[name] ?? 0, obs.confidence)
        }

        let names = Set(states.keys).union(speaking.keys)
        for name in names {
            var st = states[name] ?? NameState()
            if let conf = speaking[name] {
                if st.voiced == 0 { st.runStart = t }
                st.voiced += 1
                st.silent = 0
                st.lastVoicedT = t
                if !st.open && st.voiced >= openFrames {
                    st.open = true
                    st.segStart = st.runStart
                    st.confSum = 0
                    st.confN = 0
                }
                if st.open { st.confSum += conf; st.confN += 1 }
            } else {
                st.silent += 1
                st.voiced = 0
                if st.open && st.silent >= closeFrames {
                    closeSegment(name: name, state: st)
                    st.open = false
                }
            }
            states[name] = st
        }
    }

    /// Merge mic-VAD self spans (the user) as high-confidence segments.
    func mergeSelfSpans(_ spans: [VADSpan], selfName: String) {
        for span in spans where span.end > span.start {
            segments.append(.init(start: span.start, end: span.end,
                                  name: selfName, confidence: span.confidence, source: "mic_vad"))
        }
    }

    /// Force-close any open segments at `t` (used when a visual gap begins, so a
    /// segment isn't carried across the gap).
    func closeOpenSegments(at t: TimeInterval) {
        for (name, st) in states where st.open {
            closeSegment(name: name, state: st)
            states[name]?.open = false
            states[name]?.voiced = 0
            states[name]?.silent = 0
        }
    }

    /// Close any still-open segments at end of capture.
    func finish() {
        for (name, st) in states where st.open {
            closeSegment(name: name, state: st)
            states[name]?.open = false
        }
        segments.sort { $0.start < $1.start }
    }

    // MARK: - Internal

    private struct NameState {
        var voiced = 0
        var silent = 0
        var open = false
        var runStart: Double = 0
        var segStart: Double = 0
        var lastVoicedT: Double = 0
        var confSum: Double = 0
        var confN = 0
    }

    private func closeSegment(name: String, state st: NameState) {
        guard st.lastVoicedT > st.segStart else { return }
        let confidence = st.confN > 0 ? st.confSum / Double(st.confN) : 0.8
        segments.append(.init(start: st.segStart, end: st.lastVoicedT,
                              name: name, confidence: confidence, source: "vision"))
    }

    private func canonical(_ raw: String) -> String {
        let key = Self.normalize(raw)
        return aliases[key] ?? raw.trimmingCharacters(in: .whitespacesAndNewlines)
    }

    private static func normalize(_ s: String) -> String {
        s.lowercased().trimmingCharacters(in: .whitespacesAndNewlines)
    }
}