#include "transcriber.h"
#include "whisper.h"
#include <SDL.h>
#include <SDL_audio.h>
#include <algorithm>
#include <cmath>
#include <cstring>

Transcriber::~Transcriber() {
    cancel();
    if (m_worker.joinable()) m_worker.join();
    if (m_ctx) whisper_free(m_ctx);
}

int Transcriber::default_threads() {
    unsigned hc = std::thread::hardware_concurrency();
    if (hc <= 2) return (int)std::max(1u, hc);
    return (int)(hc / 2);
}

float Transcriber::recorded_seconds() const {
    std::lock_guard<std::mutex> lk(m_capture_mtx);
    return (float)(m_capture.size() / (double)WHISPER_SAMPLE_RATE);
}

void Transcriber::s_progress(whisper_context*, whisper_state*, int p, void* ud) {
    auto* self = static_cast<Transcriber*>(ud);
    if (self && self->m_on_progress) self->m_on_progress(p);
}

bool Transcriber::s_abort(void* ud) {
    auto* self = static_cast<Transcriber*>(ud);
    return self && self->m_abort.load();
}

bool Transcriber::preload(const WhisperConfig& cfg) {
    std::lock_guard<std::mutex> lk(m_cfg_mtx);
    m_cfg = cfg;
    if (m_cfg.n_threads <= 0) m_cfg.n_threads = default_threads();
    if (m_ctx) return true;
    whisper_context_params cp = whisper_context_default_params();
    cp.use_gpu = m_cfg.use_gpu;
    m_ctx = whisper_init_from_file_with_params(m_cfg.model_path.c_str(), cp);
    return m_ctx != nullptr;
}

static void sdl_capture_cb(void* user, Uint8* stream, int len) {
    auto* self = static_cast<Transcriber*>(user);
    self->on_audio(reinterpret_cast<float*>(stream), len / (int)sizeof(float));
}

bool Transcriber::start_recording() {
    if (m_recording.load() || m_busy.load()) return false;
    {
        std::lock_guard<std::mutex> lk(m_capture_mtx);
        m_capture.clear();
        m_capture.reserve(WHISPER_SAMPLE_RATE * 30);
    }

    SDL_AudioSpec want{}, have{};
    want.freq     = WHISPER_SAMPLE_RATE;
    want.format   = AUDIO_F32;
    want.channels = 1;
    want.samples  = 1024;
    want.callback = sdl_capture_cb;
    want.userdata = this;

    const char* dev = SDL_GetAudioDeviceName(m_cfg.capture_id, SDL_TRUE);
    m_dev = SDL_OpenAudioDevice(dev, SDL_TRUE, &want, &have, 0);
    if (!m_dev) return false;

    m_energy    = 0.0f;
    m_recording = true;
    SDL_PauseAudioDevice(m_dev, 0);
    return true;
}

void Transcriber::on_audio(const float* s, int n) {
    if (n <= 0 || !m_recording.load()) return;
    double sq = 0.0;
    for (int i = 0; i < n; ++i) sq += (double)s[i] * s[i];
    float rms = (float)std::sqrt(sq / n);
    float e   = m_energy.load();
    m_energy  = std::min(1.0f, e * 0.6f + (rms * 4.0f) * 0.4f);
    std::lock_guard<std::mutex> lk(m_capture_mtx);
    m_capture.insert(m_capture.end(), s, s + n);
}

void Transcriber::cancel() {
    if (!m_recording.load()) return;
    m_recording = false;
    if (m_dev) { SDL_PauseAudioDevice(m_dev, 1); SDL_CloseAudioDevice(m_dev); m_dev = 0; }
    std::lock_guard<std::mutex> lk(m_capture_mtx);
    m_capture.clear();
    m_energy = 0.0f;
}

void Transcriber::stop_and_transcribe() {
    if (!m_recording.load()) return;
    m_recording = false;
    if (m_dev) { SDL_PauseAudioDevice(m_dev, 1); SDL_CloseAudioDevice(m_dev); m_dev = 0; }
    m_energy = 0.0f;

    std::vector<float> audio;
    { std::lock_guard<std::mutex> lk(m_capture_mtx); audio.swap(m_capture); }

    if (audio.size() < (size_t)(WHISPER_SAMPLE_RATE * 0.3)) {
        if (m_on_result) m_on_result("");
        return;
    }
    if (m_worker.joinable()) m_worker.join();
    m_busy   = true;
    m_worker = std::thread(&Transcriber::transcribe_worker, this, std::move(audio));
}

static void trim_silence(std::vector<float>& a, float thresh = 0.01f) {
    const size_t win = 1600;
    auto loud = [&](size_t i) {
        float m = 0.f;
        for (size_t k = i; k < std::min(a.size(), i + win); ++k)
            m = std::max(m, std::fabs(a[k]));
        return m > thresh;
    };
    size_t s = 0, e = a.size();
    while (s + win < a.size() && !loud(s)) s += win;
    while (e > win && !loud(e - win)) e -= win;
    if (s + win <= e)
        a.assign(a.begin() + (s > win ? s - win : 0), a.begin() + e);
}

static std::string clean_text(std::string s) {
    const char* junk[] = {"[BLANK_AUDIO]", "[NOISE]", "(blank)", "(noise)", "[ Silence ]"};
    for (auto j : junk) {
        size_t p;
        while ((p = s.find(j)) != std::string::npos) s.erase(p, strlen(j));
    }
    size_t b = s.find_first_not_of(" \t\r\n");
    size_t e = s.find_last_not_of(" \t\r\n");
    return (b == std::string::npos) ? "" : s.substr(b, e - b + 1);
}

std::string Transcriber::run_inference(std::vector<float>& audio) {
    if (!m_ctx) return "";
    m_abort = false;
    if (m_cfg.trim_silence) trim_silence(audio);
    m_audio_seconds = (float)(audio.size() / (double)WHISPER_SAMPLE_RATE);
    whisper_full_params wp = whisper_full_default_params(WHISPER_SAMPLING_GREEDY);
    wp.print_progress   = false;
    wp.print_realtime   = false;
    wp.print_timestamps = false;
    wp.no_timestamps    = true;
    wp.translate        = false;
    wp.language         = m_cfg.language.c_str();
    wp.n_threads        = m_cfg.n_threads;
    wp.no_context       = true;
    wp.suppress_blank   = true;
    wp.suppress_nst     = true;
    wp.temperature      = 0.0f;
    wp.progress_callback           = &Transcriber::s_progress;
    wp.progress_callback_user_data = this;
    wp.abort_callback              = &Transcriber::s_abort;
    wp.abort_callback_user_data    = this;
    std::string out;
    if (whisper_full(m_ctx, wp, audio.data(), (int)audio.size()) == 0) {
        int n = whisper_full_n_segments(m_ctx);
        for (int i = 0; i < n; ++i) {
            const char* t = whisper_full_get_segment_text(m_ctx, i);
            if (t) out += t;
        }
        out = clean_text(out);
    }
    return out;
}

void Transcriber::transcribe_worker(std::vector<float> audio) {
    std::string out = run_inference(audio);
    m_busy = false;
    if (m_on_result) m_on_result(out);
}

std::string Transcriber::transcribe_sync(std::vector<float> audio) {
    return run_inference(audio);
}

bool Transcriber::reload(const WhisperConfig& cfg) {
    if (m_recording.load() || m_busy.load()) return false;
    if (m_worker.joinable()) m_worker.join();
    {
        std::lock_guard<std::mutex> lk(m_cfg_mtx);
        if (m_ctx) { whisper_free(m_ctx); m_ctx = nullptr; }
    }
    return preload(cfg);
}

std::vector<std::string> Transcriber::get_audio_devices() {
    std::vector<std::string> devices;
    if (SDL_Init(SDL_INIT_AUDIO) < 0) return devices;
    int nDevices = SDL_GetNumAudioDevices(SDL_TRUE);
    for (int i = 0; i < nDevices; ++i) {
        const char* name = SDL_GetAudioDeviceName(i, SDL_TRUE);
        if (name) devices.push_back(name);
    }
    return devices;
}
