#pragma once

#include <iostream>
#include <syncstream>

#include <string>
#include <map>
#include <queue>

#include <thread>
#include <mutex>

#include <functional>
#include <variant>
#include <optional>

#define WIN32_LEAN_AND_MEAN
#include <Windows.h>
#include <io.h>
#include <fcntl.h>

#include "json.hpp"

// A simple JSON-RPC 2.0 implementation for C++ on Windows, using nlohmann::json for JSON handling.
// See https://www.jsonrpc.org/specification for details on the protocol.
// See https://github.com/nlohmann/json for the JSON library.
// Generated by Gemini and edited by hand to fit the Windows console I/O model.
namespace jsonrpc {

    using json = nlohmann::json;

    struct Error {
        int code = 0;
        std::string message;
        json data = nullptr;

        // Hidden Friend Idiom
        friend void to_json(json& j, const Error& e) {
            j = json{ {"code", e.code}, {"message", e.message} };
            if (!e.data.is_null()) j["data"] = e.data;
        }

        friend void from_json(const json& j, Error& e) {
            j.at("code").get_to(e.code);
            j.at("message").get_to(e.message);
            if (j.contains("data")) e.data = j.at("data");
        }
    };
	// In JSON-RPC spec, id can be string, number, or null. 
    // For simplicity, we use int here.
    struct Request {
        int id = 0;
        std::string method;
        json params = nullptr;
    };

    struct Response {
        int id = 0;
        std::variant<json, Error> content;

        bool is_error() const { return std::holds_alternative<Error>(content); }

        static Response make_success(int id, json result) {
            return Response{ id, std::move(result) };
        }

        static Response make_error(int id, int code, std::string msg, json data = nullptr) {
            return Response{ id, Error{code, std::move(msg), std::move(data)} };
        }
    };

    inline void to_json(json& j, const Request& r) {
        j = json{ {"jsonrpc", "2.0"}, {"id", r.id}, {"method", r.method} };
        if (!r.params.is_null()) j["params"] = r.params;
    }

    inline void to_json(json& j, const Response& r) {
        j = json{ {"jsonrpc", "2.0"}, {"id", r.id} };
        if (std::holds_alternative<Error>(r.content)) {
            j["error"] = std::get<Error>(r.content);
        }
        else {
            j["result"] = std::get<json>(r.content);
        }
    }

	// Notification is not implemented in this simple version, 
    // but can be added similarly to Request without id.

    using IncomingMessage = std::variant<Request, Response>;

	// JSON parser for incoming messages. 
    // It can throw exceptions if the JSON is invalid 
    // or doesn't conform to the expected structure.
    struct Parser {
        static IncomingMessage parse(const json& j) {
            if (!j.is_object()) throw std::runtime_error("Invalid JSON-RPC: not an object");

            if (!j.contains("id") || !j["id"].is_number_integer()) {
                throw std::runtime_error("Invalid JSON-RPC: id must be an integer");
            }
            int id = j["id"].get<int>();
			// Request must have "method", Response must have "result" or "error"
            if (j.contains("method")) {
                return Request{ id, j["method"], j.value("params", json(nullptr)) };
            }
            // error
            if (j.contains("error")) {
                return Response{ id, j["error"].get<Error>() };
            }
            // result
            else if (j.contains("result")) {
                return Response{ id, j["result"] };
            }
			// Seems Gemini like uncheked error very much, interesting...
            throw std::runtime_error("Invalid JSON-RPC: missing method, result, or error");
        }
		// Firstly, parse the string into JSON, then parse the JSON into IncomingMessage.
        static IncomingMessage parse(const std::string& str) {
            return parse(json::parse(str));
        }
    };
}

#define WM_JSONRPC_MESSAGE (WM_USER + 101)

// Code for stdio-based JSON-RPC connection handling on Windows.
// Specifically for Win32 Message Loop integration. It reads from stdin,
// writes to stdout, and uses PostThreadMessage to notify the main thread of new messages.
namespace jsonrpc {

	// Thread-safe queue for incoming messages.
    template<typename T>
    class ThreadSafeQueue {
    private:
        std::queue<T> queue_;
        std::mutex mutex_;
    public:
        void push(T value) {
            std::lock_guard<std::mutex> lock(mutex_);
            queue_.push(std::move(value));
        }
        bool try_pop(T& value) {
            std::lock_guard<std::mutex> lock(mutex_);
            if (queue_.empty()) return false;
            value = std::move(queue_.front());
            queue_.pop();
            return true;
        }
    };
	// The main connection class that manages the JSON-RPC communication.
	// This name comes from Emacs's jsonrpc.el, which uses jsonrpc-connection 
    // as the basic class for handling JSON-RPC connections.
    class Conn {
    public:
        using RequestHandler = std::function<json(const json&)>;
        using ResponseHandler = std::function<void(const Response&)>;

        Conn() : next_id_(1), running_(false), main_thread_id_(0) {
            // Force Windows stdin/stdout into binary mode to prevent \r\n 
            // translation that can cause Content-Length miscalculation or byte reading errors.
            (void)_setmode(_fileno(stdin), _O_BINARY);
            (void)_setmode(_fileno(stdout), _O_BINARY);
        }

        ~Conn() { stop(); }

        void start() {
            if (running_) return;
			// Get main thread ID for PostThreadMessage. 
            // We assume start() is called from the main thread.
            main_thread_id_ = GetCurrentThreadId();
            running_ = true;
			// Start the reader thread that continuously reads from stdin
            // and pushes messages to the inbox queue.
            reader_thread_ = std::thread([this]() { read_loop(); });
        }

        void stop() {
            running_ = false;
            if (reader_thread_.joinable()) reader_thread_.detach();
        }
		// Register a method handler for incoming requests.
        void register_method(const std::string& name, RequestHandler handler) {
            std::lock_guard<std::mutex> lock(map_mutex_);
            method_handlers_[name] = handler;
        }
		// Send a request to the other side, with a callback for the response.        
        void send_request(const std::string& method, const json& params, ResponseHandler callback) {
            int id = next_id_++;
            {
                std::lock_guard<std::mutex> lock(map_mutex_);
                pending_callbacks_[id] = callback;
            }
            Request req{ id, method, params };
            json j;
            to_json(j, req);
            send_message(j.dump());
        }
		// Stdin Read Loop
        void process_queue() {
            IncomingMessage msg;
            while (inbox_.try_pop(msg)) {
                if (std::holds_alternative<Request>(msg)) {
                    handle_request(std::get<Request>(msg));
                }
                else if (std::holds_alternative<Response>(msg)) {
                    handle_response(std::get<Response>(msg));
                }
            }
        }

    private:
        std::atomic<bool> running_;
        std::atomic<int> next_id_;
        std::thread reader_thread_;
        DWORD main_thread_id_;

        ThreadSafeQueue<IncomingMessage> inbox_;

        std::mutex map_mutex_;
        std::map<std::string, RequestHandler> method_handlers_;
        std::map<int, ResponseHandler> pending_callbacks_;

        void send_message(const std::string& body) {
            // std::osyncstream will atomically write the buffer to the stream 
            // when it is destructed, so we don't need to manually lock.

			// Emacs's jsonrpc use a HTTP-like framing with Content-Length header,
            // so we follow the same format here:
			// Content-Length: <length>\r\n\r\n<body>
            std::osyncstream(std::cout)
                << "Content-Length: " << body.length() << "\r\n"
                << "\r\n" << body << std::flush;
        }
        void read_loop() {
            // A helper function that reads a line and automatically removes the Windows \r character.
            auto read_header_line = []() -> std::optional<std::string> {
                std::string line;
                if (!std::getline(std::cin, line)) return std::nullopt; // EOF or stream error
                if (!line.empty() && line.back() == '\r') line.pop_back();
                return line;
                };

			// If we can ensure that input always is valid JSON and well-formed, 
            // we can simplify the reading logic by just reading the Content-Length header 
            // and then reading the exact number of bytes for the body.
			// Thus, -32700 and -32600 error handling is not needed.
            while (running_) {
                int content_length = 0;

                // Read header.
                while (true) {
                    auto line = read_header_line();
                    if (!line) return;          // stream error or EOF, exit the loop
                    if (line->empty()) break;   // empty line indicates end of headers, ready to read body

                    if (line->starts_with("Content-Length: ")) {
                        // "Content-Length: " is 16 characters long.
                        try { content_length = std::stoi(line->substr(16)); }
                        catch (...) {}
                    }
                }
                // invalid content length, skip to next message
                if (content_length <= 0) continue;

                // Read whole body.
                std::vector<char> buffer(content_length);
                std::cin.read(buffer.data(), content_length);

                // Make sure we read the exact number of bytes specified by Content-Length. 
                // If not, it means the stream is broken or we reached EOF, so we should exit the loop.
                if (std::cin.gcount() != content_length) break;

                // Parse and post message to main thread
                try {
                    // Gemini says below is faster.
                    //auto msg = Parser::parse(std::string(buffer.begin(), buffer.end()));
                    auto j = json::parse(buffer.begin(), buffer.end());
                    auto msg = Parser::parse(j);
                    inbox_.push(std::move(msg));

                    if (main_thread_id_ != 0) {
                        PostThreadMessage(main_thread_id_, WM_JSONRPC_MESSAGE, 0, 0);
                    }
                }
                catch (...) {
                    // Ignore parsing errors and continue reading the next message.
                }
            }
        }

        void handle_request(const Request& req) {
            RequestHandler handler = nullptr;
            {
                std::lock_guard<std::mutex> lock(map_mutex_);
                if (method_handlers_.count(req.method)) {
                    handler = method_handlers_[req.method];
                }
            }

            bool success = false;
            json result_json;
            std::string error_msg;

            if (handler) {
                try {
                    result_json = handler(req.params);
                    success = true;
                }
                catch (const std::exception& e) {
                    error_msg = e.what();
                }
            }
            else {
                error_msg = "Method not found";
            }

            json j_resp;
            if (success) {
                to_json(j_resp, Response::make_success(req.id, result_json));
            }
            else {
				// TODO: Consider -32602 Invalid params error if the handler 
                // throws an exception related to parameter parsing or validation.
                int code = (error_msg == "Method not found") ? -32601 : -32603;
                to_json(j_resp, Response::make_error(req.id, code, error_msg));
            }
            send_message(j_resp.dump());
        }

        void handle_response(const Response& resp) {
            ResponseHandler callback = nullptr;
            {
                std::lock_guard<std::mutex> lock(map_mutex_);
                auto it = pending_callbacks_.find(resp.id);
                if (it != pending_callbacks_.end()) {
                    callback = it->second;
                    pending_callbacks_.erase(it);
                }
            }
            if (callback) callback(resp);
        }
    };
}