#include "audio.h"

#include <cstdint>
#include <cstring>
#include <cstdio>
#include <fstream>

#include <samplerate.h>
#include <pulse/pulseaudio.h>

namespace lc {

namespace {

// 单声道流式重采样器（libsamplerate）
class Resampler {
public:
    ~Resampler() { if (st_) src_delete(st_); }

    bool init(int fromRate, int toRate)
    {
        int err = 0;
        st_ = src_new(SRC_SINC_FASTEST, 1, &err);
        ratio_ = (double) toRate / (double) fromRate;
        return st_ != nullptr;
    }

    // 输入 frames 个单声道样本，追加到 out
    bool process(const float * in, size_t frames, std::vector<float> & out)
    {
        if (!st_) return false;
        if (frames == 0) return true;

        // 输出上限：比例 + 余量
        const size_t maxOut = (size_t) (frames * ratio_) + 64;
        const size_t base   = out.size();
        out.resize(base + maxOut);

        SRC_DATA d{};
        d.data_in       = in;
        d.input_frames  = (long) frames;
        d.data_out      = out.data() + base;
        d.output_frames = (long) maxOut;
        d.src_ratio     = ratio_;
        d.end_of_input  = 0;

        if (src_process(st_, &d) != 0) return false;
        out.resize(base + (size_t) d.output_frames_gen);
        return true;
    }

private:
    SRC_STATE * st_    = nullptr;
    double      ratio_ = 1.0;
};

// 读小端整数
uint32_t rd32(const uint8_t * p) { return p[0] | (p[1] << 8) | (p[2] << 16) | ((uint32_t) p[3] << 24); }
uint16_t rd16(const uint8_t * p) { return (uint16_t) (p[0] | (p[1] << 8)); }

} // namespace

// ---------------------------------------------------------------------------
// WAV 读取
// ---------------------------------------------------------------------------
bool readWavMono(const std::string & path, int targetRate,
                 std::vector<float> & out, std::string * error)
{
    const auto fail = [error](const std::string & m) {
        if (error) *error = m;
        return false;
    };

    std::ifstream f(path, std::ios::binary);
    if (!f) return fail("无法打开文件: " + path);

    std::vector<uint8_t> buf((std::istreambuf_iterator<char>(f)), std::istreambuf_iterator<char>());
    if (buf.size() < 44 || std::memcmp(buf.data(), "RIFF", 4) != 0 || std::memcmp(buf.data() + 8, "WAVE", 4) != 0) {
        return fail("不是合法的 WAV 文件: " + path);
    }

    uint16_t audioFormat = 0, channels = 0, bitsPerSample = 0;
    uint32_t sampleRate = 0;
    const uint8_t * dataPtr = nullptr;
    size_t dataLen = 0;

    size_t pos = 12;
    while (pos + 8 <= buf.size()) {
        const uint8_t * chunk = buf.data() + pos;
        const uint32_t chunkSize = rd32(chunk + 4);
        const uint8_t * body = chunk + 8;
        if (std::memcmp(chunk, "fmt ", 4) == 0) {
            if (chunkSize < 16) return fail("fmt 块过短");
            audioFormat   = rd16(body + 0);
            channels      = rd16(body + 2);
            sampleRate    = rd32(body + 4);
            bitsPerSample = rd16(body + 14);
            if (audioFormat == 0xFFFE && chunkSize >= 40) {
                audioFormat = rd16(body + 24);   // WAVE_FORMAT_EXTENSIBLE 的 subformat 前两字节
            }
        } else if (std::memcmp(chunk, "data", 4) == 0) {
            dataPtr = body;
            dataLen = chunkSize;
            if (pos + 8 + chunkSize > buf.size()) dataLen = buf.size() - (pos + 8);
        }
        pos += 8 + chunkSize + (chunkSize & 1);   // 块按偶数字节对齐
    }

    if (!dataPtr || channels == 0 || sampleRate == 0) return fail("WAV 缺少 fmt 或 data 块");

    const bool isFloat = (audioFormat == 3);
    const int bytesPerSample = bitsPerSample / 8;
    if (bytesPerSample <= 0) return fail("不支持的位深");
    const size_t frameBytes = (size_t) bytesPerSample * channels;
    const size_t frames = dataLen / frameBytes;

    // 降混为单声道 float
    std::vector<float> mono(frames);
    for (size_t i = 0; i < frames; ++i) {
        double acc = 0.0;
        const uint8_t * p = dataPtr + i * frameBytes;
        for (uint16_t c = 0; c < channels; ++c, p += bytesPerSample) {
            if (isFloat) {
                float v;
                std::memcpy(&v, p, 4);
                acc += v;
            } else if (bitsPerSample == 16) {
                acc += (int16_t) rd16(p) / 32768.0;
            } else if (bitsPerSample == 24) {
                int32_t v = p[0] | (p[1] << 8) | (p[2] << 16);
                if (v & 0x800000) v |= ~0xFFFFFF;
                acc += v / 8388608.0;
            } else if (bitsPerSample == 32) {
                acc += (int32_t) rd32(p) / 2147483648.0;
            }
        }
        mono[i] = (float) (acc / channels);
    }

    if ((int) sampleRate == targetRate) {
        out = std::move(mono);
        return true;
    }

    Resampler rs;
    if (!rs.init((int) sampleRate, targetRate)) return fail("重采样器初始化失败");
    out.clear();
    return rs.process(mono.data(), mono.size(), out);
}

// ---------------------------------------------------------------------------
// 实时采集：默认 sink 的 monitor
// ---------------------------------------------------------------------------
struct MonitorCapture::Impl {
    pa_threaded_mainloop * ml      = nullptr;
    pa_context *           ctx     = nullptr;
    pa_stream *            stream  = nullptr;
    ChunkCallback          callback;
    int                    targetRate = 16000;
    int                    deviceRate = 48000;
    int                    channels   = 2;
    Resampler              resampler;
    std::string            deviceName;
    bool                   ready = false;
};

namespace {

void storeMonitor(pa_context * /*c*/, const pa_server_info * info, void * userdata)
{
    auto * d = static_cast<MonitorCapture::Impl *>(userdata);
    if (!info || !info->default_sink_name) return;
    d->deviceName = std::string(info->default_sink_name) + ".monitor";
    pa_threaded_mainloop_signal(d->ml, 0);
}

void streamState(pa_stream * s, void * userdata)
{
    auto * d = static_cast<MonitorCapture::Impl *>(userdata);
    switch (pa_stream_get_state(s)) {
    case PA_STREAM_READY:
    case PA_STREAM_FAILED:
    case PA_STREAM_TERMINATED:
        pa_threaded_mainloop_signal(d->ml, 0);
        break;
    default: break;
    }
}

void streamRead(pa_stream * s, size_t /*nbytes*/, void * userdata)
{
    auto * d = static_cast<MonitorCapture::Impl *>(userdata);
    while (pa_stream_readable_size(s) > 0) {
        const void * data = nullptr;
        size_t bytes = 0;
        if (pa_stream_peek(s, &data, &bytes) < 0 || !data || bytes == 0) {
            pa_stream_drop(s);
            continue;
        }

        const float * interleaved = static_cast<const float *>(data);
        const size_t frames = bytes / (sizeof(float) * (size_t) d->channels);

        // 降混单声道
        std::vector<float> mono(frames);
        for (size_t i = 0; i < frames; ++i) {
            double acc = 0.0;
            for (int c = 0; c < d->channels; ++c) acc += interleaved[i * d->channels + c];
            mono[i] = (float) (acc / d->channels);
        }

        // 重采样到目标采样率
        std::vector<float> resampled;
        if (d->deviceRate == d->targetRate) {
            resampled = std::move(mono);
        } else if (!d->resampler.process(mono.data(), mono.size(), resampled)) {
            pa_stream_drop(s);
            continue;
        }

        if (d->callback && !resampled.empty()) {
            d->callback(resampled.data(), resampled.size());
        }
        pa_stream_drop(s);
    }
}

void contextState(pa_context * c, void * userdata)
{
    auto * d = static_cast<MonitorCapture::Impl *>(userdata);
    switch (pa_context_get_state(c)) {
    case PA_CONTEXT_READY:   d->ready = true;  pa_threaded_mainloop_signal(d->ml, 0); break;
    case PA_CONTEXT_FAILED:
    case PA_CONTEXT_TERMINATED:
        pa_threaded_mainloop_signal(d->ml, 0);
        break;
    default: break;
    }
}

} // namespace

MonitorCapture::MonitorCapture() : d(new Impl) {}

MonitorCapture::~MonitorCapture()
{
    stop();
    delete d;
}

std::string MonitorCapture::deviceName() const { return d->deviceName; }

bool MonitorCapture::start(int targetRate, ChunkCallback callback, std::string * error)
{
    const auto fail = [error](const std::string & m) {
        if (error) *error = m;
        return false;
    };

    d->targetRate = targetRate;
    d->callback   = std::move(callback);

    d->ml = pa_threaded_mainloop_new();
    if (!d->ml) return fail("pa_threaded_mainloop_new 失败");

    d->ctx = pa_context_new(pa_threaded_mainloop_get_api(d->ml), "live-caption");
    if (!d->ctx) return fail("pa_context_new 失败");

    pa_context_set_state_callback(d->ctx, contextState, d);

    if (pa_threaded_mainloop_start(d->ml) < 0) return fail("pa_threaded_mainloop_start 失败");

    pa_threaded_mainloop_lock(d->ml);

    if (pa_context_connect(d->ctx, nullptr, PA_CONTEXT_NOFLAGS, nullptr) < 0) {
        pa_threaded_mainloop_unlock(d->ml);
        return fail("pa_context_connect 失败（PulseAudio/PipeWire 未运行？）");
    }

    // 等 context READY
    for (;;) {
        const pa_context_state_t st = pa_context_get_state(d->ctx);
        if (st == PA_CONTEXT_READY) break;
        if (!PA_CONTEXT_IS_GOOD(st)) {
            pa_threaded_mainloop_unlock(d->ml);
            return fail("无法连接到 PulseAudio/PipeWire 服务");
        }
        pa_threaded_mainloop_wait(d->ml);
    }

    // 查默认 sink，拼出它的 monitor 源
    d->deviceName.clear();
    pa_operation * op = pa_context_get_server_info(d->ctx, storeMonitor, d);
    while (pa_operation_get_state(op) == PA_OPERATION_RUNNING) {
        pa_threaded_mainloop_wait(d->ml);
    }
    pa_operation_unref(op);

    if (d->deviceName.empty()) {
        pa_threaded_mainloop_unlock(d->ml);
        return fail("拿不到默认输出设备名");
    }

    // 以设备实际速率请求（默认 48k）
    pa_sample_spec ss{};
    ss.format   = PA_SAMPLE_FLOAT32LE;
    ss.rate     = 48000;
    ss.channels = 2;

    d->stream = pa_stream_new(d->ctx, "live-caption-monitor", &ss, nullptr);
    if (!d->stream) {
        pa_threaded_mainloop_unlock(d->ml);
        return fail("pa_stream_new 失败");
    }

    pa_stream_set_read_callback(d->stream, streamRead, d);
    pa_stream_set_state_callback(d->stream, streamState, d);

    pa_buffer_attr attr{};
    attr.maxlength = (uint32_t) -1;
    attr.tlength   = (uint32_t) -1;
    attr.prebuf    = (uint32_t) -1;
    attr.minreq    = (uint32_t) -1;
    attr.fragsize  = (uint32_t) (pa_frame_size(&ss) * 1024);   // ~21ms

    if (pa_stream_connect_record(d->stream, d->deviceName.c_str(), &attr, PA_STREAM_ADJUST_LATENCY) < 0) {
        pa_threaded_mainloop_unlock(d->ml);
        return fail("连接 monitor 源失败: " + d->deviceName);
    }

    // 等 stream READY，才能确认 monitor 源确实可用
    for (;;) {
        const pa_stream_state_t st = pa_stream_get_state(d->stream);
        if (st == PA_STREAM_READY) break;
        if (st == PA_STREAM_FAILED || st == PA_STREAM_TERMINATED) {
            pa_threaded_mainloop_unlock(d->ml);
            return fail("monitor 源不可用: " + d->deviceName + " (" +
                        std::string(pa_strerror(pa_context_errno(d->ctx))) + ")");
        }
        pa_threaded_mainloop_wait(d->ml);
    }

    // 更新为实际生效的格式
    const pa_sample_spec * actual = pa_stream_get_sample_spec(d->stream);
    if (actual) {
        d->deviceRate = (int) actual->rate;
        d->channels   = (int) actual->channels;
    }
    if (!d->resampler.init(d->deviceRate, d->targetRate)) {
        pa_threaded_mainloop_unlock(d->ml);
        return fail("重采样器初始化失败");
    }

    pa_threaded_mainloop_unlock(d->ml);
    return true;
}

void MonitorCapture::stop()
{
    if (!d->ml) return;

    pa_threaded_mainloop_lock(d->ml);
    if (d->stream) {
        pa_stream_disconnect(d->stream);
        pa_stream_unref(d->stream);
        d->stream = nullptr;
    }
    if (d->ctx) {
        pa_context_disconnect(d->ctx);
        pa_context_unref(d->ctx);
        d->ctx = nullptr;
    }
    pa_threaded_mainloop_unlock(d->ml);

    pa_threaded_mainloop_stop(d->ml);
    pa_threaded_mainloop_free(d->ml);
    d->ml = nullptr;
    d->callback = nullptr;
}

} // namespace lc