#include "lemon/realtime_session.h"
#include "lemon/router.h"
#include <algorithm>
#include <random>
#include <chrono>
#include <iostream>
#include <cmath>
#include <thread>
#include <lemon/utils/aixlog.hpp>

#ifdef _WIN32
#include <windows.h>
#else
#include <unistd.h>
#endif

namespace lemon {

RealtimeSessionManager::RealtimeSessionManager(Router* router)
    : router_(router) {
}

RealtimeSessionManager::~RealtimeSessionManager() {
    // Wait for pending transcriptions to complete
    {
        std::lock_guard<std::mutex> lock(transcriptions_mutex_);
        for (auto& f : pending_transcriptions_) {
            if (f.valid()) f.wait();
        }
        pending_transcriptions_.clear();
    }

    // Close all sessions
    std::lock_guard<std::mutex> lock(sessions_mutex_);
    sessions_.clear();
}

std::string RealtimeSessionManager::generate_session_id() {
    // Generate a random session ID
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, 15);

    const char* hex_chars = "0123456789abcdef";
    std::string id = "sess_";

    for (int i = 0; i < 24; i++) {
        id += hex_chars[dis(gen)];
    }

    return id;
}

void RealtimeSessionManager::apply_turn_detection_config(
    std::shared_ptr<RealtimeSession> session,
    const json& turn_detection) {
    if (turn_detection.is_null()) {
        session->turn_detection_enabled = false;
        session->turn_detection_config = nullptr;
        session->vad.reset();
        session->vad_speech_window_open = false;
        session->last_interim_transcription_ms = 0;
        session->interim_generation.fetch_add(1, std::memory_order_relaxed);
        return;
    }

    SimpleVAD::Config vad_config;
    if (session->turn_detection_config.is_object()) {
        vad_config.energy_threshold = session->turn_detection_config.value(
            "threshold",
            vad_config.energy_threshold
        );
        vad_config.min_silence_ms = session->turn_detection_config.value(
            "silence_duration_ms",
            vad_config.min_silence_ms
        );
        vad_config.min_speech_ms = session->turn_detection_config.value(
            "prefix_padding_ms",
            vad_config.min_speech_ms
        );
    }

    if (turn_detection.is_object()) {
        if (turn_detection.contains("threshold")) {
            vad_config.energy_threshold = turn_detection["threshold"].get<float>();
        }
        if (turn_detection.contains("silence_duration_ms")) {
            vad_config.min_silence_ms = turn_detection["silence_duration_ms"].get<int>();
        }
        if (turn_detection.contains("prefix_padding_ms")) {
            vad_config.min_speech_ms = turn_detection["prefix_padding_ms"].get<int>();
        }
    }

    session->vad.set_config(vad_config);
    session->turn_detection_config = {
        {"threshold", vad_config.energy_threshold},
        {"silence_duration_ms", vad_config.min_silence_ms},
        {"prefix_padding_ms", vad_config.min_speech_ms}
    };
    session->turn_detection_enabled = true;
}

std::string RealtimeSessionManager::create_session(
    std::function<void(const json&)> send_callback,
    const json& config
) {
    std::string session_id = generate_session_id();

    auto session = std::make_shared<RealtimeSession>(session_id);
    session->send_message = std::move(send_callback);

    // Apply initial configuration
    if (config.contains("model")) {
        session->model = config["model"].get<std::string>();
    }

    // Configure VAD if specified
    if (config.contains("turn_detection")) {
        apply_turn_detection_config(session, config["turn_detection"]);
    }

    // Attempt to connect to a streaming backend if one is available
    connect_streaming_backend(session);

    {
        std::lock_guard<std::mutex> lock(sessions_mutex_);
        sessions_[session_id] = std::move(session);
    }

    // Send session created message (OpenAI-compatible)
    json created_msg = {
        {"type", "session.created"},
        {"session", {
            {"id", session_id}
        }}
    };

    // Get the session and send
    auto sess = get_session(session_id);
    if (sess && sess->send_message) {
        sess->send_message(created_msg);
    }

    return session_id;
}

void RealtimeSessionManager::update_session(const std::string& session_id, const json& config) {
    auto session = get_session(session_id);
    if (!session) {
        return;
    }

    if (config.contains("model")) {
        session->model = config["model"].get<std::string>();
    }

    if (config.contains("turn_detection")) {
        apply_turn_detection_config(session, config["turn_detection"]);
    }

    // Reconnect streaming backend if model changed
    connect_streaming_backend(session);

    // Send session updated message (OpenAI-compatible)
    if (session->send_message) {
        json updated_msg = {
            {"type", "session.updated"},
            {"session", {
                {"id", session_id},
                {"model", session->model},
                {"turn_detection", session->turn_detection_config}
            }}
        };
        session->send_message(updated_msg);
    }
}

void RealtimeSessionManager::append_audio(const std::string& session_id, const std::string& base64_audio) {
    auto session = get_session(session_id);
    if (!session || !session->session_active) {
        return;
    }

    // If connected to a streaming backend, forward audio directly
    if (session->use_streaming_backend.load()) {
        forward_streaming_audio(session, base64_audio);
        return;
    }

    // Append to buffer
    session->audio_buffer.append(base64_audio);

    // Log buffer growth periodically (every ~5 seconds at 256ms chunks ≈ every 20 chunks)
    static int chunk_count = 0;
    if (++chunk_count % 20 == 1) {
        LOG(DEBUG, "RealtimeSession") << "Audio buffer: " << session->audio_buffer.duration_ms()
                  << "ms (" << session->audio_buffer.sample_count() << " samples)" << std::endl;
    }

    // In manual-commit mode, buffer audio only and skip server-side VAD/interim work.
    if (session->turn_detection_enabled.load()) {
        process_vad(session);
    }
}

void RealtimeSessionManager::process_vad(std::shared_ptr<RealtimeSession> session) {
    // Get recent audio for VAD processing (last 100ms)
    auto recent_audio = session->audio_buffer.get_recent_samples(100);
    if (recent_audio.empty()) {
        return;
    }

    // Log RMS periodically for threshold tuning
    {
        float sum_sq = 0.0f;
        for (float s : recent_audio) sum_sq += s * s;
        float rms = std::sqrt(sum_sq / static_cast<float>(recent_audio.size()));
        static int vad_log_count = 0;
        if (++vad_log_count % 20 == 1) {
            LOG(DEBUG, "RealtimeSession") << "VAD: RMS=" << rms
                      << " speech_active=" << session->vad.is_speech_active() << std::endl;
        }
    }

    SimpleVAD::Event event = session->vad.process(recent_audio, StreamingAudioBuffer::SAMPLE_RATE);

    switch (event) {
        case SimpleVAD::Event::SpeechStart: {
            LOG(DEBUG, "RealtimeSession") << "VAD: SpeechStart detected" << std::endl;
            session->audio_start_ms = session->vad.speech_start_ms();
            session->vad_speech_window_open = true;
            session->last_interim_transcription_ms = 0;  // Reset interim tracking for new utterance

            if (session->send_message) {
                json msg = {
                    {"type", "input_audio_buffer.speech_started"},
                    {"audio_start_ms", session->audio_start_ms}
                };
                session->send_message(msg);
            }
            break;
        }

        case SimpleVAD::Event::SpeechEnd: {
            LOG(DEBUG, "RealtimeSession") << "VAD: SpeechEnd detected, triggering transcription" << std::endl;
            int64_t audio_end_ms = session->vad.speech_end_ms();

            if (session->send_message) {
                json msg = {
                    {"type", "input_audio_buffer.speech_stopped"},
                    {"audio_end_ms", audio_end_ms}
                };
                session->send_message(msg);
            }

            // Trigger final transcription (clears buffer)
            session->vad_speech_window_open = false;
            transcribe_and_send(session);
            break;
        }

        case SimpleVAD::Event::None:
        default:
            // Speech is ongoing — check if we should fire an interim transcription
            if (session->vad.is_speech_active()) {
                maybe_interim_transcribe(session);
            }
            break;
    }
}

void RealtimeSessionManager::maybe_interim_transcribe(std::shared_ptr<RealtimeSession> session) {
    if (!session || session->audio_buffer.empty()) return;

    int buffer_ms = session->audio_buffer.duration_ms();

    // Determine how much new audio has arrived since the last interim transcription.
    // On the first interim of an utterance last_interim_transcription_ms is 0, so we
    // compare against the raw buffer duration.
    int since_last = (session->last_interim_transcription_ms == 0)
        ? buffer_ms
        : buffer_ms - session->last_interim_transcription_ms;

    if (since_last >= INTERIM_TRANSCRIPTION_CHUNK_MS) {
        transcribe_interim(session);
    }
}

void RealtimeSessionManager::transcribe_interim(std::shared_ptr<RealtimeSession> session) {
    if (!session || session->audio_buffer.empty()) return;

    // Avoid overlapping interim transcriptions for the same session
    bool expected = false;
    if (!session->interim_in_flight.compare_exchange_strong(expected, true)) {
        return;  // Another interim is already running
    }

    // Snapshot the buffer WITHOUT clearing it
    auto wav_data = session->audio_buffer.get_wav_padded(500);
    std::string model = session->model;
    session->last_interim_transcription_ms = session->audio_buffer.duration_ms();
    const std::uint64_t generation =
        session->interim_generation.load(std::memory_order_relaxed);

    LOG(DEBUG, "RealtimeSession") << "Firing interim transcription at "
              << session->last_interim_transcription_ms << "ms" << std::endl;

    auto future = std::async(std::launch::async,
        [this, session, wav_data = std::move(wav_data), model = std::move(model), generation]() {
            transcribe_wav(
                session,
                wav_data,
                model,
                /*is_interim=*/true,
                generation);
            session->interim_in_flight.store(false);
        });

    {
        std::lock_guard<std::mutex> lock(transcriptions_mutex_);
        pending_transcriptions_.erase(
            std::remove_if(pending_transcriptions_.begin(), pending_transcriptions_.end(),
                [](const std::future<void>& f) {
                    return f.wait_for(std::chrono::seconds(0)) == std::future_status::ready;
                }),
            pending_transcriptions_.end()
        );
        pending_transcriptions_.push_back(std::move(future));
    }
}

void RealtimeSessionManager::commit_audio(const std::string& session_id) {
    auto session = get_session(session_id);
    if (!session) {
        return;
    }

    // If streaming backend is active, forward commit and skip buffering
    if (session->use_streaming_backend.load()) {
        forward_streaming_commit(session);
        return;
    }

    if (session->turn_detection_enabled.load() && !session->vad_speech_window_open.load()) {
        if (!session->audio_buffer.empty()) {
            LOG(DEBUG, "RealtimeSession")
                << "Ignoring commit with no active VAD speech window; clearing buffered audio"
                << std::endl;
        }
        session->audio_buffer.clear();
        session->vad.reset();
        session->last_interim_transcription_ms = 0;
        session->interim_generation.fetch_add(1, std::memory_order_relaxed);

        if (session->send_message) {
            json msg = {
                {"type", "input_audio_buffer.cleared"}
            };
            session->send_message(msg);
        }
        return;
    }

    if (session->audio_buffer.empty()) {
        return;
    }

    // Send committed message
    if (session->send_message) {
        json msg = {
            {"type", "input_audio_buffer.committed"}
        };
        session->send_message(msg);
    }

    // Trigger transcription
    session->vad_speech_window_open = false;
    transcribe_and_send(session);
}

void RealtimeSessionManager::clear_audio(const std::string& session_id) {
    auto session = get_session(session_id);
    if (!session) {
        return;
    }

    // Forward clear to the streaming backend, keeping the stream open so
    // subsequent appends keep streaming. The backend replies with
    // input_audio_buffer.cleared, which is forwarded to the client.
    if (session->use_streaming_backend.load()) {
        forward_streaming_clear(session);
        return;
    }

    session->audio_buffer.clear();
    session->vad.reset();
    session->vad_speech_window_open = false;
    session->last_interim_transcription_ms = 0;
    session->interim_generation.fetch_add(1, std::memory_order_relaxed);

    if (session->send_message) {
        json msg = {
            {"type", "input_audio_buffer.cleared"}
        };
        session->send_message(msg);
    }
}

void RealtimeSessionManager::transcribe_and_send(std::shared_ptr<RealtimeSession> session) {
    if (!session || session->audio_buffer.empty()) {
        return;
    }

    // Snapshot WAV data and clear buffer on the callback thread (no data race)
    auto wav_data = session->audio_buffer.get_wav_padded(500);
    std::string model = session->model;
    session->audio_buffer.clear();
    session->vad.reset();
    session->vad_speech_window_open = false;
    session->last_interim_transcription_ms = 0;  // Reset for next utterance
    // A final result supersedes any interim jobs that have not started yet.
    session->interim_generation.fetch_add(1, std::memory_order_relaxed);

    // Dispatch transcription to worker thread so it doesn't block the WebSocket callback
    auto future = std::async(std::launch::async,
        [this, session, wav_data = std::move(wav_data), model = std::move(model)]() {
            transcribe_wav(session, wav_data, model);
        });

    // Track future for clean shutdown
    {
        std::lock_guard<std::mutex> lock(transcriptions_mutex_);
        // Remove completed futures
        pending_transcriptions_.erase(
            std::remove_if(pending_transcriptions_.begin(), pending_transcriptions_.end(),
                [](const std::future<void>& f) {
                    return f.wait_for(std::chrono::seconds(0)) == std::future_status::ready;
                }),
            pending_transcriptions_.end()
        );
        pending_transcriptions_.push_back(std::move(future));
    }
}

void RealtimeSessionManager::transcribe_wav(
    std::shared_ptr<RealtimeSession> session,
    std::vector<uint8_t> wav_data, std::string model,
    bool is_interim,
    std::uint64_t interim_generation) {
    try {
        const char* tag = is_interim ? "interim" : "final";
        std::unique_lock<std::mutex> transcription_lock(session->transcription_mutex);

        if (!session->session_active.load()) {
            LOG(DEBUG, "RealtimeSession") << "Skipping " << tag
                      << " transcription for closed session" << std::endl;
            return;
        }

        if (is_interim &&
            session->interim_generation.load(std::memory_order_relaxed) != interim_generation) {
            LOG(DEBUG, "RealtimeSession") << "Skipping stale interim transcription"
                      << std::endl;
            return;
        }

        // Convert WAV bytes to a string for the router (expects file_data as string)
        std::string file_data(reinterpret_cast<const char*>(wav_data.data()), wav_data.size());

        // Build transcription request
        json request = {
            {"model", model},
            {"file_data", file_data},
            {"filename", "realtime_audio.wav"}
        };

        // Call router for transcription
        LOG(DEBUG, "RealtimeSession") << "Calling Whisper " << tag << " transcription ("
                  << wav_data.size() << " bytes)..." << std::endl;
        json response = router_->audio_transcriptions(request);

        // Drop stale result (including stale errors) instead of publishing it.
        if (is_interim &&
            session->interim_generation.load(std::memory_order_relaxed) != interim_generation) {
            LOG(DEBUG, "RealtimeSession") << "Discarding stale interim transcription result"
                      << std::endl;
            return;
        }

        LOG(DEBUG, "RealtimeSession") << "Whisper " << tag << " response: " << response.dump() << std::endl;

        if (response.contains("error")) {
            const json& error = response["error"];
            std::string message = "Whisper backend returned an error";
            if (error.is_object()) {
                message = error.value("message", message);
            } else if (error.is_string()) {
                message = error.get<std::string>();
            }

            const std::string prefix = "Transcription failed: ";
            if (message.rfind(prefix, 0) == 0) {
                message.erase(0, prefix.size());
            }
            throw std::runtime_error(message);
        }

        if (!response.contains("text") || !response["text"].is_string()) {
            throw std::runtime_error(
                "Whisper backend response did not contain a string 'text' field");
        }

        // Send transcription result if session is still active
        if (session->send_message && session->session_active.load()) {
            std::string transcript = response["text"].get<std::string>();

            LOG(DEBUG, "RealtimeSession") << "Sending " << tag << " transcript to client: \""
                      << transcript << "\"" << std::endl;

            if (is_interim) {
                // Interim/partial result — client should treat as replaceable
                json msg = {
                    {"type", "conversation.item.input_audio_transcription.delta"},
                    {"delta", transcript}
                };
                session->send_message(msg);
            } else {
                // Final result — speech segment is complete
                json msg = {
                    {"type", "conversation.item.input_audio_transcription.completed"},
                    {"transcript", transcript}
                };
                session->send_message(msg);
            }
        }

    } catch (const std::exception& e) {
        LOG(ERROR, "RealtimeSession") << "Transcription error: " << e.what() << std::endl;

        if (session->send_message && session->session_active.load()) {
            json error_msg = {
                {"type", "error"},
                {"error", {
                    {"message", std::string("Transcription failed: ") + e.what()},
                    {"type", "transcription_error"}
                }}
            };
            session->send_message(error_msg);
        }
    }
}

void RealtimeSessionManager::connect_streaming_backend(std::shared_ptr<RealtimeSession> session) {
    if (!session || !router_) {
        return;
    }

    std::string address = router_->get_streaming_transcription_address(session->model);
    if (address.empty()) {
        // Backend does not support streaming transcription
        disconnect_streaming_backend(session);
        return;
    }

    {
        // If already connected to the same address, nothing to do
        std::lock_guard<std::mutex> lock(session->streaming_mutex);
        if (session->use_streaming_backend.load() && session->streaming_client &&
            session->streaming_client->is_connected()) {
            return;
        }
    }

    // Disconnect any existing streaming connection
    disconnect_streaming_backend(session);

    // Capture weak_ptr: the session owns the client, and the client's read
    // thread owns this callback — a shared_ptr capture would be a cycle.
    std::weak_ptr<RealtimeSession> weak_session = session;
    auto client = std::make_unique<utils::TcpJsonlClient>();
    auto callback = [weak_session](const json& msg) {
        auto session = weak_session.lock();
        if (!session || !session->session_active.load() || !session->send_message) {
            return;
        }
        // Forward backend events to the WebSocket client. The streaming
        // backend speaks OpenAI Realtime event types; pass through the text
        // events (delta/completed) and the stream-level audio events
        // (speech_started/speech_stopped/committed) the UI listens for.
        std::string event_type = msg.value("type", "");
        if (event_type == "conversation.item.input_audio_transcription.delta" ||
            event_type == "conversation.item.input_audio_transcription.completed" ||
            event_type == "input_audio_buffer.speech_started" ||
            event_type == "input_audio_buffer.speech_stopped" ||
            event_type == "input_audio_buffer.committed" ||
            event_type == "input_audio_buffer.cleared" ||
            event_type == "session.updated") {
            session->send_message(msg);
        }
    };

    // Retry with backoff: the backend's TCP listener starts on a separate
    // thread from its HTTP health endpoint, so a session created right after
    // model load can otherwise race it and silently lose streaming.
    bool ok = false;
    for (int attempt = 0; attempt < 10 && !ok; ++attempt) {
        if (attempt > 0) {
            std::this_thread::sleep_for(std::chrono::milliseconds(100));
        }
        ok = client->connect(address, callback);
    }

    if (ok) {
        std::lock_guard<std::mutex> lock(session->streaming_mutex);
        session->streaming_client = std::move(client);
        session->use_streaming_backend.store(true);
        LOG(INFO, "RealtimeSession") << "Connected to streaming backend at " << address << std::endl;
    } else {
        LOG(WARNING, "RealtimeSession") << "Failed to connect to streaming backend at " << address << std::endl;
    }
}

void RealtimeSessionManager::disconnect_streaming_backend(std::shared_ptr<RealtimeSession> session) {
    if (!session) {
        return;
    }
    session->use_streaming_backend.store(false);
    std::lock_guard<std::mutex> lock(session->streaming_mutex);
    if (session->streaming_client) {
        session->streaming_client->close();
        session->streaming_client.reset();
    }
}

void RealtimeSessionManager::forward_streaming_audio(std::shared_ptr<RealtimeSession> session,
                                                     const std::string& base64_audio) {
    if (!session) {
        return;
    }
    std::lock_guard<std::mutex> lock(session->streaming_mutex);
    if (!session->streaming_client || !session->streaming_client->is_connected()) {
        return;
    }
    json msg = {
        {"type", "input_audio_buffer.append"},
        {"audio", base64_audio}
    };
    session->streaming_client->send(msg);
}

void RealtimeSessionManager::forward_streaming_commit(std::shared_ptr<RealtimeSession> session) {
    if (!session) {
        return;
    }
    std::lock_guard<std::mutex> lock(session->streaming_mutex);
    if (!session->streaming_client || !session->streaming_client->is_connected()) {
        return;
    }
    json msg = {
        {"type", "input_audio_buffer.commit"}
    };
    session->streaming_client->send(msg);
}

void RealtimeSessionManager::forward_streaming_clear(std::shared_ptr<RealtimeSession> session) {
    if (!session) {
        return;
    }
    std::lock_guard<std::mutex> lock(session->streaming_mutex);
    if (!session->streaming_client || !session->streaming_client->is_connected()) {
        return;
    }
    json msg = {
        {"type", "input_audio_buffer.clear"}
    };
    session->streaming_client->send(msg);
}

void RealtimeSessionManager::close_session(const std::string& session_id) {
    std::shared_ptr<RealtimeSession> session;
    {
        std::lock_guard<std::mutex> lock(sessions_mutex_);
        auto it = sessions_.find(session_id);
        if (it != sessions_.end()) {
            session = it->second;
            session->session_active = false;
            session->interim_generation.fetch_add(1, std::memory_order_relaxed);
            sessions_.erase(it);
        }
    }
    if (session) {
        disconnect_streaming_backend(session);
    }
}

bool RealtimeSessionManager::session_exists(const std::string& session_id) const {
    std::lock_guard<std::mutex> lock(sessions_mutex_);
    return sessions_.find(session_id) != sessions_.end();
}

std::shared_ptr<RealtimeSession> RealtimeSessionManager::get_session(const std::string& session_id) {
    std::lock_guard<std::mutex> lock(sessions_mutex_);

    auto it = sessions_.find(session_id);
    if (it != sessions_.end()) {
        return it->second;
    }
    return nullptr;
}

} // namespace lemon
