#include "AiAudioInputPort.h"
#include "AiHttpClient.h"
#include <objbase.h>
#include <algorithm>
#include <sstream>
#include <cstring>
#include <cmath>

static short AiClamp16(int v)
{
    if (v > 32767) return 32767;
    if (v < -32768) return -32768;
    return (short)v;
}

static int AiBytesPerSample(int bitsPerSample)
{
    if (bitsPerSample == 8) return 1;
    if (bitsPerSample == 16) return 2;
    if (bitsPerSample == 24) return 3;
    if (bitsPerSample == 32) return 4;
    return 0;
}

static int AiReadS24LE(const unsigned char* p)
{
    int v = (int)p[0] | ((int)p[1] << 8) | ((int)p[2] << 16);
    if (v & 0x800000) v |= ~0xFFFFFF;
    return v;
}

static bool AiLooksLikeFloat32(const unsigned char* data, int dataSize, int channels)
{
    if (!data || dataSize < 4 || channels <= 0) return false;
    int samples = dataSize / 4;
    if (samples > 512) samples = 512;
    int plausible = 0;
    int nonZero = 0;
    for (int i = 0; i < samples; ++i) {
        float f = 0.0f;
        memcpy(&f, data + (size_t)i * 4, 4);
        if (std::isfinite(f) && f >= -1.25f && f <= 1.25f) ++plausible;
        if (std::fabs(f) > 1.0e-8f) ++nonZero;
    }
    // Silence is all zero and can be treated as either int32 or float32.  When non-zero,
    // a real WASAPI float stream normally has most values in [-1, 1].
    if (nonZero == 0) return true;
    return plausible >= (samples * 9) / 10;
}

static short AiReadSampleAsS16(const unsigned char* p, int bytesPerSample, int bitsPerSample, bool assumeFloat32)
{
    if (!p) return 0;
    int sample = 0;
    if (bitsPerSample == 8 && bytesPerSample == 1) {
        sample = ((int)(*p) - 128) << 8;
    } else if (bitsPerSample == 16 && bytesPerSample == 2) {
        sample = *(const short*)p;
    } else if (bitsPerSample == 24 && bytesPerSample == 3) {
        sample = AiReadS24LE(p) >> 8;
    } else if (bitsPerSample == 32 && bytesPerSample == 4) {
        if (assumeFloat32) {
            float f = 0.0f;
            memcpy(&f, p, 4);
            if (!std::isfinite(f)) f = 0.0f;
            if (f > 1.0f) f = 1.0f;
            if (f < -1.0f) f = -1.0f;
            sample = (int)(f * 32767.0f);
        } else {
            sample = (*(const int*)p) >> 16;
        }
    }
    return AiClamp16(sample);
}

AiAudioInputPort::AiAudioInputPort()
    : pendingStartTimestamp100ns_(0), nextChunkIndex_(0), running_(false), stopping_(false),
      inputFormatLogged_(false), pushedInputFrames_(0), convertedInputSamples_(0),
      droppedQueueChunks_(0), maxObservedQueueDepth_(0), expectedNextInputTimestamp100ns_(0),
      timestampGapWarnings_(0), producedChunks_(0), sentChunks_(0), sendFailures_(0),
      receivedResults_(0), maxSendRoundtripMs_(0), lastProducerReportChunk_(0),
      lastSenderReportChunk_(0)
{
}

AiAudioInputPort::~AiAudioInputPort()
{
    Stop();
}

void AiAudioInputPort::SetResultCallback(ResultCallback cb)
{
    std::lock_guard<std::mutex> guard(lock_);
    resultCb_ = cb;
}

void AiAudioInputPort::Start(const AiAudioInputConfig& cfg)
{
    Stop();
    {
        std::lock_guard<std::mutex> guard(lock_);
        config_ = cfg;
        if (config_.endpoint.empty()) config_.endpoint = L"ws://127.0.0.1:8766/v1/ai/stream";
        if (config_.targetSampleRate <= 0) config_.targetSampleRate = 16000;
        // Step10: use smaller chunks for real-time ASR while keeping ASR windows on the service side.
        if (config_.chunkMs <= 0) config_.chunkMs = 200;
        if (config_.chunkMs < 100) config_.chunkMs = 100;
        if (config_.chunkMs > 1000) config_.chunkMs = 1000;
        pendingSamples_.clear();
        pendingStartTimestamp100ns_ = 0;
        queue_.clear();
        nextChunkIndex_ = 0;
        stopping_.store(false);
        running_.store(true);

        inputFormatLogged_ = false;
        pushedInputFrames_ = 0;
        convertedInputSamples_ = 0;
        droppedQueueChunks_ = 0;
        maxObservedQueueDepth_ = 0;
        expectedNextInputTimestamp100ns_ = 0;
        timestampGapWarnings_ = 0;
        producedChunks_ = 0;
        sentChunks_ = 0;
        sendFailures_ = 0;
        receivedResults_ = 0;
        maxSendRoundtripMs_ = 0;
        lastProducerReportChunk_ = 0;
        lastSenderReportChunk_ = 0;
    }
    worker_ = std::thread(&AiAudioInputPort::WorkerThread, this);
}

void AiAudioInputPort::Stop()
{
    if (!running_.load() && !worker_.joinable()) return;
    stopping_.store(true);
    running_.store(false);
    cv_.notify_all();
    if (worker_.joinable()) worker_.join();
    std::lock_guard<std::mutex> guard(lock_);
    queue_.clear();
    pendingSamples_.clear();
    pendingStartTimestamp100ns_ = 0;
}

bool AiAudioInputPort::IsRunning() const
{
    return running_.load();
}

void AiAudioInputPort::PostDebugResult(const std::string& text, long long timestamp100ns, long long chunkIndex)
{
    ResultCallback cb;
    {
        std::lock_guard<std::mutex> guard(lock_);
        cb = resultCb_;
    }
    if (!cb) return;
    AiDemoResult r;
    r.typeUtf8 = "audio.debug";
    r.ok = true;
    r.statusUtf8 = "ok";
    r.timestamp100ns = timestamp100ns;
    r.startTimestamp100ns = timestamp100ns;
    r.endTimestamp100ns = timestamp100ns;
    r.chunkIndex = chunkIndex;
    r.summaryUtf8 = text;
    cb(r);
}

void AiAudioInputPort::MaybeLogInputFormatLocked(const HWS_AUDIO_FRAME* frame, int dstSamples)
{
    if (!frame || inputFormatLogged_) return;
    inputFormatLogged_ = true;
    std::ostringstream os;
    os << "AI audio input normalized: src="
       << frame->sampleRate << "Hz "
       << frame->channels << "ch "
       << frame->bitsPerSample << "bit data=" << frame->dataSize
       << " bytes -> 16000Hz mono pcm16, firstConvertedSamples=" << dstSamples
       << ", chunkMs=" << config_.chunkMs;
    ResultCallback cb = resultCb_;
    if (cb) {
        AiDemoResult r;
        r.typeUtf8 = "audio.debug";
        r.ok = true;
        r.statusUtf8 = "ok";
        r.timestamp100ns = frame->timestamp100ns;
        r.summaryUtf8 = os.str();
        cb(r);
    }
}

void AiAudioInputPort::MaybeCheckInputTimestampLocked(const HWS_AUDIO_FRAME* frame)
{
    if (!frame || frame->sampleRate <= 0 || frame->channels <= 0 || frame->bitsPerSample <= 0) return;
    const int bytesPerSample = AiBytesPerSample(frame->bitsPerSample);
    if (bytesPerSample <= 0) return;
    const int srcFrameCount = frame->dataSize / (bytesPerSample * frame->channels);
    if (srcFrameCount <= 0) return;
    const long long duration100ns = ((long long)srcFrameCount * 10000000LL) / (long long)frame->sampleRate;

    if (expectedNextInputTimestamp100ns_ != 0 && frame->timestamp100ns != 0) {
        long long gap = frame->timestamp100ns - expectedNextInputTimestamp100ns_;
        if (gap < 0) gap = -gap;
        if (gap > 1000000LL && timestampGapWarnings_ < 8) { // >100 ms
            ++timestampGapWarnings_;
            std::ostringstream os;
            os << "AI audio timestamp gap warning: expected=" << expectedNextInputTimestamp100ns_
               << " actual=" << frame->timestamp100ns << " gapMs=" << (gap / 10000.0);
            ResultCallback cb = resultCb_;
            if (cb) {
                AiDemoResult r;
                r.typeUtf8 = "audio.debug";
                r.ok = true;
                r.statusUtf8 = "warning";
                r.timestamp100ns = frame->timestamp100ns;
                r.summaryUtf8 = os.str();
                cb(r);
            }
        }
    }
    expectedNextInputTimestamp100ns_ = frame->timestamp100ns + duration100ns;
}

bool AiAudioInputPort::ConvertFrameToMono16k(const HWS_AUDIO_FRAME* frame, std::vector<short>* outSamples)
{
    if (!frame || !frame->data || !outSamples) return false;
    outSamples->clear();
    if (frame->sampleRate <= 0 || frame->channels <= 0 || frame->dataSize <= 0) return false;

    const int srcChannels = frame->channels;
    const int srcRate = frame->sampleRate;
    const int dstRate = config_.targetSampleRate > 0 ? config_.targetSampleRate : 16000;
    const int bytesPerSample = AiBytesPerSample(frame->bitsPerSample);
    if (bytesPerSample <= 0) return false;

    const int srcFrameCount = frame->dataSize / (bytesPerSample * srcChannels);
    if (srcFrameCount <= 0) return false;

    // For multi-channel capture-card audio, the useful signal is normally in the first two channels.
    // Averaging all 4/6/8 channels can reduce speech level if the extra channels are silent.
    int mixChannels = srcChannels;
    if (mixChannels > 2) mixChannels = 2;

    const bool assumeFloat32 = (frame->bitsPerSample == 32 && bytesPerSample == 4)
        ? AiLooksLikeFloat32(frame->data, frame->dataSize, srcChannels)
        : false;

    std::vector<short> mono48;
    mono48.resize((size_t)srcFrameCount);
    for (int i = 0; i < srcFrameCount; ++i) {
        int sum = 0;
        const unsigned char* frameBase = frame->data + (size_t)i * srcChannels * bytesPerSample;
        for (int ch = 0; ch < mixChannels; ++ch) {
            const unsigned char* p = frameBase + (size_t)ch * bytesPerSample;
            sum += (int)AiReadSampleAsS16(p, bytesPerSample, frame->bitsPerSample, assumeFloat32);
        }
        mono48[(size_t)i] = AiClamp16(sum / mixChannels);
    }

    if (srcRate == dstRate) {
        *outSamples = std::move(mono48);
        return true;
    }

    int dstFrameCount = (int)(((long long)srcFrameCount * dstRate + srcRate - 1) / srcRate);
    if (dstFrameCount <= 0) dstFrameCount = 1;
    outSamples->resize((size_t)dstFrameCount);

    // Linear interpolation resampler. This is still lightweight but smoother than nearest-neighbor.
    for (int i = 0; i < dstFrameCount; ++i) {
        double srcPos = ((double)i * (double)srcRate) / (double)dstRate;
        int idx0 = (int)srcPos;
        if (idx0 < 0) idx0 = 0;
        if (idx0 >= srcFrameCount) idx0 = srcFrameCount - 1;
        int idx1 = idx0 + 1;
        if (idx1 >= srcFrameCount) idx1 = srcFrameCount - 1;
        double frac = srcPos - (double)idx0;
        double v = (double)mono48[(size_t)idx0] * (1.0 - frac) + (double)mono48[(size_t)idx1] * frac;
        (*outSamples)[(size_t)i] = AiClamp16((int)(v >= 0.0 ? v + 0.5 : v - 0.5));
    }
    return true;
}

void AiAudioInputPort::QueueReadyChunksLocked()
{
    const int sampleRate = config_.targetSampleRate > 0 ? config_.targetSampleRate : 16000;
    const int chunkMs = config_.chunkMs > 0 ? config_.chunkMs : 200;
    int chunkSamples = (sampleRate * chunkMs) / 1000;
    if (chunkSamples <= 0) chunkSamples = sampleRate / 5;

    while ((int)pendingSamples_.size() >= chunkSamples) {
        AiCapturedAudioChunk chunk;
        chunk.audioSource = config_.audioSource;
        chunk.sampleRate = sampleRate;
        chunk.channels = 1;
        chunk.bitsPerSample = 16;
        chunk.sampleCount = chunkSamples;
        chunk.timestamp100ns = pendingStartTimestamp100ns_;
        chunk.chunkIndex = nextChunkIndex_++;
        chunk.srcLangUtf8 = config_.srcLangUtf8;
        chunk.dstLangUtf8 = config_.dstLangUtf8;
        chunk.pcm16.resize((size_t)chunkSamples * 2);
        memcpy(chunk.pcm16.data(), pendingSamples_.data(), chunk.pcm16.size());

        pendingSamples_.erase(pendingSamples_.begin(), pendingSamples_.begin() + chunkSamples);
        pendingStartTimestamp100ns_ += ((long long)chunkSamples * 10000000LL) / (long long)sampleRate;

        // Audio should not be silently dropped. Keep a much larger queue than older builds
        // because the WebSocket worker may block while faster-whisper processes a segment.
        const size_t maxQueueChunks = 256;
        while (queue_.size() >= maxQueueChunks) {
            queue_.pop_front();
            ++droppedQueueChunks_;
        }
        queue_.push_back(std::move(chunk));
        ++producedChunks_;
        if ((long long)queue_.size() > maxObservedQueueDepth_) maxObservedQueueDepth_ = (long long)queue_.size();
    }
}

void AiAudioInputPort::PushAudioFrame(const HWS_AUDIO_FRAME* frame)
{
    if (!frame || !running_.load()) return;

    std::vector<short> mono16k;
    if (!ConvertFrameToMono16k(frame, &mono16k) || mono16k.empty()) return;

    ResultCallback debugCb;
    std::string debugText;
    std::string debugStatus = "ok";
    long long debugTs = frame->timestamp100ns;

    {
        std::lock_guard<std::mutex> guard(lock_);
        if (!running_.load()) return;
        ++pushedInputFrames_;
        convertedInputSamples_ += (long long)mono16k.size();
        MaybeLogInputFormatLocked(frame, (int)mono16k.size());
        MaybeCheckInputTimestampLocked(frame);

        if (pendingSamples_.empty()) {
            pendingStartTimestamp100ns_ = frame->timestamp100ns;
        }
        pendingSamples_.insert(pendingSamples_.end(), mono16k.begin(), mono16k.end());
        QueueReadyChunksLocked();

        const bool overflowReport = (droppedQueueChunks_ > 0 &&
            (droppedQueueChunks_ == 1 || (droppedQueueChunks_ % 20) == 0));
        const bool periodicReport = (producedChunks_ > 0 &&
            (producedChunks_ - lastProducerReportChunk_ >= 25));
        const bool pressureReport = (queue_.size() >= 10 && producedChunks_ != lastProducerReportChunk_);
        if (overflowReport || periodicReport || pressureReport) {
            lastProducerReportChunk_ = producedChunks_;
            debugCb = resultCb_;
            debugStatus = overflowReport ? "warning" : "ok";
            std::ostringstream os;
            os << "AI audio producer stats: inputFrames=" << pushedInputFrames_
               << " producedChunks=" << producedChunks_
               << " queueDepth=" << queue_.size()
               << " maxQueueDepth=" << maxObservedQueueDepth_
               << " droppedQueueChunks=" << droppedQueueChunks_
               << " pendingSamples=" << pendingSamples_.size()
               << " lastSrc=" << frame->sampleRate << "Hz/" << frame->channels
               << "ch/" << frame->bitsPerSample << "bit"
               << " lastMono16kSamples=" << mono16k.size();
            debugText = os.str();
        }
    }
    cv_.notify_one();

    if (debugCb) {
        AiDemoResult r;
        r.typeUtf8 = "audio.debug";
        r.ok = true;
        r.statusUtf8 = debugStatus;
        r.timestamp100ns = debugTs;
        r.summaryUtf8 = debugText;
        debugCb(r);
    }
}

void AiAudioInputPort::WorkerThread()
{
    HRESULT cohr = CoInitializeEx(NULL, COINIT_MULTITHREADED);
    bool needCoUninit = SUCCEEDED(cohr);
    AiWebSocketClient wsClient;

    for (;;) {
        AiCapturedAudioChunk chunk;
        AiAudioInputConfig cfg;
        ResultCallback cb;
        size_t queueDepthAfterPop = 0;

        {
            std::unique_lock<std::mutex> guard(lock_);
            cv_.wait(guard, [this] { return stopping_.load() || !queue_.empty(); });
            if (stopping_.load() && queue_.empty()) break;
            chunk = std::move(queue_.front());
            queue_.pop_front();
            queueDepthAfterPop = queue_.size();
            cfg = config_;
            cb = resultCb_;
        }

        AiDemoResult result;
        result.typeUtf8 = "audio.error";
        result.ok = false;
        result.statusUtf8 = "error";
        result.timestamp100ns = chunk.timestamp100ns;
        result.startTimestamp100ns = chunk.timestamp100ns;
        result.endTimestamp100ns = chunk.timestamp100ns + ((long long)chunk.sampleCount * 10000000LL) / (long long)chunk.sampleRate;
        result.audioSource = chunk.audioSource;
        result.chunkIndex = chunk.chunkIndex;

        std::wstring errorText;
        if (!wsClient.IsConnected() || _wcsicmp(wsClient.Endpoint().c_str(), cfg.endpoint.c_str()) != 0) {
            if (!wsClient.Connect(cfg.endpoint, &errorText)) {
                result.summaryUtf8 = AiWideToUtf8(errorText.empty() ? L"AI audio WebSocket connect failed." : errorText);
                if (cb) cb(result);
                Sleep(100);
                continue;
            }
        }

        std::string response;
        ULONGLONG sendStartMs = GetTickCount64();
        if (!wsClient.SendAudioChunk(chunk.audioSource,
                                     chunk.timestamp100ns,
                                     chunk.chunkIndex,
                                     chunk.sampleRate,
                                     chunk.channels,
                                     chunk.bitsPerSample,
                                     chunk.sampleCount,
                                     chunk.pcm16,
                                     chunk.srcLangUtf8,
                                     chunk.dstLangUtf8,
                                     &response,
                                     &errorText)) {
            {
                std::lock_guard<std::mutex> guard(lock_);
                ++sendFailures_;
            }
            result.summaryUtf8 = AiWideToUtf8(errorText.empty() ? L"AI audio WebSocket request failed." : errorText);
            if (cb) cb(result);
            wsClient.Close();
            Sleep(50);
            continue;
        }
        ULONGLONG sendEndMs = GetTickCount64();
        long long roundtripMs = (long long)(sendEndMs - sendStartMs);

        {
            std::lock_guard<std::mutex> guard(lock_);
            ++sentChunks_;
            ++receivedResults_;
            if (roundtripMs > maxSendRoundtripMs_) maxSendRoundtripMs_ = roundtripMs;
        }

        if (cb && (roundtripMs > 800 || queueDepthAfterPop >= 10 || (chunk.chunkIndex - lastSenderReportChunk_) >= 25)) {
            lastSenderReportChunk_ = chunk.chunkIndex;
            AiDemoResult dbg;
            dbg.typeUtf8 = "audio.debug";
            dbg.ok = true;
            dbg.statusUtf8 = (roundtripMs > 1500 || queueDepthAfterPop >= 20) ? "warning" : "ok";
            dbg.timestamp100ns = chunk.timestamp100ns;
            dbg.audioSource = chunk.audioSource;
            dbg.chunkIndex = chunk.chunkIndex;
            std::ostringstream os;
            os << "AI audio sender stats: sentChunk=" << chunk.chunkIndex
               << " rttMs=" << roundtripMs
               << " queueDepthAfterPop=" << queueDepthAfterPop
               << " payloadBytes=" << chunk.pcm16.size()
               << " sampleCount=" << chunk.sampleCount
               << " maxRttMs=" << maxSendRoundtripMs_;
            dbg.summaryUtf8 = os.str();
            cb(dbg);
        }

        result.ok = true;
        result.rawJsonUtf8 = response;
        result.typeUtf8 = AiExtractJsonString(response, "type");
        if (result.typeUtf8.empty()) result.typeUtf8 = "audio.transcript_partial";
        result.statusUtf8 = AiExtractJsonString(response, "status");
        if (result.statusUtf8.empty()) result.statusUtf8 = "ok";
        result.textUtf8 = AiExtractJsonString(response, "text");
        result.translationUtf8 = AiExtractJsonString(response, "translation");
        result.isFinal = (response.find("\"isFinal\": true") != std::string::npos) ||
                         (response.find("\"isFinal\":true") != std::string::npos);
        result.srcLangUtf8 = AiExtractJsonString(response, "srcLang");
        result.dstLangUtf8 = AiExtractJsonString(response, "dstLang");
        result.summaryUtf8 = AiExtractJsonString(response, "summary");
        if (result.summaryUtf8.empty()) {
            result.summaryUtf8 = !result.translationUtf8.empty() ? result.translationUtf8 : result.textUtf8;
        }
        if (result.summaryUtf8.empty()) result.summaryUtf8 = "audio stream result";
        if (cb) cb(result);
    }

    wsClient.Close();
    if (needCoUninit) CoUninitialize();
}
