deepgramservice.cpp

#include "deepgramservice.h"

#include <QJsonArray>
#include <QJsonDocument>
#include <QJsonObject>
#include <QNetworkProxy>
#include <QNetworkRequest>
#include <QRegularExpression>
#include <QUrl>
#include <QUrlQuery>
#include <cstring>

DeepgramService::DeepgramService(QObject *parent)
    : SpeechService(parent) {
    connect(&socket, &QWebSocket::connected, this, &DeepgramService::handleConnected);
    connect(&socket, &QWebSocket::disconnected, this, &DeepgramService::handleDisconnected);
    connect(&socket, &QWebSocket::textMessageReceived, this, &DeepgramService::handleTextMessage);
    connect(&socket, &QWebSocket::errorOccurred, this, &DeepgramService::handleSocketError);

    // Deepgram recommends keepalive every 5-10 seconds
    keepaliveTimer.setInterval(8000);
    keepaliveTimer.setSingleShot(false);
    connect(&keepaliveTimer, &QTimer::timeout, this, &DeepgramService::sendKeepalive);
}

QString DeepgramService::providerId() const {
    return QStringLiteral("deepgram");
}

void DeepgramService::start(const AppConfig &config, const AudioSpec &spec) {
    start(config, spec, 0);
}

void DeepgramService::start(const AppConfig &config, const AudioSpec &spec, int sessionId) {
    serviceSessionId = sessionId;
    
    resetState();
    stopRequested = false;
    currentSettings = config.deepgram;
    currentSpec = spec;

    if (currentSettings.apiKey.trimmed().isEmpty()) {
        emit errorOccurred(tr("Deepgram API key is missing."));
        return;
    }

    if (currentSpec.sampleRate <= 0 || currentSpec.channels <= 0 || currentSpec.format.isEmpty()) {
        emit errorOccurred(tr("Invalid audio format settings."));
        return;
    }

    bool littleEndian = true;
    if (!parsePcmFormat(currentSpec.format, &inputBitsPerSample, &littleEndian, &inputSampleType)) {
        emit errorOccurred(tr("Unsupported audio format for Deepgram."));
        return;
    }
    if (!littleEndian) {
        emit errorOccurred(tr("Only little-endian PCM is supported."));
        return;
    }

    needsConversion = !(inputSampleType == SampleType::Signed && inputBitsPerSample == 16);
    sessionActive = true;

    // If already connected, just signal ready
    if (socket.state() == QAbstractSocket::ConnectedState) {
        emit readyForAudio();
        return;
    }

    // If currently connecting, wait for connection
    if (socket.state() == QAbstractSocket::ConnectingState) {
        return;
    }

    // Open new connection
    stopKeepaliveTimer();
    if (socket.state() != QAbstractSocket::UnconnectedState) {
        socket.abort();
    }

    if (!applyProxy(currentSettings.proxy))
        return;

    const QUrl url = buildWebSocketUrl(currentSettings, currentSpec);
    if (!url.isValid()) {
        emit errorOccurred(tr("Invalid Deepgram WebSocket URL."));
        return;
    }

    QNetworkRequest request(url);
    const QString authHeader = QStringLiteral("Token ") + currentSettings.apiKey.trimmed();
    request.setRawHeader("Authorization", authHeader.toUtf8());
    
    socket.open(request);
}

void DeepgramService::stop() {
    stopRequested = true;
    sessionActive = false;
    // Don't close the socket - keep it alive for next session
}

void DeepgramService::sendAudio(const QByteArray &data) {
    if (!sessionActive || socket.state() != QAbstractSocket::ConnectedState || data.isEmpty())
        return;
    
    // Convert audio if needed
    const QByteArray payload = needsConversion ? convertToPcm16(data) : data;
    if (payload.isEmpty())
        return;
    
    socket.sendBinaryMessage(payload);
    
    // Reset keepalive timer when audio is sent
    if (keepaliveTimer.isActive()) {
        keepaliveTimer.stop();
        keepaliveTimer.start();
    }
}

bool DeepgramService::isActive() const {
    return sessionActive && socket.state() == QAbstractSocket::ConnectedState;
}

void DeepgramService::disconnect() {
    sessionActive = false;
    stopKeepaliveTimer();
    
    if (socket.state() == QAbstractSocket::ConnectedState) {
        const QJsonObject closeMsg{{"type", "CloseStream"}};
        const QJsonDocument doc(closeMsg);
        socket.sendTextMessage(QString::fromUtf8(doc.toJson(QJsonDocument::Compact)));
        
        QTimer::singleShot(100, this, [this]() {
            socket.close();
        });
    } else if (socket.state() != QAbstractSocket::UnconnectedState) {
        socket.abort();
    }
}

void DeepgramService::handleConnected() {
    startKeepaliveTimer();
    emit readyForAudio();
}

void DeepgramService::handleDisconnected() {
    stopKeepaliveTimer();
    if (sessionActive) {
        // Unexpected disconnect during active session
        sessionActive = false;
        emit errorOccurred(tr("Connection lost."));
    }
}

void DeepgramService::handleTextMessage(const QString &message) {
    QJsonParseError parseError{};
    const QJsonDocument doc = QJsonDocument::fromJson(message.toUtf8(), &parseError);
    if (parseError.error != QJsonParseError::NoError || !doc.isObject()) {
        return;
    }

    const QJsonObject root = doc.object();
    const QString type = root.value("type").toString();
    
    // Handle error messages
    if (type == QStringLiteral("Error") || root.contains("err_code")) {
        QString errorMsg = root.value("description").toString();
        if (errorMsg.isEmpty()) {
            errorMsg = root.value("err_msg").toString(tr("Deepgram error."));
        }
        emit errorOccurred(errorMsg);
        return;
    }
    
    // Handle metadata (connection established)
    if (type == QStringLiteral("Metadata")) {
        // Connection is ready, already emitted readyForAudio in handleConnected
        return;
    }
    
    // Handle UtteranceEnd
    if (type == QStringLiteral("UtteranceEnd")) {
        // Could be used to finalize current utterance
        return;
    }
    
    // Handle SpeechStarted
    if (type == QStringLiteral("SpeechStarted")) {
        // Speech detection event
        return;
    }
    
    // Handle Results
    if (type == QStringLiteral("Results")) {
        const QJsonObject channel = root.value("channel").toObject();
        const QJsonArray alternatives = channel.value("alternatives").toArray();
        
        if (alternatives.isEmpty())
            return;
        
        const QJsonObject bestAlternative = alternatives.first().toObject();
        const QString transcript = bestAlternative.value("transcript").toString().trimmed();
        
        if (transcript.isEmpty())
            return;
        
        const bool isFinal = root.value("is_final").toBool(false);
        const bool speechFinal = root.value("speech_final").toBool(false);
        
        // If is_final or speech_final, accumulate to committed text
        if (isFinal || speechFinal) {
            committedText += transcript;
            if (!transcript.endsWith(QLatin1Char(' ')) && !transcript.endsWith(QLatin1Char('.'))) {
                committedText += QLatin1Char(' ');
            }
            currentText.clear();
        } else {
            // Interim result
            currentText = transcript;
        }
        
        // Emit combined result
        const QString fullText = committedText + currentText;
        if (!fullText.isEmpty()) {
            emit transcriptUpdated(fullText);
        }
    }
}

void DeepgramService::handleSocketError(QAbstractSocket::SocketError) {
    emit errorOccurred(socket.errorString());
}

void DeepgramService::resetState() {
    committedText.clear();
    currentText.clear();
    stopRequested = false;
    needsConversion = false;
    inputBitsPerSample = 0;
    inputSampleType = SampleType::Signed;
    sessionActive = false;
}

void DeepgramService::startKeepaliveTimer() {
    keepaliveTimer.start();
}

void DeepgramService::stopKeepaliveTimer() {
    keepaliveTimer.stop();
}

void DeepgramService::sendKeepalive() {
    if (socket.state() != QAbstractSocket::ConnectedState)
        return;
    
    const QJsonObject keepalive{{"type", "KeepAlive"}};
    const QJsonDocument doc(keepalive);
    socket.sendTextMessage(QString::fromUtf8(doc.toJson(QJsonDocument::Compact)));
}

bool DeepgramService::applyProxy(const QString &proxyString) {
    const QString trimmed = proxyString.trimmed();
    if (trimmed.isEmpty()) {
        socket.setProxy(QNetworkProxy::NoProxy);
        return true;
    }

    QUrl url(trimmed);
    if (url.scheme().isEmpty())
        url = QUrl(QStringLiteral("http://") + trimmed);

    QNetworkProxy proxy;
    if (url.scheme().startsWith("socks", Qt::CaseInsensitive))
        proxy.setType(QNetworkProxy::Socks5Proxy);
    else
        proxy.setType(QNetworkProxy::HttpProxy);

    proxy.setHostName(url.host());
    proxy.setPort(static_cast<quint16>(url.port(443)));
    proxy.setUser(url.userName());
    proxy.setPassword(url.password());

    if (proxy.hostName().isEmpty() || proxy.port() == 0) {
        emit errorOccurred(tr("Invalid proxy address."));
        socket.setProxy(QNetworkProxy::NoProxy);
        return false;
    }

    socket.setProxy(proxy);
    return true;
}

QUrl DeepgramService::buildWebSocketUrl(const DeepgramSettings &settings, const AudioSpec &spec) const {
    QString endpoint = settings.endpoint.trimmed();
    if (endpoint.isEmpty()) {
        endpoint = QStringLiteral("wss://api.deepgram.com/v1/listen");
    }
    
    QUrl url(endpoint);
    if (url.scheme().isEmpty()) {
        url = QUrl(QStringLiteral("wss://") + endpoint);
    }
    
    QUrlQuery query;
    
    // Model (required by some Deepgram configurations)
    QString model = settings.model.trimmed();
    if (!model.isEmpty()) {
        query.addQueryItem("model", model);
    } else {
        // Default to nova-3 if not specified
        query.addQueryItem("model", "nova-3");
    }
    
    // Language (BCP-47 tag)
    QString language = settings.language.trimmed();
    if (!language.isEmpty()) {
        query.addQueryItem("language", language);
    }
    
    // Audio encoding and sample rate
    query.addQueryItem("encoding", "linear16");
    query.addQueryItem("sample_rate", QString::number(spec.sampleRate));
    query.addQueryItem("channels", QString::number(spec.channels));
    
    // Punctuation
    if (settings.punctuate) {
        query.addQueryItem("punctuate", "true");
    }
    
    // Interim results
    if (settings.interimResults) {
        query.addQueryItem("interim_results", "true");
    }
    
    // Smart formatting
    if (settings.smartFormat) {
        query.addQueryItem("smart_format", "true");
    }
    
    url.setQuery(query);
    return url;
}

bool DeepgramService::parsePcmFormat(const QString &format, int *bitsPerSample,
                                     bool *littleEndian, SampleType *sampleType) const {
    if (bitsPerSample)
        *bitsPerSample = 0;
    if (littleEndian)
        *littleEndian = true;
    if (sampleType)
        *sampleType = SampleType::Signed;

    const QString trimmed = format.trimmed();
    if (trimmed.isEmpty())
        return false;

    QRegularExpression regex(QStringLiteral("^pcm_([usf])(\\d+)(le|be)?$"),
                             QRegularExpression::CaseInsensitiveOption);
    const QRegularExpressionMatch match = regex.match(trimmed);
    if (!match.hasMatch())
        return false;

    const QString type = match.captured(1).toLower();
    const int bits = match.captured(2).toInt();
    const QString endian = match.captured(3).toLower();

    if (!endian.isEmpty() && endian != QStringLiteral("le"))
        return false;

    if (sampleType) {
        if (type == QStringLiteral("f"))
            *sampleType = SampleType::Float;
        else if (type == QStringLiteral("u"))
            *sampleType = SampleType::Unsigned;
        else
            *sampleType = SampleType::Signed;
    }

    if (bitsPerSample)
        *bitsPerSample = bits;
    if (littleEndian)
        *littleEndian = true;

    return bits > 0;
}

QByteArray DeepgramService::convertToPcm16(const QByteArray &data) const {
    if (inputBitsPerSample <= 0)
        return QByteArray();

    const int bytesPerSample = inputBitsPerSample / 8;
    if (bytesPerSample <= 0)
        return QByteArray();

    const int sampleCount = data.size() / bytesPerSample;
    if (sampleCount <= 0)
        return QByteArray();

    QByteArray out(sampleCount * 2, 0);
    const char *src = data.constData();
    char *dst = out.data();

    auto writeSample = [dst](int index, qint16 value) {
        dst[index * 2] = static_cast<char>(value & 0xff);
        dst[index * 2 + 1] = static_cast<char>((value >> 8) & 0xff);
    };

    switch (inputSampleType) {
        case SampleType::Signed: {
            if (inputBitsPerSample == 8) {
                for (int i = 0; i < sampleCount; ++i) {
                    const qint8 sample = static_cast<qint8>(src[i]);
                    writeSample(i, static_cast<qint16>(sample) << 8);
                }
            } else if (inputBitsPerSample == 16) {
                for (int i = 0; i < sampleCount; ++i) {
                    const int offset = i * 2;
                    const qint16 sample = static_cast<qint16>(
                        static_cast<quint8>(src[offset]) |
                        (static_cast<quint8>(src[offset + 1]) << 8));
                    writeSample(i, sample);
                }
            } else if (inputBitsPerSample == 24) {
                for (int i = 0; i < sampleCount; ++i) {
                    const int offset = i * 3;
                    const quint32 packed = static_cast<quint8>(src[offset]) |
                                           (static_cast<quint8>(src[offset + 1]) << 8) |
                                           (static_cast<quint8>(src[offset + 2]) << 16);
                    qint32 value = (packed & 0x800000)
                                       ? static_cast<qint32>(packed | 0xFF000000)
                                       : static_cast<qint32>(packed);
                    writeSample(i, static_cast<qint16>(value >> 8));
                }
            } else if (inputBitsPerSample == 32) {
                for (int i = 0; i < sampleCount; ++i) {
                    const int offset = i * 4;
                    const qint32 value = static_cast<qint32>(
                        static_cast<quint8>(src[offset]) |
                        (static_cast<quint8>(src[offset + 1]) << 8) |
                        (static_cast<quint8>(src[offset + 2]) << 16) |
                        (static_cast<quint8>(src[offset + 3]) << 24));
                    writeSample(i, static_cast<qint16>(value >> 16));
                }
            } else {
                return QByteArray();
            }
            break;
        }
        case SampleType::Unsigned: {
            if (inputBitsPerSample == 8) {
                for (int i = 0; i < sampleCount; ++i) {
                    const quint8 sample = static_cast<quint8>(src[i]);
                    const qint16 value = static_cast<qint16>(static_cast<int>(sample) - 128) << 8;
                    writeSample(i, value);
                }
            } else if (inputBitsPerSample == 16) {
                for (int i = 0; i < sampleCount; ++i) {
                    const int offset = i * 2;
                    const quint16 sample = static_cast<quint16>(
                        static_cast<quint8>(src[offset]) |
                        (static_cast<quint8>(src[offset + 1]) << 8));
                    const qint32 value = static_cast<qint32>(sample) - 32768;
                    writeSample(i, static_cast<qint16>(value));
                }
            } else if (inputBitsPerSample == 24) {
                for (int i = 0; i < sampleCount; ++i) {
                    const int offset = i * 3;
                    const quint32 sample = static_cast<quint8>(src[offset]) |
                                           (static_cast<quint8>(src[offset + 1]) << 8) |
                                           (static_cast<quint8>(src[offset + 2]) << 16);
                    const qint32 value = static_cast<qint32>(sample) - 0x800000;
                    writeSample(i, static_cast<qint16>(value >> 8));
                }
            } else if (inputBitsPerSample == 32) {
                for (int i = 0; i < sampleCount; ++i) {
                    const int offset = i * 4;
                    const quint32 sample = static_cast<quint32>(
                        static_cast<quint8>(src[offset]) |
                        (static_cast<quint8>(src[offset + 1]) << 8) |
                        (static_cast<quint8>(src[offset + 2]) << 16) |
                        (static_cast<quint8>(src[offset + 3]) << 24));
                    const qint64 value = static_cast<qint64>(sample) - 0x80000000LL;
                    writeSample(i, static_cast<qint16>(value >> 16));
                }
            } else {
                return QByteArray();
            }
            break;
        }
        case SampleType::Float: {
            if (inputBitsPerSample == 32) {
                for (int i = 0; i < sampleCount; ++i) {
                    const int offset = i * 4;
                    float value = 0.0f;
                    memcpy(&value, src + offset, sizeof(float));
                    const float clamped = qBound(-1.0f, value, 1.0f);
                    writeSample(i, static_cast<qint16>(clamped * 32767.0f));
                }
            } else if (inputBitsPerSample == 64) {
                for (int i = 0; i < sampleCount; ++i) {
                    const int offset = i * 8;
                    double value = 0.0;
                    memcpy(&value, src + offset, sizeof(double));
                    const double clamped = qBound(-1.0, value, 1.0);
                    writeSample(i, static_cast<qint16>(clamped * 32767.0));
                }
            } else {
                return QByteArray();
            }
            break;
        }
    }

    return out;
}