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). Dedupe rule: a final for a
/// role closes any open partial row for that role.
@MainActor
final class TranscriptHandler: ObservableObject {
    @Published private(set) var events: [TranscriptEvent] = []

    let history = TranscriptHistory()

    func receiveSegment(id: String, role: TranscriptEvent.Role, text: String, isFinal: Bool) {
        guard !text.isEmpty else { return }
        upsert(id: id, role: role, text: text, isFinal: isFinal)
        if isFinal {
            history.add(role: role.rawValue, text: text)
        }
    }

    func receiveDataTranscript(role: TranscriptEvent.Role, text: String) {
        closePartial(for: role)
        upsert(id: "final-\(role.rawValue)-\(UUID().uuidString)", role: role, text: text, isFinal: true)
        history.add(role: role.rawValue, text: text)
    }

    func clear() {
        events.removeAll()
    }

    func clearHistory() {
        history.clear()
    }

    private func closePartial(for role: TranscriptEvent.Role) {
        for index in events.indices where !events[index].isFinal && events[index].role == role {
            events[index].isFinal = true
        }
    }

    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: Date()))
        }
    }
}

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
            for event in events {
                self?.receiveSegment(id: event.id, role: event.role, text: event.text, isFinal: event.isFinal)
            }
        }
    }
}