#include "lemon/streaming_proxy.h"
#include <sstream>
#include <iostream>
#include <chrono>
#include <cstring>
#include <stdexcept>
#include <curl/curl.h>
#include <lemon/utils/aixlog.hpp>

namespace lemon {

namespace {

void extract_telemetry_from_chunk(const nlohmann::json& chunk, StreamingProxy::TelemetryData& telemetry) {
    nlohmann::json usage;
    if (chunk.contains("usage")) {
        usage = chunk["usage"];
    } else if (chunk.contains("response") && chunk["response"].is_object() && chunk["response"].contains("usage")) {
        usage = chunk["response"]["usage"];
    }

    if (usage.is_object()) {
        if (usage.contains("prompt_tokens")) {
            telemetry.input_tokens = usage["prompt_tokens"].get<int>();
        } else if (usage.contains("input_tokens")) {
            telemetry.input_tokens = usage["input_tokens"].get<int>();
        }
        if (usage.contains("prompt_tokens") || usage.contains("input_tokens")) {
            telemetry.prompt_tokens = telemetry.input_tokens;
        }
        if (usage.contains("completion_tokens")) {
            telemetry.output_tokens = usage["completion_tokens"].get<int>();
        } else if (usage.contains("output_tokens")) {
            telemetry.output_tokens = usage["output_tokens"].get<int>();
        }
        if (usage.contains("prefill_duration_ttft")) {
            telemetry.time_to_first_token = usage["prefill_duration_ttft"].get<double>();
        }
        if (usage.contains("decoding_speed_tps")) {
            telemetry.tokens_per_second = usage["decoding_speed_tps"].get<double>();
        }
        if (usage.contains("prompt_tokens_details") && usage["prompt_tokens_details"].is_object() &&
            usage["prompt_tokens_details"].contains("cached_tokens") &&
            usage["prompt_tokens_details"]["cached_tokens"].is_number()) {
            telemetry.cache_tokens = usage["prompt_tokens_details"]["cached_tokens"].get<int>();
        } else if (usage.contains("input_tokens_details") && usage["input_tokens_details"].is_object() &&
                   usage["input_tokens_details"].contains("cached_tokens") &&
                   usage["input_tokens_details"]["cached_tokens"].is_number()) {
            // Responses API usage shape.
            telemetry.cache_tokens = usage["input_tokens_details"]["cached_tokens"].get<int>();
        }
    }

    nlohmann::json timings;
    if (chunk.contains("timings")) {
        timings = chunk["timings"];
    } else if (chunk.contains("response") && chunk["response"].is_object() && chunk["response"].contains("timings")) {
        timings = chunk["response"]["timings"];
    }

    if (timings.is_object()) {
        if (timings.contains("prompt_n")) {
            telemetry.input_tokens = timings["prompt_n"].get<int>();
        }
        if (timings.contains("predicted_n")) {
            telemetry.output_tokens = timings["predicted_n"].get<int>();
        }
        if (timings.contains("prompt_ms")) {
            telemetry.time_to_first_token = timings["prompt_ms"].get<double>() / 1000.0;
        }
        if (timings.contains("predicted_per_second")) {
            telemetry.tokens_per_second = timings["predicted_per_second"].get<double>();
        }
        if (timings.contains("cache_n") && timings["cache_n"].is_number()) {
            telemetry.cache_tokens = timings["cache_n"].get<int>();
        }
    }
}

std::string lower_copy(std::string s) {
    std::transform(s.begin(), s.end(), s.begin(),
                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
    return s;
}

bool is_gpu_hang_or_compute_error(const std::string& message) {
    const std::string lowered = lower_copy(message);
    return lowered.find("compute error") != std::string::npos ||
           lowered.find("gpu hang") != std::string::npos;
}

} // namespace


void StreamingProxy::forward_sse_stream(
    const std::string& backend_url,
    const std::string& request_body,
    httplib::DataSink& sink,
    std::function<void(const TelemetryData&)> on_complete,
    long timeout_seconds,
    std::function<void()> on_chunk,
    long heartbeat_interval_ms) {

    TelemetryData telemetry;
    try {
        auto req_json = json::parse(request_body);
        if (req_json.contains("model") && req_json["model"].is_string()) {
            telemetry.model_name = req_json["model"].get<std::string>();
        }
    } catch (...) {}
    std::string line_buffer;
    bool stream_error = false;
    bool has_done_marker = false;
    bool has_first_token = false;
    double time_to_first_token = 0.0;
    const auto start_time = std::chrono::steady_clock::now();
    auto last_activity_time = start_time;

    int backend_status = 200;
    std::string error_body;
    static constexpr size_t max_error_body = 64 * 1024;

    auto process_line = [&telemetry](const std::string& line) {
        std::string json_str;
        if (line.find("data: ") == 0) {
            json_str = line.substr(6);
        } else if (line.find("ChatCompletionChunk: ") == 0) {
            json_str = line.substr(21);
        }
        if (!json_str.empty() && json_str != "[DONE]") {
            try {
                auto chunk = json::parse(json_str);
                extract_telemetry_from_chunk(chunk, telemetry);
            } catch (...) {}
        }
    };

    utils::HttpResponse result = utils::HttpClient::post_stream(
        backend_url,
        request_body,
        [&sink, &line_buffer, &has_done_marker, &has_first_token, &time_to_first_token,
         &start_time, &last_activity_time, &on_chunk, &process_line, &backend_status, &error_body](const char* data, size_t length) {
            last_activity_time = std::chrono::steady_clock::now();

            if (backend_status != 200) {
                if (error_body.size() < max_error_body) {
                    error_body.append(data, std::min(length, max_error_body - error_body.size()));
                }
                return true;
            }

            if (on_chunk) {
                on_chunk();
            }

            line_buffer.append(data, length);
            process_sse_lines(line_buffer, process_line);

            std::string chunk(data, length);
            if (!has_first_token && chunk.find("data: ") != std::string::npos) {
                has_first_token = true;
                time_to_first_token = std::chrono::duration<double>(
                    std::chrono::steady_clock::now() - start_time).count();
            }

            if (chunk.find("data: [DONE]") != std::string::npos) {
                has_done_marker = true;
            }

            if (!sink.write(data, length)) {
                return false;
            }

            return true;
        },
        {},
        timeout_seconds,
        [&backend_status](int status) { backend_status = status; },
        utils::HttpSecurityPolicy::TrustedLoopback,
        [&sink, &last_activity_time, &backend_status, heartbeat_interval_ms]() {
            if (sink.is_writable && !sink.is_writable()) {
                return true;
            }

            if (heartbeat_interval_ms > 0 && backend_status == 200) {
                const auto now = std::chrono::steady_clock::now();
                const auto elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(
                    now - last_activity_time).count();
                if (elapsed >= heartbeat_interval_ms) {
                    static constexpr const char* heartbeat = ": ping\n\n";
                    if (!sink.write(heartbeat, std::strlen(heartbeat))) {
                        return true;
                    }
                    last_activity_time = now;
                }
            }

            return false;
        }
    );

    const bool client_disconnected =
        result.curl_code == CURLE_WRITE_ERROR ||
        result.curl_code == CURLE_ABORTED_BY_CALLBACK;
    const bool transport_interrupted =
        result.curl_code == CURLE_PARTIAL_FILE || result.curl_code == CURLE_RECV_ERROR;

    if (result.curl_code != CURLE_OK) {
        if (client_disconnected) {
            stream_error = true;
            LOG(WARNING, "StreamingProxy") << "Client disconnected during SSE stream (CURL error: " << result.curl_error << ")" << std::endl;
            telemetry.error_message = "Client disconnected during stream";
        } else if (transport_interrupted) {
            if (!has_done_marker) {
                // This is the important crash path: HTTP headers may have been sent and
                // some bytes may even have reached the client, but the SSE protocol never
                // completed. Do not synthesize [DONE], because that hides backend crashes
                // from the router and leaves stale loaded-model state behind.
                throw std::runtime_error(
                    "backend connection failed during SSE stream before DONE: CURL error: " +
                    result.curl_error);
            }
        } else {
            stream_error = true;
            LOG(ERROR, "StreamingProxy") << "SSE stream failed: CURL error: " << result.curl_error << std::endl;
            telemetry.error_message = "SSE stream failed: CURL error: " + result.curl_error;
        }
    }

    if (!client_disconnected &&
        (result.status_code != 200 || backend_status != 200)) {
        const int status = backend_status != 200 ? backend_status : result.status_code;
        LOG(ERROR, "StreamingProxy") << "Backend returned error: " << status
                                     << (error_body.empty() ? "" : ": " + error_body) << std::endl;
        telemetry.error_message = "Backend returned error status code: " + std::to_string(status);

        if (is_gpu_hang_or_compute_error(error_body)) {
            throw std::runtime_error("backend compute error / GPU hang during streaming: " + error_body);
        }

        stream_error = true;

        // The response is already committed as 200 text/event-stream, so an
        // unframed error body is dropped by every spec-compliant client parser.
        // No [DONE] follows, matching OpenAI's behavior for in-stream errors.
        json payload;
        try {
            payload = json::parse(error_body);
        } catch (...) {
            payload = nullptr;
        }
        if (!payload.is_object() || !payload.contains("error")) {
            std::string message = error_body.empty()
                ? "backend returned HTTP " + std::to_string(status)
                : error_body;
            payload = json{{"error", {{"message", message},
                                      {"type", "backend_error"},
                                      {"status_code", status}}}};
        } else if (payload["error"].is_object() && !payload["error"].contains("status_code")) {
            // A backend's own error object carries no transport status, so
            // adapters downstream would have to guess one.
            payload["error"]["status_code"] = status;
        }
        const std::string event = "data: " + payload.dump() + "\n\n";
        sink.write(event.data(), event.size());
    }

    if (!stream_error) {
        // Ensure [DONE] marker is sent only for clean transports. If the transport
        // was interrupted before [DONE], the block above throws and recovery is
        // handled by WrappedServer/Router instead of pretending success.
        if (!has_done_marker) {
            LOG(WARNING, "StreamingProxy") << "WARNING: Backend did not send [DONE] marker, adding it" << std::endl;
            const char* done_marker = "data: [DONE]\n\n";
            sink.write(done_marker, strlen(done_marker));
        }

        sink.done();

        LOG(INFO, "Server") << "Streaming completed - 200 OK" << std::endl;

        if (!line_buffer.empty()) {
            if (line_buffer.back() == '\r') {
                line_buffer.pop_back();
            }
            process_line(line_buffer);
        }

        if (telemetry.time_to_first_token <= 0.0) {
            telemetry.time_to_first_token = time_to_first_token;
        }
        if (telemetry.tokens_per_second <= 0.0 && telemetry.output_tokens > 0) {
            double total_duration = std::chrono::duration<double>(
                std::chrono::steady_clock::now() - start_time).count();
            double decode_duration = total_duration - telemetry.time_to_first_token;
            if (decode_duration > 0.0) {
                telemetry.tokens_per_second = telemetry.output_tokens / decode_duration;
            }
        }
        telemetry.print();

        if (on_complete) {
            on_complete(telemetry);
        }
    } else {
        sink.done();
        if (on_complete) {
            on_complete(telemetry);
        }
    }
}

void StreamingProxy::forward_byte_stream(
    const std::string& backend_url,
    const std::string& request_body,
    httplib::DataSink& sink,
    long timeout_seconds,
    std::function<void()> on_chunk) {

    bool stream_error = false;

    // On a non-200 the backend body is an error description, not payload bytes:
    // divert it here instead of the client sink, so it can be reshaped into a
    // JSON error below rather than served as successful media.
    int backend_status = 200;
    std::string error_body;
    static constexpr size_t max_error_body = 64 * 1024;

    utils::HttpResponse result = utils::HttpClient::post_stream(
        backend_url,
        request_body,
        [&sink, &on_chunk, &backend_status, &error_body](const char* data, size_t length) {
            if (backend_status != 200) {
                if (error_body.size() < max_error_body) {
                    error_body.append(data, std::min(length, max_error_body - error_body.size()));
                }
                return true;
            }

            if (on_chunk) {
                on_chunk();
            }

            if (!sink.write(data, length)) {
                return false;
            }

            return true;
        },
        {},
        timeout_seconds,
        [&backend_status](int status) { backend_status = status; },
        utils::HttpSecurityPolicy::TrustedLoopback
    );

    const bool transport_interrupted =
        result.curl_code == CURLE_PARTIAL_FILE || result.curl_code == CURLE_RECV_ERROR;

    if (result.curl_code != CURLE_OK) {
        stream_error = true;
        if (result.curl_code == CURLE_WRITE_ERROR) {
            LOG(WARNING, "StreamingProxy") << "Client disconnected during byte stream (CURL error: " << result.curl_error << ")" << std::endl;
        } else if (transport_interrupted) {
            // Keep byte streams consistent with SSE: an interrupted transport is a
            // backend failure, not a clean stream completion. The caller will mark
            // the backend unavailable and reload after the current response unwinds.
            throw std::runtime_error(
                "backend connection failed during byte stream: CURL error: " +
                result.curl_error);
        } else {
            LOG(ERROR, "StreamingProxy") << "Byte stream failed: CURL error: " << result.curl_error << std::endl;
        }
    }

    if (result.status_code != 200 || backend_status != 200) {
        const int status = backend_status != 200 ? backend_status : result.status_code;
        LOG(ERROR, "StreamingProxy") << "Backend returned error " << status
                                     << (error_body.empty() ? "" : ": " + error_body) << std::endl;

        if (is_gpu_hang_or_compute_error(error_body)) {
            throw std::runtime_error("backend compute error / GPU hang during byte stream: " + error_body);
        }

        stream_error = true;

        json payload;
        try {
            payload = json::parse(error_body);
        } catch (...) {
            payload = nullptr;
        }
        if (!payload.is_object() || !payload.contains("error")) {
            std::string message = error_body.empty()
                ? "backend returned HTTP " + std::to_string(status)
                : error_body;
            payload = json{{"error", {{"message", message},
                                      {"type", "backend_error"},
                                      {"status_code", status}}}};
        } else if (payload["error"].is_object() && !payload["error"].contains("status_code")) {
            payload["error"]["status_code"] = status;
        }
        const std::string out = payload.dump();
        sink.write(out.data(), out.size());
    }

    if (!stream_error) {
        LOG(INFO, "Server") << "Streaming completed - 200 OK" << std::endl;
    }
    sink.done();
}

StreamingProxy::TelemetryData StreamingProxy::extract_telemetry(const nlohmann::json& payload) {
    TelemetryData telemetry;
    extract_telemetry_from_chunk(payload, telemetry);
    return telemetry;
}

StreamingProxy::TelemetryData StreamingProxy::parse_telemetry(const std::string& buffer) {
    TelemetryData telemetry;

    std::istringstream stream(buffer);
    std::string line;
    json last_chunk_with_usage;

    while (std::getline(stream, line)) {
        std::string json_str;
        if (line.find("data: ") == 0) {
            json_str = line.substr(6);
        } else if (line.find("ChatCompletionChunk: ") == 0) {
            json_str = line.substr(21);
        }

        if (!json_str.empty() && json_str != "[DONE]") {
            try {
                auto chunk = json::parse(json_str);
                bool has_usage = chunk.contains("usage") || chunk.contains("timings");
                if (!has_usage && chunk.contains("response") && chunk["response"].is_object()) {
                    has_usage = chunk["response"].contains("usage") || chunk["response"].contains("timings");
                }
                if (has_usage) {
                    last_chunk_with_usage = chunk;
                }
            } catch (...) {}
        }
    }

    if (!last_chunk_with_usage.empty()) {
        try {
            extract_telemetry_from_chunk(last_chunk_with_usage, telemetry);
        } catch (const std::exception& e) {
            LOG(ERROR, "StreamingProxy") << "Error parsing telemetry: " << e.what() << std::endl;
        }
    }

    return telemetry;
}

void StreamingProxy::process_sse_lines(std::string& line_buffer, std::function<void(const std::string&)> line_callback) {
    size_t pos;
    while ((pos = line_buffer.find('\n')) != std::string::npos) {
        std::string line = line_buffer.substr(0, pos);
        line_buffer.erase(0, pos + 1);
        if (!line.empty() && line.back() == '\r') {
            line.pop_back();
        }
        line_callback(line);
    }
}

void StreamingProxy::accumulate_responses_delta(const nlohmann::json& parsed, std::string& accumulated_text) {
    if (parsed.contains("choices") && parsed["choices"].is_array() && !parsed["choices"].empty()) {
        auto choice = parsed["choices"][0];
        if (choice.is_object() && choice.contains("delta")) {
            auto delta = choice["delta"];
            if (delta.is_object() && delta.contains("content") && delta["content"].is_string()) {
                accumulated_text += delta["content"].get<std::string>();
            }
        }
    }
    if (parsed.contains("response") && parsed["response"].is_string()) {
        accumulated_text += parsed["response"].get<std::string>();
    }
    // Supports Responses API type-restricted delta vs. backward compatible fallback.
    if (parsed.contains("delta")) {
        bool should_extract_delta = true;
        if (parsed.contains("type")) {
            should_extract_delta = (parsed["type"] == "response.output_text.delta");
        }
        if (should_extract_delta) {
            if (parsed["delta"].is_string()) {
                accumulated_text += parsed["delta"].get<std::string>();
            } else if (parsed["delta"].is_object() && parsed["delta"].contains("text") && parsed["delta"]["text"].is_string()) {
                accumulated_text += parsed["delta"]["text"].get<std::string>();
            }
        }
    }
}

} // namespace lemon
