#include "translationservice.h"

#include "apikeystore.h"
#include "qtpotatosettings.h"
#include "translationpresets.h"

#include <KLocalizedString>

#include <QCoroFuture>
#include <QCoroTimer>

#include <QCryptographicHash>
#include <QDir>
#include <QFileInfo>
#include <QHash>
#include <QLocale>
#include <QNetworkAccessManager>
#include <QPointer>
#include <QSaveFile>
#include <QStandardPaths>
#include <QtConcurrentRun>

#include <chrono>
#include <vector>

using namespace Qt::Literals::StringLiterals;
using namespace std::chrono_literals;

namespace
{
// Source lines before a batch that the model sees for context.
constexpr qsizetype ContextLines = 6;
constexpr int MaxAttempts = 3;
// Requests in flight at once. Batches only use source lines as context, so they don't
// depend on each other.
constexpr int ParallelRequests = 3;
// Bump when the prompt changes so old cached translations are not reused.
constexpr auto CacheVersion = "llm-v2";

// "Vietnamese" -> "vi", so the file ends in ".vi.srt" like other players expect.
QString languageSuffix(const QString &language)
{
    for (int i = QLocale::C + 1; i <= QLocale::LastLanguage; ++i) {
        const auto candidate = static_cast<QLocale::Language>(i);
        if (QLocale::languageToString(candidate).compare(language, Qt::CaseInsensitive) == 0) {
            const QString code = QLocale::languageToCode(candidate, QLocale::ISO639Part1);
            return code.isEmpty() ? QLocale::languageToCode(candidate) : code;
        }
    }
    QString suffix = language.toLower();
    suffix.removeIf([](QChar c) {
        return !c.isLetterOrNumber() && c != u'-';
    });
    return suffix.isEmpty() ? u"translated"_s : suffix;
}

QString cachePath(const QUrl &videoUrl, const QString &key, const QString &language)
{
    const QString dir = QStandardPaths::writableLocation(QStandardPaths::CacheLocation) + u"/translations"_s;
    const QString baseName = videoUrl.isLocalFile() ? QFileInfo(videoUrl.toLocalFile()).completeBaseName() : u"subtitles"_s;
    return u"%1/%2.%3.%4.srt"_s.arg(dir, baseName, key, languageSuffix(language));
}
}

struct TranslationService::Job {
    std::atomic_bool cancelled = false;
    std::unique_ptr<LlmClient> client;
    QString mediaTitle;
};

struct TranslationService::Batches {
    QStringList lines;
    QStringList translations;
    qsizetype batchSize = 0;
    qsizetype nextStart = 0;
    qsizetype doneLines = 0;
    QString error;
};

TranslationService::TranslationService(QObject *parent)
    : QObject(parent)
    , m_network(new QNetworkAccessManager(this))
{
}

TranslationService::~TranslationService()
{
    cancel();
}

bool TranslationService::isBusy() const
{
    return m_busy;
}

double TranslationService::progress() const
{
    return m_progress;
}

QString TranslationService::statusText() const
{
    return m_statusText;
}

QString TranslationService::errorText() const
{
    return m_errorText;
}

QStringList TranslationService::availableModels() const
{
    return m_availableModels;
}

QVariantList TranslationService::translationPresets() const
{
    QVariantList result;
    for (const auto &preset : TranslationPresets::all()) {
        result.append(QVariantMap{{u"id"_s, preset.id}, {u"name"_s, preset.name}, {u"description"_s, preset.description}, {u"guidance"_s, preset.guidance}});
    }
    return result;
}

bool TranslationService::isFetchingModels() const
{
    return m_fetchingModels;
}

QVariantList TranslationService::glossaryEntries() const
{
    QVariantList result;
    for (const auto &entry : m_glossary.entries) result.append(QVariantMap{{u"source"_s, entry.source}, {u"target"_s, entry.target}});
    return result;
}

QString TranslationService::glossaryDirectory() const
{
    return m_glossary.directory;
}

QString TranslationService::glossaryError() const
{
    return m_glossaryError;
}

void TranslationService::loadGlossary(const QUrl &videoUrl)
{
    const QString directory = videoUrl.isLocalFile() ? QFileInfo(videoUrl.toLocalFile()).absolutePath() : QString();
    m_glossary = directory.isEmpty() ? GlossaryStore::Glossary{} : GlossaryStore::load(directory);
    if (!m_glossaryError.isEmpty()) {
        m_glossaryError.clear();
        Q_EMIT glossaryErrorChanged();
    }
    Q_EMIT glossaryChanged();
}

bool TranslationService::saveGlossary(const QVariantList &entries)
{
    QList<GlossaryStore::Entry> result;
    for (const QVariant &value : entries) {
        const QVariantMap map = value.toMap();
        result.append({map.value(u"source"_s).toString().trimmed(), map.value(u"target"_s).toString().trimmed()});
    }

    GlossaryStore::Glossary candidate = m_glossary;
    candidate.entries = result;
    if (!candidate.directory.isEmpty()) {
        QString error;
        if (!GlossaryStore::save(candidate, &error)) {
            const QString message = error.isEmpty() ? i18n("Could not save the translation glossary.")
                                                     : i18n("Could not save the translation glossary: %1", error);
            if (m_glossaryError != message) {
                m_glossaryError = message;
                Q_EMIT glossaryErrorChanged();
            }
            return false;
        }
    }
    m_glossary = candidate;
    if (!m_glossaryError.isEmpty()) {
        m_glossaryError.clear();
        Q_EMIT glossaryErrorChanged();
    }
    Q_EMIT glossaryChanged();
    return true;
}

void TranslationService::clearGlossary()
{
    m_glossary.entries.clear();
    if (!m_glossary.directory.isEmpty()) {
        GlossaryStore::save(m_glossary);
    }
    Q_EMIT glossaryChanged();
}

void TranslationService::translate(const QUrl &videoUrl, const QVariantMap &track, const QString &mediaTitle)
{
    cancel();
    m_job = std::make_shared<Job>();
    m_job->mediaTitle = mediaTitle;
    run(videoUrl, track);
}

void TranslationService::cancel()
{
    if (!m_job) {
        return;
    }
    // The running coroutine keeps its own reference and stops at its next check.
    m_job->cancelled = true;
    if (m_job->client) {
        m_job->client->abort();
    }
    m_job.reset();
    setBusy(false);
    setStatusText({});
}

void TranslationService::fetchModels()
{
    if (!m_fetchingModels) {
        runFetchModels();
    }
}

void TranslationService::clearError()
{
    setErrorText({});
}

QCoro::Task<std::pair<LlmClient::Config, QString>> TranslationService::loadConfig()
{
    const auto settings = QtPotatoSettings::self();
    LlmClient::Config config{
        LlmClient::normalizedBaseUrl(settings->apiBaseUrl()),
        {},
        settings->model().trimmed(),
        settings->targetLanguage().trimmed(),
        settings->translationPreset().trimmed(),
        settings->customInstructions(),
        {},
    };

    if (!config.baseUrl.isValid() || config.baseUrl.scheme().isEmpty() || config.baseUrl.host().isEmpty()) {
        co_return std::pair{config, i18n("Set the API base URL in Settings first.")};
    }

    const ApiKeyStore::ReadResult key = co_await ApiKeyStore::readKey();
    if (!key.error.isEmpty()) {
        co_return std::pair{config, i18n("Could not read the API key from the wallet: %1", key.error)};
    }
    config.apiKey = key.key;
    co_return std::pair{config, QString()};
}

QCoro::Task<Subtitles::ExtractResult> TranslationService::loadCues(std::shared_ptr<Job> job, QUrl videoUrl, QVariantMap track)
{
    if (track.value(u"external"_s).toBool()) {
        const QString path = track.value(u"externalFilename"_s).toString();
        co_return co_await QtConcurrent::run([path] {
            QString error;
            QList<SubtitleCue> cues = Subtitles::parseFile(path, &error);
            return Subtitles::ExtractResult{cues, error};
        });
    }

    if (!videoUrl.isLocalFile()) {
        co_return Subtitles::ExtractResult{{}, i18n("Only subtitles of local files can be translated.")};
    }
    const int streamIndex = track.value(u"ffIndex"_s, -1).toInt();
    if (streamIndex < 0) {
        co_return Subtitles::ExtractResult{{}, i18n("mpv did not report where this subtitle track is stored.")};
    }

    const QString path = videoUrl.toLocalFile();
    co_return co_await QtConcurrent::run([path, streamIndex, job] {
        return Subtitles::extractEmbedded(path, streamIndex, job->cancelled);
    });
}

QCoro::Task<> TranslationService::run(QUrl videoUrl, QVariantMap track)
{
    const std::shared_ptr<Job> job = m_job;

    setErrorText({});
    setProgress(0);
    setBusy(true);
    setStatusText(i18n("Reading subtitles…"));

    auto [config, configError] = co_await loadConfig();
    if (job->cancelled) {
        co_return;
    }
    if (!configError.isEmpty()) {
        fail(job, configError);
        co_return;
    }
    if (config.model.isEmpty()) {
        fail(job, i18n("Choose a model in Settings first."));
        co_return;
    }

    Subtitles::ExtractResult source = co_await loadCues(job, videoUrl, track);
    if (job->cancelled) {
        co_return;
    }
    if (!source.error.isEmpty()) {
        fail(job, source.error);
        co_return;
    }
    const QList<SubtitleCue> cues = Subtitles::mergeRepeated(source.cues);
    if (cues.isEmpty()) {
        fail(job, i18n("This subtitle track has no text to translate."));
        co_return;
    }

    // Episodes in one folder share a glossary. Streams have no folder, so they get none.
    const QString seriesDirectory = videoUrl.isLocalFile() ? QFileInfo(videoUrl.toLocalFile()).absolutePath() : QString();
    const bool useGlossary = !seriesDirectory.isEmpty();
    GlossaryStore::Glossary glossary = useGlossary ? GlossaryStore::load(seriesDirectory) : GlossaryStore::Glossary{};
    const QString sourceSrt = Subtitles::toSrt(cues);

    // Same subtitles, model, language, prompt settings and glossary give the same file, so it is translated once.
    const auto outputPathFor = [&](const GlossaryStore::Glossary &g) {
        QCryptographicHash hash(QCryptographicHash::Sha1);
        hash.addData(CacheVersion);
        hash.addData(config.model.toUtf8());
        hash.addData(config.targetLanguage.toUtf8());
        hash.addData(config.translationPreset.toUtf8());
        hash.addData(config.customInstructions.toUtf8());
        hash.addData(GlossaryStore::serialize(g).toUtf8());
        hash.addData(sourceSrt.toUtf8());
        return cachePath(videoUrl, QString::fromLatin1(hash.result().toHex().left(12)), config.targetLanguage);
    };
    QString outputPath = outputPathFor(glossary);

    if (QFileInfo::exists(outputPath)) {
        m_job.reset();
        setBusy(false);
        setStatusText({});
        Q_EMIT translated(QUrl::fromLocalFile(outputPath), true);
        co_return;
    }

    if (useGlossary && QtPotatoSettings::self()->autoBuildGlossary()) {
        setStatusText(i18n("Building glossary…"));
        QStringList extractionLines;
        QHash<QString, bool> seen;
        for (const SubtitleCue &cue : cues) {
            if (!seen.contains(cue.text)) { seen.insert(cue.text, true); extractionLines.append(cue.text); }
        }
        constexpr qsizetype MaxGlossaryChars = 30000;
        // Long subtitles are sampled evenly so names from the whole episode are seen.
        qsizetype allChars = 0;
        for (const QString &line : std::as_const(extractionLines)) {
            allChars += line.size() + 1;
        }
        const qsizetype step = std::max<qsizetype>(1, (allChars + MaxGlossaryChars - 1) / MaxGlossaryChars);
        QStringList sampled;
        qsizetype total = 0;
        for (qsizetype i = 0; i < extractionLines.size(); i += step) {
            if (total + extractionLines.at(i).size() + 1 > MaxGlossaryChars) {
                break;
            }
            sampled.append(extractionLines.at(i));
            total += extractionLines.at(i).size() + 1;
        }
        job->client = std::make_unique<LlmClient>(m_network, config);
        const LlmClient::BatchResult extracted = co_await job->client->extractGlossary(sampled, job->mediaTitle);
        if (job->cancelled) {
            co_return;
        }
        if (extracted.error.isEmpty()) {
            QList<GlossaryStore::Entry> additions;
            for (const QString &encoded : extracted.translations) {
                const qsizetype tab = encoded.indexOf(u'\t');
                if (tab > 0) additions.append({encoded.left(tab), encoded.mid(tab + 1)});
            }
            glossary = GlossaryStore::merge(glossary, additions);
            GlossaryStore::save(glossary);
        } else {
            setStatusText(i18n("Glossary unavailable; translating without it."));
        }
        outputPath = outputPathFor(glossary);
        if (QFileInfo::exists(outputPath)) {
            m_job.reset();
            setBusy(false);
            setStatusText({});
            Q_EMIT translated(QUrl::fromLocalFile(outputPath), true);
            co_return;
        }
    }

    // Repeated lines (songs, catchphrases) are sent once.
    QStringList uniqueLines;
    QHash<QString, qsizetype> uniqueIndex;
    QList<qsizetype> cueToUnique;
    for (const SubtitleCue &cue : cues) {
        auto it = uniqueIndex.constFind(cue.text);
        if (it == uniqueIndex.cend()) {
            it = uniqueIndex.insert(cue.text, uniqueLines.size());
            uniqueLines.append(cue.text);
        }
        cueToUnique.append(*it);
    }

    QStringList glossaryLines;
    for (const auto &entry : std::as_const(glossary.entries)) {
        glossaryLines.append(entry.source + u" → "_s + entry.target);
    }
    config.glossaryText = glossaryLines.join(u'\n');

    job->client = std::make_unique<LlmClient>(m_network, config);
    auto batches = std::make_shared<Batches>();
    batches->lines = uniqueLines;
    batches->translations.resize(uniqueLines.size());
    batches->batchSize = QtPotatoSettings::self()->batchSize();
    setStatusText(i18n("Translating… %1 of %2 lines", 0, uniqueLines.size()));

    // Tasks start running as soon as they are created; awaiting them just waits for the end.
    std::vector<QCoro::Task<>> workers;
    for (int i = 0; i < ParallelRequests; ++i) {
        workers.push_back(translateBatches(job, batches));
    }
    for (QCoro::Task<> &worker : workers) {
        co_await std::move(worker);
    }
    if (job->cancelled) {
        co_return;
    }
    if (!batches->error.isEmpty()) {
        fail(job, batches->error);
        co_return;
    }

    QList<SubtitleCue> output = cues;
    for (qsizetype i = 0; i < output.size(); ++i) {
        output[i].text = batches->translations.at(cueToUnique.at(i));
    }

    QDir().mkpath(QFileInfo(outputPath).absolutePath());
    QSaveFile file(outputPath);
    if (!file.open(QIODevice::WriteOnly) || file.write(Subtitles::toSrt(output).toUtf8()) < 0 || !file.commit()) {
        fail(job, i18n("Could not save the translation to %1: %2", outputPath, file.errorString()));
        co_return;
    }

    m_job.reset();
    setBusy(false);
    setStatusText({});
    Q_EMIT translated(QUrl::fromLocalFile(outputPath), false);
}

QCoro::Task<LlmClient::BatchResult> TranslationService::translateLines(std::shared_ptr<Job> job, QStringList lines, QStringList context)
{
    LlmClient::BatchResult result;
    for (int attempt = 1; attempt <= MaxAttempts; ++attempt) {
        result = co_await job->client->translate(lines, context, job->mediaTitle);
        if (job->cancelled || result.error.isEmpty() || !result.retryable || attempt == MaxAttempts) {
            break;
        }
        co_await QCoro::sleepFor(2s * attempt);
        if (job->cancelled) {
            co_return result;
        }
    }
    if (job->cancelled || !result.error.isEmpty()) {
        co_return result;
    }

    QList<qsizetype> missing;
    for (qsizetype i = 0; i < result.translations.size(); ++i) {
        if (result.translations.at(i).isNull()) {
            missing.append(i);
        }
    }
    if (missing.isEmpty()) {
        co_return result;
    }
    if (lines.size() == 1) {
        co_return LlmClient::BatchResult{{}, i18n("The model did not translate this line: %1", lines.first()), false};
    }

    if (missing.size() == lines.size()) {
        // Nothing came back usable; smaller requests are easier for the model.
        const qsizetype half = lines.size() / 2;
        const LlmClient::BatchResult first = co_await translateLines(job, lines.first(half), context);
        if (job->cancelled || !first.error.isEmpty()) {
            co_return first;
        }
        QStringList secondContext = context + lines.first(half);
        secondContext = secondContext.mid(std::max<qsizetype>(0, secondContext.size() - ContextLines));
        const LlmClient::BatchResult second = co_await translateLines(job, lines.sliced(half), secondContext);
        if (job->cancelled || !second.error.isEmpty()) {
            co_return second;
        }
        co_return LlmClient::BatchResult{first.translations + second.translations, {}, false};
    }

    // Ask again only for the lines the model skipped.
    QStringList retryLines;
    for (const qsizetype i : missing) {
        retryLines.append(lines.at(i));
    }
    const LlmClient::BatchResult retry = co_await translateLines(job, retryLines, context);
    if (job->cancelled || !retry.error.isEmpty()) {
        co_return retry;
    }
    for (qsizetype i = 0; i < missing.size(); ++i) {
        result.translations[missing.at(i)] = retry.translations.at(i);
    }
    co_return result;
}

QCoro::Task<> TranslationService::translateBatches(std::shared_ptr<Job> job, std::shared_ptr<Batches> batches)
{
    // After an error, requests already sent may finish, but no new ones start.
    while (!job->cancelled && batches->error.isEmpty() && batches->nextStart < batches->lines.size()) {
        const qsizetype start = batches->nextStart;
        batches->nextStart += batches->batchSize;

        const qsizetype contextStart = std::max<qsizetype>(0, start - ContextLines);
        const LlmClient::BatchResult result =
            co_await translateLines(job, batches->lines.mid(start, batches->batchSize), batches->lines.mid(contextStart, start - contextStart));
        if (job->cancelled) {
            co_return;
        }
        if (!result.error.isEmpty()) {
            if (batches->error.isEmpty()) {
                batches->error = result.error;
            }
            co_return;
        }

        for (qsizetype i = 0; i < result.translations.size(); ++i) {
            batches->translations[start + i] = result.translations.at(i);
        }
        batches->doneLines += result.translations.size();
        setProgress(double(batches->doneLines) / batches->lines.size());
        setStatusText(i18n("Translating… %1 of %2 lines", batches->doneLines, batches->lines.size()));
    }
}

QCoro::Task<> TranslationService::runFetchModels()
{
    QPointer<TranslationService> self(this);
    setFetchingModels(true);
    setErrorText({});

    const auto [config, configError] = co_await loadConfig();
    if (!self) {
        co_return;
    }
    if (!configError.isEmpty()) {
        setErrorText(configError);
        setFetchingModels(false);
        co_return;
    }

    LlmClient client(m_network, config);
    const LlmClient::ModelsResult result = co_await client.listModels();
    if (!self) {
        co_return;
    }
    if (!result.error.isEmpty()) {
        setErrorText(i18n("Could not list models: %1", result.error));
    } else if (result.models != m_availableModels) {
        m_availableModels = result.models;
        Q_EMIT availableModelsChanged();
    }
    setFetchingModels(false);
}

void TranslationService::fail(const std::shared_ptr<Job> &job, const QString &error)
{
    if (m_job != job) {
        return;
    }
    m_job.reset();
    setErrorText(error);
    setBusy(false);
    setStatusText({});
}

void TranslationService::setBusy(bool busy)
{
    if (busy != m_busy) {
        m_busy = busy;
        Q_EMIT busyChanged();
    }
}

void TranslationService::setProgress(double progress)
{
    if (progress != m_progress) {
        m_progress = progress;
        Q_EMIT progressChanged();
    }
}

void TranslationService::setStatusText(const QString &text)
{
    if (text != m_statusText) {
        m_statusText = text;
        Q_EMIT statusTextChanged();
    }
}

void TranslationService::setErrorText(const QString &text)
{
    if (text != m_errorText) {
        m_errorText = text;
        Q_EMIT errorTextChanged();
    }
}

void TranslationService::setFetchingModels(bool fetching)
{
    if (fetching != m_fetchingModels) {
        m_fetchingModels = fetching;
        Q_EMIT fetchingModelsChanged();
    }
}