import Foundation
import LiveKit

struct TranscriptEvent: Equatable, Identifiable {
    enum Role: String {
        case user
        case agent
    }

    let id: String
    let role: Role
    var text: String
    var isFinal: Bool
    let timestamp: Date
}

/// Merges the primary transcription segments (Room delegate) with the
/// secondary `transcript` data topic (finals only).
///
/// **Single canonical collection:** `events` is the one visible transcript.
/// It is initialized from the persisted history snapshot at init time and
/// updated in place by live events. The view renders only `events`.
///
/// **Cross-source dedupe:** Records `(role, normalized text, source, timestamp)`
/// and suppresses only the *opposite-source* copy within `dedupeWindow` seconds.
/// Same-source repeated utterances are never suppressed, even inside the window.
@MainActor
final class TranscriptHandler: ObservableObject {
    /// The single canonical visible transcript. Initialized from persisted
    /// history; live events append/update here. Trimmed to 50 finals.
    @Published private(set) var events: [TranscriptEvent] = []

    var acceptsRoom: ((Room) -> Bool)?
    let history: TranscriptHistory

    // MARK: Cross-source dedupe state

    enum Source {
        case primary   // LiveKit transcription segments
        case data      // `transcript` data topic
    }

    struct RecentFinal {
        let source: Source
        let timestamp: Date
    }

    /// Keyed by normalized (role, text). Records the most recent final for
    /// each key and which source produced it. Only the opposite-source copy
    /// is suppressed within the window; same-source repeats pass through.
    private var recentFinals: [String: RecentFinal] = [:]
    static let dedupeWindow: TimeInterval = 5
    static let maxEvents = 50

    /// Injectable clock for deterministic tests. Defaults to `Date()`.
    var now: () -> Date = { Date() }

    convenience init() {
        self.init(history: TranscriptHistory())
    }

    init(history: TranscriptHistory) {
        self.history = history
        // Initialize the canonical events from the persisted snapshot so the
        // UI shows prior-session transcripts without a separate "restored" store.
        events = history.entries.map { entry in
            TranscriptEvent(
                id: "hist-\(entry.timestamp.timeIntervalSince1970)-\(entry.role)",
                role: entry.role == "agent" ? .agent : .user,
                text: entry.text,
                isFinal: true,
                timestamp: entry.timestamp
            )
        }
    }

    // MARK: Ingest

    func receiveSegment(id: String, role: TranscriptEvent.Role, text: String, isFinal: Bool) {
        guard !text.isEmpty else { return }
        if isFinal {
            // Cross-source dedupe: suppress only if the *data* source already
            // recorded this exact final within the window. Same-source (primary)
            // repeats are always allowed. When suppressing, consume (remove) the
            // matched record so a second same-source repeat from primary is not
            // also suppressed — the next arrival establishes a new pair.
            let key = Self.dedupeKey(role: role, text: text)
            if let recent = recentFinals[key],
               recent.source == .data,
               now().timeIntervalSince(recent.timestamp) < Self.dedupeWindow {
                // Consume: remove the matched opposite-source record.
                recentFinals[key] = nil
                return
            }
        }
        upsert(id: id, role: role, text: text, isFinal: isFinal)
        if isFinal {
            recordFinal(role: role, text: text, source: .primary)
            trimEvents()
            syncPersistence()
        }
    }

    func receiveDataTranscript(role: TranscriptEvent.Role, text: String) {
        guard !text.isEmpty else { return }
        // When a data final arrives for a role that has an open partial, replace
        // the partial's text with the final and mark it final in place. The data
        // topic is the authoritative final; the user sees one row either way.
        if let index = events.lastIndex(where: { !$0.isFinal && $0.role == role }) {
            events[index].text = text
            events[index].isFinal = true
            recordFinal(role: role, text: text, source: .data)
            trimEvents()
            syncPersistence()
            return
        }
        // No open partial: check whether the *primary* source already recorded
        // this exact final within the window (primary→data order). Same-source
        // (data) repeats are always allowed. When suppressing, consume (remove)
        // the matched record so a second same-source repeat from data is not
        // also suppressed — the next arrival establishes a new pair.
        let key = Self.dedupeKey(role: role, text: text)
        if let recent = recentFinals[key],
           recent.source == .primary,
           now().timeIntervalSince(recent.timestamp) < Self.dedupeWindow {
            // Consume: remove the matched opposite-source record.
            recentFinals[key] = nil
            return
        }
        upsert(id: "final-\(role.rawValue)-\(UUID().uuidString)", role: role, text: text, isFinal: true)
        recordFinal(role: role, text: text, source: .data)
        trimEvents()
        syncPersistence()
    }

    /// Mirrors the canonical finals into the persisted snapshot so the UI and
    /// storage never diverge. Same-source repeats are preserved verbatim.
    private func syncPersistence() {
        let finals = events.filter(\.isFinal).map { (role: $0.role.rawValue, text: $0.text, timestamp: $0.timestamp) }
        history.sync(finals: finals)
    }

    // MARK: Clear

    /// Clears visible rows, dedupe state, and persistence atomically.
    func clearAll() {
        events.removeAll()
        recentFinals.removeAll()
        history.clear()
    }

    // MARK: Private

    private func upsert(id: String, role: TranscriptEvent.Role, text: String, isFinal: Bool) {
        if let index = events.firstIndex(where: { $0.id == id }) {
            events[index].text = text
            events[index].isFinal = isFinal
        } else {
            events.append(TranscriptEvent(id: id, role: role, text: text, isFinal: isFinal, timestamp: now()))
        }
    }

    private func recordFinal(role: TranscriptEvent.Role, text: String, source: Source) {
        let key = Self.dedupeKey(role: role, text: text)
        recentFinals[key] = RecentFinal(source: source, timestamp: now())
        // Prune expired entries to bound the dictionary.
        let current = now()
        for (k, recent) in recentFinals where current.timeIntervalSince(recent.timestamp) >= Self.dedupeWindow {
            recentFinals[k] = nil
        }
    }

    /// Keeps at most `maxEvents` finals in the canonical collection. Partial
    /// (non-final) rows are never trimmed; only surplus old finals are dropped.
    private func trimEvents() {
        let finalCount = events.filter(\.isFinal).count
        guard finalCount > Self.maxEvents else { return }
        // Drop oldest finals until we're back to maxEvents.
        var toRemove = finalCount - Self.maxEvents
        var writeIndex = 0
        for event in events {
            if event.isFinal && toRemove > 0 {
                toRemove -= 1
                continue
            }
            events[writeIndex] = event
            writeIndex += 1
        }
        events.removeSubrange(writeIndex...)
    }

    private static func dedupeKey(role: TranscriptEvent.Role, text: String) -> String {
        let normalized = text.trimmingCharacters(in: .whitespacesAndNewlines)
            .replacingOccurrences(of: "\\s+", with: " ", options: .regularExpression)
        return "\(role.rawValue):\(normalized)"
    }
}

extension TranscriptHandler: RoomDelegate {
    nonisolated func room(_ room: Room, participant: Participant, trackPublication: TrackPublication, didReceiveTranscriptionSegments segments: [TranscriptionSegment]) {
        let role: TranscriptEvent.Role = participant.isAgent ? .agent : .user
        let isLocal = participant.identity?.stringValue == room.localParticipant.identity?.stringValue
        let resolvedRole: TranscriptEvent.Role = isLocal ? .user : role
        let events = segments.map {
            TranscriptEvent(id: $0.id, role: resolvedRole, text: $0.text, isFinal: $0.isFinal, timestamp: $0.lastReceivedTime)
        }
        Task { @MainActor [weak self] in
            guard let self, self.acceptsRoom?(room) ?? true else { return }
            for event in events {
                self.receiveSegment(id: event.id, role: event.role, text: event.text, isFinal: event.isFinal)
            }
        }
    }
}