#include "ImpulseMultiChannelConverter.h"

#include "IprhParser.h"

#include <QDebug>
#include <QDir>
#include <QFile>
#include <QFileInfo>
#include <QLocale>
#include <QMap>
#include <QRegularExpression>
#include <QSharedPointer>
#include <QStringList>
#include <QTextStream>
#include <algorithm>
#include <limits>

bool ImpulseMultiChannelConverter::isMultiChannelImpulseHeader(const QString &headerPath,
                                                               QString *dirPath,
                                                               QString *baseName)
{
    QFileInfo fi(headerPath);
    QRegularExpression reMultiChannel(QStringLiteral("^(.*)_A(\\d+)$"),
                                      QRegularExpression::CaseInsensitiveOption);
    QRegularExpressionMatch match = reMultiChannel.match(fi.completeBaseName());
    if (!match.hasMatch())
        return false;

    const QString surveyBase = match.captured(1);
    QDir dir(fi.absolutePath());
    const QStringList channels = dir.entryList(QStringList() << surveyBase + QStringLiteral("_A*.iprh"),
                                               QDir::Files);
    if (channels.size() <= 1)
        return false;

    if (dirPath)
        *dirPath = fi.absolutePath();
    if (baseName)
        *baseName = surveyBase;
    return true;
}

bool ImpulseMultiChannelConverter::buildPlan(const QString &dirPath,
                                             const QString &baseName,
                                             ConversionPlan &plan,
                                             QString *errorMessage)
{
    plan = ConversionPlan{};
    plan.dirPath = dirPath;
    plan.baseName = baseName;

    QDir dir(dirPath);
    dir.setFilter(QDir::Files | QDir::NoDotAndDotDot);

    QRegularExpression reChannel(QStringLiteral("^%1_A(\\d+)\\.iprh$").arg(QRegularExpression::escape(baseName)),
                                 QRegularExpression::CaseInsensitiveOption);
    QMap<int, QString> channelHeaders;
    for (const QFileInfo &fi : dir.entryInfoList()) {
        QRegularExpressionMatch match = reChannel.match(fi.fileName());
        if (!match.hasMatch())
            continue;

        const int chNum = match.captured(1).toInt();
        const QString iprhPath = fi.absoluteFilePath();
        QString iprbPath = iprhPath;
        iprbPath.replace(QStringLiteral(".iprh"), QStringLiteral(".iprb"), Qt::CaseInsensitive);
        if (QFile::exists(iprbPath)) {
            channelHeaders.insert(chNum, iprhPath);
        } else {
            qDebug() << "Missing .iprb for" << iprhPath;
        }
    }

    if (channelHeaders.isEmpty()) {
        if (errorMessage)
            *errorMessage = QStringLiteral("未找到 Impulse 多通道文件：%1").arg(baseName);
        return false;
    }

    GPRDataModel::Header masterHeader;
    if (!IprhParser::parseHeaderOnly(channelHeaders.first(), masterHeader)) {
        if (errorMessage)
            *errorMessage = QStringLiteral("无法解析主通道头文件：%1").arg(channelHeaders.first());
        return false;
    }

    QVector<float> xOffsets;
    QVector<float> yOffsets;
    xOffsets.reserve(channelHeaders.size());
    yOffsets.reserve(channelHeaders.size());

    qint64 minTraceCount = std::numeric_limits<qint64>::max();
    double antennaSeparation = 0.0;

    for (auto it = channelHeaders.begin(); it != channelHeaders.end(); ++it) {
        ChannelInfo info;
        info.channelNumber = it.key();
        info.iprhPath = it.value();
        info.iprbPath = info.iprhPath;
        info.iprbPath.replace(QStringLiteral(".iprh"), QStringLiteral(".iprb"), Qt::CaseInsensitive);

        if (!IprhParser::parseHeaderOnly(info.iprhPath, info.header)) {
            if (errorMessage)
                *errorMessage = QStringLiteral("无法解析通道头文件：%1").arg(info.iprhPath);
            return false;
        }

        if (info.header.samplesPerTrace != masterHeader.samplesPerTrace) {
            if (errorMessage)
                *errorMessage = QStringLiteral("多通道 SAMPLES 不一致：%1").arg(info.iprhPath);
            return false;
        }
        if (!qFuzzyCompare(info.header.timeIntervalNs, masterHeader.timeIntervalNs) &&
            info.header.timeIntervalNs > 0 && masterHeader.timeIntervalNs > 0) {
            qDebug() << "Warning: Inconsistent TIME INTERVAL across channels" << info.iprhPath;
        }
        if (!qFuzzyCompare(info.header.distanceInc, masterHeader.distanceInc) &&
            info.header.distanceInc > 0 && masterHeader.distanceInc > 0) {
            qDebug() << "Warning: Inconsistent DISTANCE INTERVAL across channels" << info.iprhPath;
        }

        if (info.header.rawParams.contains(QStringLiteral("CH_X_OFFSET"))) {
            info.xOffset = info.header.rawParams.value(QStringLiteral("CH_X_OFFSET")).toFloat();
        } else if (info.header.rawParams.contains(QStringLiteral("CH_OFFSET_X"))) {
            info.xOffset = info.header.rawParams.value(QStringLiteral("CH_OFFSET_X")).toFloat();
        } else if (!info.header.chXOffsets.isEmpty()) {
            info.xOffset = info.header.chXOffsets.first();
        }
        if (info.header.rawParams.contains(QStringLiteral("CH_Y_OFFSET"))) {
            info.yOffset = info.header.rawParams.value(QStringLiteral("CH_Y_OFFSET")).toFloat();
        } else if (info.header.rawParams.contains(QStringLiteral("CH_OFFSET_Y"))) {
            info.yOffset = info.header.rawParams.value(QStringLiteral("CH_OFFSET_Y")).toFloat();
        } else if (!info.header.chYOffsets.isEmpty()) {
            info.yOffset = info.header.chYOffsets.first();
        }
        xOffsets.append(info.xOffset);
        yOffsets.append(info.yOffset);

        if (info.header.rawParams.contains(QStringLiteral("ANTENNA SEPARATION"))) {
            antennaSeparation = info.header.rawParams.value(QStringLiteral("ANTENNA SEPARATION")).toDouble();
        }

        const qint64 traceBytes = static_cast<qint64>(info.header.samplesPerTrace) * sizeof(qint16);
        if (traceBytes <= 0) {
            if (errorMessage)
                *errorMessage = QStringLiteral("通道采样数无效：%1").arg(info.iprhPath);
            return false;
        }
        info.traceCount = QFileInfo(info.iprbPath).size() / traceBytes;
        if (info.traceCount <= 0) {
            if (errorMessage)
                *errorMessage = QStringLiteral("通道无有效数据：%1").arg(info.iprbPath);
            return false;
        }
        minTraceCount = std::min(minTraceCount, info.traceCount);
        plan.channels.append(info);
    }

    if (plan.channels.isEmpty() || minTraceCount <= 0 || minTraceCount == std::numeric_limits<qint64>::max()) {
        if (errorMessage)
            *errorMessage = QStringLiteral("Impulse 多通道没有有效 trace：%1").arg(baseName);
        return false;
    }

    for (const ChannelInfo &info : plan.channels) {
        if (info.traceCount != minTraceCount) {
            qDebug() << "Warning: Inconsistent trace count across channels. Using minimum:"
                     << minTraceCount << "channel has:" << info.traceCount << info.iprhPath;
        }
    }

    bool allZeroOffsets = true;
    for (float value : xOffsets) {
        if (!qFuzzyIsNull(value)) {
            allZeroOffsets = false;
            break;
        }
    }
    if (allZeroOffsets && antennaSeparation > 1e-6 && plan.channels.size() > 1) {
        const double totalWidth = (plan.channels.size() - 1) * antennaSeparation;
        const double startX = -totalWidth / 2.0;
        for (int i = 0; i < xOffsets.size(); ++i) {
            xOffsets[i] = static_cast<float>(startX + i * antennaSeparation);
            plan.channels[i].xOffset = xOffsets[i];
        }
        qDebug() << "Computed symmetric X offsets from antenna separation:" << antennaSeparation;
    }

    plan.channelCount = plan.channels.size();
    plan.tracesPerChannel = static_cast<int>(minTraceCount);
    plan.samplesPerTrace = masterHeader.samplesPerTrace;
    plan.traceByteSize = plan.samplesPerTrace * static_cast<qint64>(sizeof(qint16));
    plan.totalOutputTraces = static_cast<qint64>(plan.tracesPerChannel) * plan.channelCount;
    plan.outputHeader = masterHeader;
    plan.outputHeader.numberOfChannels = plan.channelCount;
    plan.outputHeader.numTraces = static_cast<int>(plan.totalOutputTraces);
    plan.outputHeader.chXOffsets = xOffsets;
    plan.outputHeader.chYOffsets = yOffsets;
    if (plan.outputHeader.timeWindowNs <= 0.0 && plan.outputHeader.timeIntervalNs > 0.0) {
        plan.outputHeader.timeWindowNs = plan.outputHeader.samplesPerTrace * plan.outputHeader.timeIntervalNs;
    }

    const QString basePath = dir.absoluteFilePath(baseName + QStringLiteral("_mala_converted"));
    plan.outputRadPath = basePath + QStringLiteral(".rad");
    plan.outputRd3Path = basePath + QStringLiteral(".rd3");
    return true;
}

bool ImpulseMultiChannelConverter::convertStreaming(const ConversionPlan &plan,
                                                    const Options &options,
                                                    QString *radFilePath,
                                                    QString *errorMessage,
                                                    CancelFn cancel,
                                                    ProgressFn progress)
{
    if (plan.channelCount <= 0 || plan.tracesPerChannel <= 0 || plan.traceByteSize <= 0) {
        if (errorMessage)
            *errorMessage = QStringLiteral("Impulse 转换计划无效");
        return false;
    }

    if (options.reuseExistingIfValid && convertedFilesAreValid(plan)) {
        if (radFilePath)
            *radFilePath = plan.outputRadPath;
        return true;
    }

    if (!options.overwriteExisting && (QFile::exists(plan.outputRadPath) || QFile::exists(plan.outputRd3Path))) {
        if (errorMessage)
            *errorMessage = QStringLiteral("转换输出文件已存在：%1").arg(plan.outputRadPath);
        return false;
    }

    const QString tmpRd3Path = plan.outputRd3Path + QStringLiteral(".tmp");
    const QString tmpRadPath = plan.outputRadPath + QStringLiteral(".tmp");
    QFile::remove(tmpRd3Path);
    QFile::remove(tmpRadPath);

    QVector<QSharedPointer<QFile>> inputFiles;
    inputFiles.reserve(plan.channels.size());
    for (const ChannelInfo &channel : plan.channels) {
        auto file = QSharedPointer<QFile>::create(channel.iprbPath);
        if (!file->open(QIODevice::ReadOnly)) {
            if (errorMessage)
                *errorMessage = QStringLiteral("无法读取 Impulse 通道数据：%1").arg(channel.iprbPath);
            return false;
        }
        inputFiles.append(file);
    }

    QFile rd3File(tmpRd3Path);
    if (!rd3File.open(QIODevice::WriteOnly | QIODevice::Truncate)) {
        if (errorMessage)
            *errorMessage = QStringLiteral("无法写入转换后的 RD3 文件：%1").arg(tmpRd3Path);
        return false;
    }

    const qint64 bytesPerTracePosition = plan.traceByteSize * plan.channelCount;
    const qint64 chunkBudget = qMax<qint64>(bytesPerTracePosition, options.maxChunkBytes);
    const int chunkTraces = qMax<qint64>(1, chunkBudget / bytesPerTracePosition);
    QVector<QByteArray> channelBuffers(plan.channelCount);

    qint64 tracesDone = 0;
    while (tracesDone < plan.tracesPerChannel) {
        if (cancel && cancel()) {
            rd3File.close();
            QFile::remove(tmpRd3Path);
            QFile::remove(tmpRadPath);
            if (errorMessage)
                *errorMessage = QStringLiteral("Impulse 多通道转换已取消");
            return false;
        }

        const int currentChunk = qMin<qint64>(chunkTraces, plan.tracesPerChannel - tracesDone);
        const qint64 expectedChannelBytes = currentChunk * plan.traceByteSize;
        for (int ch = 0; ch < plan.channelCount; ++ch) {
            channelBuffers[ch] = inputFiles[ch]->read(expectedChannelBytes);
            if (channelBuffers[ch].size() != expectedChannelBytes) {
                rd3File.close();
                QFile::remove(tmpRd3Path);
                if (errorMessage) {
                    *errorMessage = QStringLiteral("读取通道数据不完整：%1").arg(plan.channels[ch].iprbPath);
                }
                return false;
            }
        }

        for (int localTrace = 0; localTrace < currentChunk; ++localTrace) {
            const qint64 offset = localTrace * plan.traceByteSize;
            for (int ch = 0; ch < plan.channelCount; ++ch) {
                const qint64 written = rd3File.write(channelBuffers[ch].constData() + offset, plan.traceByteSize);
                if (written != plan.traceByteSize) {
                    rd3File.close();
                    QFile::remove(tmpRd3Path);
                    if (errorMessage)
                        *errorMessage = QStringLiteral("写入转换后的 RD3 文件失败：%1").arg(tmpRd3Path);
                    return false;
                }
            }
        }

        tracesDone += currentChunk;
        if (progress) {
            progress(Progress{tracesDone, plan.tracesPerChannel,
                              QStringLiteral("正在转换 Impulse 多通道 %1/%2")
                                  .arg(tracesDone)
                                  .arg(plan.tracesPerChannel)});
        }
    }

    rd3File.close();
    if (rd3File.error() != QFile::NoError) {
        QFile::remove(tmpRd3Path);
        if (errorMessage)
            *errorMessage = QStringLiteral("保存转换后的 RD3 文件失败：%1").arg(rd3File.errorString());
        return false;
    }

    ConversionPlan tmpPlan = plan;
    tmpPlan.outputRadPath = tmpRadPath;
    if (!writeRadHeader(tmpRadPath, tmpPlan, errorMessage)) {
        QFile::remove(tmpRd3Path);
        QFile::remove(tmpRadPath);
        return false;
    }

    QFile::remove(plan.outputRd3Path);
    QFile::remove(plan.outputRadPath);
    if (!QFile::rename(tmpRd3Path, plan.outputRd3Path)) {
        QFile::remove(tmpRd3Path);
        QFile::remove(tmpRadPath);
        if (errorMessage)
            *errorMessage = QStringLiteral("无法替换转换后的 RD3 文件：%1").arg(plan.outputRd3Path);
        return false;
    }
    if (!QFile::rename(tmpRadPath, plan.outputRadPath)) {
        QFile::remove(plan.outputRd3Path);
        QFile::remove(tmpRadPath);
        if (errorMessage)
            *errorMessage = QStringLiteral("无法替换转换后的 RAD 文件：%1").arg(plan.outputRadPath);
        return false;
    }

    if (radFilePath)
        *radFilePath = plan.outputRadPath;

    qDebug() << "Impulse multi-channel streaming conversion OK:" << plan.outputRadPath << plan.outputRd3Path;
    return true;
}

bool ImpulseMultiChannelConverter::writeRadHeader(const QString &radFilePath,
                                                  const ConversionPlan &plan,
                                                  QString *errorMessage)
{
    QFile radFile(radFilePath);
    if (!radFile.open(QIODevice::WriteOnly | QIODevice::Text | QIODevice::Truncate)) {
        if (errorMessage)
            *errorMessage = QStringLiteral("无法写入转换后的 RAD 文件：%1").arg(radFilePath);
        return false;
    }

    QTextStream rad(&radFile);
    const auto &header = plan.outputHeader;
    rad << "# Converted from Impulse multi-channel survey: " << plan.baseName << '\n';
    rad << "SAMPLES: " << header.samplesPerTrace << '\n';
    rad << "LAST TRACE: " << plan.totalOutputTraces << '\n';
    rad << "TIMEWINDOW: " << QLocale::c().toString(header.timeWindowNs, 'g', 12) << '\n';
    rad << "TIME INTERVAL: " << QLocale::c().toString(header.timeIntervalNs, 'g', 12) << '\n';
    rad << "DISTANCE INTERVAL: " << QLocale::c().toString(header.distanceInc, 'g', 12) << '\n';
    rad << "ANTENNAS: " << QLocale::c().toString(header.antennaFreq, 'g', 12) << '\n';
    if (!header.antennaType.isEmpty()) rad << "ANTENNA: " << header.antennaType << '\n';
    if (!header.date.isEmpty()) rad << "DATE: " << header.date << '\n';
    if (!header.timeStr.isEmpty()) rad << "TIME: " << header.timeStr << '\n';
    rad << "NUMBER_OF_CH: " << qMax(1, header.numberOfChannels) << '\n';
    rad << "CH_X_OFFSETS: " << formatFloatList(header.chXOffsets) << '\n';
    rad << "CH_Y_OFFSETS: " << formatFloatList(header.chYOffsets) << '\n';
    if (!header.units.isEmpty()) rad << "UNITS: " << header.units << '\n';
    rad << "START POSITION: " << QLocale::c().toString(header.startPosition, 'g', 12) << '\n';
    const double stopPosition = header.stopPosition > header.startPosition
            ? header.stopPosition
            : header.startPosition + plan.tracesPerChannel * header.distanceInc;
    rad << "STOP POSITION: " << QLocale::c().toString(stopPosition, 'g', 12) << '\n';
    radFile.close();
    return true;
}

bool ImpulseMultiChannelConverter::convertedFilesAreValid(const ConversionPlan &plan)
{
    QFileInfo radInfo(plan.outputRadPath);
    QFileInfo rd3Info(plan.outputRd3Path);
    if (!radInfo.exists() || !rd3Info.exists())
        return false;

    const qint64 expectedBytes = plan.totalOutputTraces * plan.traceByteSize;
    return rd3Info.size() == expectedBytes;
}

QString ImpulseMultiChannelConverter::formatFloatList(const QVector<float> &values)
{
    QStringList parts;
    parts.reserve(values.size());
    for (float value : values) {
        parts.append(QLocale::c().toString(value, 'g', 10));
    }
    return parts.join(' ');
}