#include "AngelscriptDebugger.hpp"
#include "ASVariableFormatter.hpp"
#include "DebugAdapterProtocol/BaseProtocol.hpp"
#include "DebugAdapterProtocol/Events.hpp"
#include "DebugAdapterProtocol/Requests.hpp"

void AngelscriptDebugger::Run(uint16_t port) {
    _server = new asio::ip::tcp::acceptor(_ioContext, asio::ip::tcp::endpoint(asio::ip::tcp::v4(), port));
    std::thread([this]() { AcceptLoop(); }).detach();
}

template <typename... Args> std::string string_format(const std::string& format, Args... args) {
    int size_s = std::snprintf(nullptr, 0, format.c_str(), args...) + 1; // Extra space for '\0'
    if (size_s <= 0) {
        throw std::runtime_error("Error during formatting.");
    }
    auto size = static_cast<size_t>(size_s);
    auto buf = std::make_unique<char[]>(size);
    std::snprintf(buf.get(), size, format.c_str(), args...);
    return std::string(buf.get(), buf.get() + size - 1); // We don't want the '\0' inside
}

std::string AngelscriptDebugger::GetResolvedScriptPath(const char* scriptSection) {
    if (_scriptPath.has_value()) {
        return string_format(_scriptPath.value(), scriptSection);
    }
    return scriptSection;
}

void AngelscriptDebugger::on_line_callback(asIScriptContext* ctx, AngelscriptDebugger* d) {
    if (ctx->GetState() == asEXECUTION_SUSPENDED) {
        return;
    }
    if (!d->HasDebuggerAttached()) {
        return;
    }
    const char* scriptSection = nullptr;
    int column = 0;
    int line = ctx->GetLineNumber(0, &column, &scriptSection);
    if (line == 0) {
        return;
    }

    if (d->_next && d->_nextContext == ctx && d->_nextDepth >= ctx->GetCallstackSize()) {
        d->_next = false;
        ctx->Suspend();
        d->_pausedContexts.push_back(ctx);
        auto* o = new DebugAdapterProtocol::StoppedEvent(new DebugAdapterProtocol::StoppedEventBody("step", "", ""));

        d->Send(o);
        return;
    }
    if (d->_stepInto && d->_nextContext == ctx) {
        d->_stepInto = false;
        ctx->Suspend();
        d->_pausedContexts.push_back(ctx);
        auto* o = new DebugAdapterProtocol::StoppedEvent(new DebugAdapterProtocol::StoppedEventBody("step", "", ""));

        d->Send(o);
        return;
    }
    if (d->_stepOut && d->_nextContext == ctx && d->_nextDepth > ctx->GetCallstackSize()) {
        d->_stepOut = false;
        ctx->Suspend();
        d->_pausedContexts.push_back(ctx);
        auto* o = new DebugAdapterProtocol::StoppedEvent(new DebugAdapterProtocol::StoppedEventBody("step", "", ""));

        d->Send(o);
        return;
    }

    auto resolvedScriptPath = d->GetResolvedScriptPath(scriptSection);
    auto sectionBreakpoints = d->_breakpoints.find(resolvedScriptPath);
    if (sectionBreakpoints != d->_breakpoints.end()) {
        if (sectionBreakpoints->second.contains(line)) {
            d->EnterBreakpoint(ctx, resolvedScriptPath, line);
        }
    }
}

size_t GetBreakpointHash(const std::string& section, size_t line) {
    using std::hash;
    using std::size_t;
    using std::string;

    return ((hash<string>()(section) ^ (hash<size_t>()(line) << 1)) >> 1);
}

void AngelscriptDebugger::on_exception_callback(asIScriptContext* ctx, AngelscriptDebugger* d) {
    if (!d->HasDebuggerAttached()) {
        return;
    }
    ctx->Suspend();
    d->_pausedContexts.push_back(ctx);

    const char* scriptSection = nullptr;
    int column = 0;
    int line = ctx->GetLineNumber(0, &column, &scriptSection);
    if (line == 0)
        return;

    auto exception = ctx->GetExceptionString();

    auto* o = new DebugAdapterProtocol::StoppedEvent(
        new DebugAdapterProtocol::StoppedEventBody("exception", "Paused on exception", exception));

    d->Send(o);
}

void AngelscriptDebugger::EnterBreakpoint(asIScriptContext* ctx, const std::string& section, size_t line) {
    ctx->Suspend();
    _pausedContexts.push_back(ctx);

    auto* o = new DebugAdapterProtocol::StoppedEvent(
        new DebugAdapterProtocol::StoppedEventBody("breakpoint", GetBreakpointHash(section, line)));

    Send(o);
}

void AngelscriptDebugger::RegisterContext(asIScriptContext* ctx) {
    ctx->SetLineCallback(asFUNCTION(on_line_callback), this, asCALL_CDECL);
    ctx->SetExceptionCallback(asFUNCTION(on_exception_callback), this, asCALL_CDECL);
}

[[noreturn]] void AngelscriptDebugger::AcceptLoop() {
    while (true) {
        auto* client = new asio::ip::tcp::socket(_ioContext);
        _server->accept(*client);
        std::thread([this, client]() { ClientLoop(*client); }).detach();
    }
}

inline std::string trim(const std::string& s) {
    auto wsfront = std::find_if_not(s.begin(), s.end(), [](int c) { return std::isspace(c); });
    auto wsback = std::find_if_not(s.rbegin(), s.rend(), [](int c) { return std::isspace(c); }).base();
    return (wsback <= wsfront ? std::string() : std::string(wsfront, wsback));
}

void AngelscriptDebugger::ClientLoop(asio::ip::tcp::socket& client) {
    std::unordered_map<std::string, std::string> headers;
    bool in_header_mode = true;
    asio::streambuf buffer;
    while (true) {
        if (!client.is_open()) {
            break;
        }
        if (client.available() == 0 && buffer.size() == 0) {
            continue;
        }
        if (client.available() > 0 && buffer.size() == 0) {
            asio::error_code error;
            asio::read_until(client, buffer, "\r\n", error);
        }

        if (in_header_mode) {
            std::istream str(&buffer);
            std::string s;
            std::getline(str, s);
            if (s == "\r") {
                if (headers.contains("content-length")) {
                    in_header_mode = false;
                }
                continue;
            }
            auto delimiter = s.find(':');
            std::string key = trim(s.substr(0, delimiter));
            std::transform(key.begin(), key.end(), key.begin(), [](unsigned char c) { return std::tolower(c); });
            std::string value = trim(s.substr(delimiter + 1));
            headers[key] = value;
        } else {
            auto contentLength = headers["content-length"];
            auto size = (size_t)std::strtol(contentLength.c_str(), nullptr, 10);
            std::cout << "message length: " << size << std::endl;
            std::stringstream msg;
            int64_t to_fetch = (int64_t)size - (int64_t)buffer.size();
            auto to_read = std::min(buffer.size(), size);
            asio::streambuf::const_buffers_type constBuffer = buffer.data();
            std::copy(asio::buffers_begin(constBuffer), asio::buffers_end(constBuffer),
                      std::ostream_iterator<char>(msg));
            buffer.consume(to_read);
            if (to_fetch > 0) {
                std::vector<uint8_t> v(to_fetch);
                asio::read(client, asio::buffer(v), asio::transfer_exactly(to_fetch));
                msg << std::string(v.begin(), v.end());
            }

            auto m = msg.str();
            std::cout << "found message: " << m << std::endl;
            nlohmann::json j = nlohmann::json::parse(m.begin(), m.end());
            auto protocolMessage = DebugAdapterProtocol::ProtocolMessage::FromJson(j);
            OnMessage(&client, protocolMessage);

            in_header_mode = true;
        }
    }
}

void AngelscriptDebugger::Send(DebugAdapterProtocol::ProtocolMessage* msg) {
    auto body = msg->ToJson().dump();
    std::vector<uint8_t> vec(body.begin(), body.end());
    for (const auto& conn : _connections) {
        if (conn->is_open()) {
            asio::error_code ignored_error;
            asio::write(*conn, asio::buffer("Content-Length: " + std::to_string(vec.size()) + "\r\n\r\n"),
                        ignored_error);
            asio::write(*conn, asio::buffer(vec), ignored_error);
        }
    }
    std::cout << "Sent message: " << body << std::endl;
    delete msg;
}

void AngelscriptDebugger::Send(asio::ip::tcp::socket* client, DebugAdapterProtocol::ProtocolMessage* msg) {
    auto body = msg->ToJson().dump();
    std::vector<uint8_t> vec(body.begin(), body.end());
    if (client->is_open()) {
        asio::error_code ignored_error;
        asio::write(*client, asio::buffer("Content-Length: " + std::to_string(vec.size()) + "\r\n\r\n"), ignored_error);
        asio::write(*client, asio::buffer(vec), ignored_error);
    }
    std::cout << "Sent message: " << body << std::endl;
    delete msg;
}

void AngelscriptDebugger::OnMessage(asio::ip::tcp::socket* client, DebugAdapterProtocol::ProtocolMessage* msg) {
    if (msg->type == "request") {
        OnRequest(client, dynamic_cast<DebugAdapterProtocol::Request*>(msg));
    }
    delete msg;
}

const char empty[] = "";
void AngelscriptDebugger::OnRequest(asio::ip::tcp::socket* client, DebugAdapterProtocol::Request* msg) {
    if (msg->GetCommand() == "setBreakpoints") {
        auto* t = dynamic_cast<DebugAdapterProtocol::SetBreakpointsRequest*>(msg);
        auto& args = t->arguments.value();
        auto& path = args->source.path.value();
        auto* response = new DebugAdapterProtocol::SetBreakpointsResponse(msg->seq);
        auto* body = new DebugAdapterProtocol::SetBreakpointsResponseBody();
        response->body = body;
        std::unordered_set<size_t> sectionBreakpoints;
        if (args->breakpoints.has_value()) {
            for (auto& bp : args->breakpoints.value()) {
                sectionBreakpoints.insert(bp.line);
                body->breakpoints.emplace_back(GetBreakpointHash(path, bp.line), args->source, bp.line);
            }
        }
        _breakpoints[path] = sectionBreakpoints;
        Send(client, response);
    } else if (msg->GetCommand() == "initialize") {
        auto response = new DebugAdapterProtocol::InitializeResponse(msg->seq);
        auto* body = new DebugAdapterProtocol::InitializeResponseBody();
        response->body = body;
        Send(client, response);

        auto r = new DebugAdapterProtocol::InitializedEvent();
        Send(client, r);

    } else if (msg->GetCommand() == "disconnect") {
        auto response = new DebugAdapterProtocol::DefinedResponse<DebugAdapterProtocol::ResponseBody, empty>(msg->seq);
        Send(client, response);
        _connections.erase(client);
    } else if (msg->GetCommand() == "configurationDone") {
        _connections.insert(client);
        Send(client, new DebugAdapterProtocol::DefinedResponse<DebugAdapterProtocol::ResponseBody, empty>(msg->seq));
    } else if (msg->GetCommand() == "attach") {
        auto t = dynamic_cast<DebugAdapterProtocol::AttachRequest*>(msg);
        _scriptPath = t->arguments.value()->scriptPath;
        Send(client, new DebugAdapterProtocol::DefinedResponse<DebugAdapterProtocol::ResponseBody, empty>(msg->seq));
    } else if (msg->GetCommand() == "threads") {
        auto response = new DebugAdapterProtocol::ThreadsResponse(msg->seq);
        auto* body = new DebugAdapterProtocol::ThreadsResponseBody();
        response->body = body;
        body->threads = {DebugAdapterProtocol::Thread(0, "main")};
        Send(client, response);
    } else if (msg->GetCommand() == "continue") {
        Continue();
        Send(client, new DebugAdapterProtocol::DefinedResponse<DebugAdapterProtocol::ResponseBody, empty>(msg->seq));
    } else if (msg->GetCommand() == "stackTrace") {
        auto ctx = _pausedContexts[0];
        auto response = new DebugAdapterProtocol::StackTraceResponse(msg->seq);
        auto stackTrace = std::vector<DebugAdapterProtocol::StackFrame>();
        stackTrace.reserve(ctx->GetCallstackSize());
        for (asUINT i = 0; i < ctx->GetCallstackSize(); ++i) {
            auto func = ctx->GetFunction(i);
            const char* scriptSection = nullptr;
            int column = 0;
            int line;
            if (i == 0 && ctx->GetExceptionFunction() != nullptr) {
                line = ctx->GetExceptionLineNumber(&column, &scriptSection);
            } else {
                line = ctx->GetLineNumber(i, &column, &scriptSection);
            }

            stackTrace.emplace_back(i, func->GetName(),
                                    DebugAdapterProtocol::Source(GetResolvedScriptPath(scriptSection)), line, column);
        }
        auto body = new DebugAdapterProtocol::StackTraceResponseBody(stackTrace);
        response->body = body;
        Send(client, response);
    } else if (msg->GetCommand() == "scopes") {
        auto t = dynamic_cast<DebugAdapterProtocol::ScopesRequest*>(msg);
        auto ctx = _pausedContexts[0];
        auto response = new DebugAdapterProtocol::ScopesResponse(msg->seq);
        auto scopes = std::vector<DebugAdapterProtocol::Scope>();
        auto frameId = t->arguments.value()->frameId;
        auto varCount = 0;
        for (int i = 0; i < ctx->GetVarCount(t->arguments.value()->frameId); ++i) {
            if (ctx->IsVarInScope(i, frameId)) {
                varCount++;
            }
        }
        _storedVariableReferences.emplace_back(std::make_unique<StackScope>(frameId, 0));
        scopes.emplace_back("Locals", _storedVariableReferences.size(), "locals", varCount);
        auto body = new DebugAdapterProtocol::ScopesResponseBody(scopes);
        response->body = body;
        Send(client, response);
    } else if (msg->GetCommand() == "variables") {
        auto t = dynamic_cast<DebugAdapterProtocol::VariablesRequest*>(msg);
        auto ctx = _pausedContexts[0];
        auto response = new DebugAdapterProtocol::VariablesResponse(msg->seq);
        auto variables = std::vector<DebugAdapterProtocol::Variable>();
        auto reference = t->arguments.value()->variablesReference - 1;
        if (reference >= _storedVariableReferences.size())
            return;

        auto& variant = _storedVariableReferences[reference];
        auto e = ctx->GetEngine();
        // We currently can't use asIScriptContext::PushBack on suspended script contexts due to limitations in
        // angelscript. Hence we make a new contenxt
        auto c = e->CreateContext();

        if (holds_alternative<std::unique_ptr<StackScope>>(variant)) {
            if (ctx->GetExceptionString() != nullptr) {
                variables.push_back(
                    DebugAdapterProtocol::Variable("exception", ctx->GetExceptionString(), "exception", {}));
            }
            auto frameId = std::get<std::unique_ptr<StackScope>>(variant)->StackLevel;
            // auto scopeType = get<StackScope>(variant).Type;
            auto varCount = ctx->GetVarCount(frameId);
            for (int i = 0; i < varCount; ++i) {
                if (!ctx->IsVarInScope(i, frameId)) {
                    continue;
                }
                auto name = ctx->GetVarName(i, frameId);
                auto type = ctx->GetVarTypeId(i, frameId);
                auto val = ctx->GetAddressOfVar(i, frameId);
                variables.push_back(ASVariableFormatter::GetAsDAPVariable(e, c, this, name, type, val));
            }
        } else {
            auto& v = std::get<std::unique_ptr<PointerVariable>>(variant);
            try {
                ASVariableFormatter::GetChildDAPVariables(variables, e, c, this, v->TypeID, v->Address);
            } catch (const std::exception& e) {
                std::cout << e.what() << std::endl;
            }
        }
        c->Release();
        auto body = new DebugAdapterProtocol::VariablesResponseBody(variables);
        response->body = body;
        Send(client, response);
    } else if (msg->GetCommand() == "next") {
        Next(_pausedContexts.at(0));
    } else if (msg->GetCommand() == "stepIn") {
        StepInto(_pausedContexts.at(0));
    } else if (msg->GetCommand() == "stepOut") {
        StepOut(_pausedContexts.at(0));
    } else {
        std::cout << "Unhandled message: " << msg->GetCommand() << std::endl;
        Send(client, new DebugAdapterProtocol::DefinedResponse<DebugAdapterProtocol::ResponseBody, empty>(msg->seq));
    }
}
void AngelscriptDebugger::Next(asIScriptContext* ctx) {
    _nextContext = ctx;
    _next = true;
    _nextDepth = ctx->GetCallstackSize();
    _storedVariableReferences.clear();
    while (!_pausedContexts.empty()) {
        auto* c = _pausedContexts[_pausedContexts.size() - 1];
        _pausedContexts.pop_back();
        std::thread([](asIScriptContext* c) { c->Execute(); }, c).detach();
    }
}

void AngelscriptDebugger::StepInto(asIScriptContext* ctx) {
    _nextContext = ctx;
    _stepInto = true;
    _nextDepth = ctx->GetCallstackSize();
    _storedVariableReferences.clear();
    while (!_pausedContexts.empty()) {
        auto* c = _pausedContexts[_pausedContexts.size() - 1];
        _pausedContexts.pop_back();
        std::thread([](asIScriptContext* c) { c->Execute(); }, c).detach();
    }
}
void AngelscriptDebugger::StepOut(asIScriptContext* ctx) {
    _nextContext = ctx;
    _stepOut = true;
    _nextDepth = ctx->GetCallstackSize();
    _storedVariableReferences.clear();
    while (!_pausedContexts.empty()) {
        auto* c = _pausedContexts[_pausedContexts.size() - 1];
        _pausedContexts.pop_back();
        std::thread([](asIScriptContext* c) { c->Execute(); }, c).detach();
    }
}