#include "StdioTransport.hpp"
#include <spdlog/sinks/stdout_color_sinks.h>

#include <sstream>
#include <cstdlib>
#include <chrono>
#include <thread>
#include <cerrno>
#include <cstring>

namespace celuna::mcp {

StdioTransport::StdioTransport(const MCPServerConfig& config) : m_config(config) {
    m_logger = spdlog::get("MCP:" + config.name);
    if (!m_logger) {
        m_logger = spdlog::stdout_color_mt("MCP:" + config.name);
    }
}

StdioTransport::~StdioTransport() {
    stop();
}

bool StdioTransport::start() {
    if (m_running) {
        return true;
    }

    m_logger->info("Starting MCP server: {} {}", m_config.command,
                   [&]() {
                       std::string args;
                       for (const auto& a : m_config.args) args += a + " ";
                       return args;
                   }());

    if (!spawnProcess()) {
        m_logger->error("Failed to spawn MCP server process");
        return false;
    }

    m_running = true;

    // Start reader thread
    m_readerThread = std::thread(&StdioTransport::readerLoop, this);

    m_logger->info("MCP server started successfully");
    return true;
}

void StdioTransport::stop() {
    if (!m_running) {
        return;
    }

    m_running = false;

    // Close pipes and terminate process first, so reader thread can exit
#ifdef _WIN32
    if (m_stdinWrite) {
        CloseHandle(m_stdinWrite);
        m_stdinWrite = nullptr;
    }
    if (m_processHandle) {
        TerminateProcess(m_processHandle, 0);
        CloseHandle(m_processHandle);
        m_processHandle = nullptr;
    }
    if (m_stdoutRead) {
        CloseHandle(m_stdoutRead);
        m_stdoutRead = nullptr;
    }
#else
    if (m_stdinFd >= 0) {
        close(m_stdinFd);
        m_stdinFd = -1;
    }
    if (m_pid > 0) {
        kill(m_pid, SIGTERM);
        // Don't wait here - let the reader thread see EOF first
    }
    if (m_stdoutFd >= 0) {
        close(m_stdoutFd);
        m_stdoutFd = -1;
    }
#endif

    // Now wait for reader thread (should exit quickly since stdout is closed)
    if (m_readerThread.joinable()) {
        m_readerThread.join();
    }

    // Final cleanup - wait for child process
#ifndef _WIN32
    if (m_pid > 0) {
        waitpid(m_pid, nullptr, 0);
        m_pid = -1;
    }
#endif

    m_logger->info("MCP server stopped");
}

bool StdioTransport::isRunning() const {
    return m_running;
}

bool StdioTransport::spawnProcess() {
#ifdef _WIN32
    // Windows implementation
    SECURITY_ATTRIBUTES sa;
    sa.nLength = sizeof(SECURITY_ATTRIBUTES);
    sa.bInheritHandle = TRUE;
    sa.lpSecurityDescriptor = nullptr;

    HANDLE stdinRead, stdoutWrite;

    // Create pipes
    if (!CreatePipe(&stdinRead, &m_stdinWrite, &sa, 0)) {
        return false;
    }
    if (!CreatePipe(&m_stdoutRead, &stdoutWrite, &sa, 0)) {
        CloseHandle(stdinRead);
        CloseHandle(m_stdinWrite);
        return false;
    }

    // Don't inherit our end of the pipes
    SetHandleInformation(m_stdinWrite, HANDLE_FLAG_INHERIT, 0);
    SetHandleInformation(m_stdoutRead, HANDLE_FLAG_INHERIT, 0);

    // Build command line
    std::string cmdLine = m_config.command;
    for (const auto& arg : m_config.args) {
        cmdLine += " " + arg;
    }

    // Setup environment
    std::string envBlock;
    for (const auto& [key, value] : m_config.env) {
        std::string resolvedValue = value;
        // Resolve ${VAR} references
        if (value.find("${") != std::string::npos) {
            size_t start = value.find("${");
            size_t end = value.find("}");
            if (end > start) {
                std::string envVar = value.substr(start + 2, end - start - 2);
                const char* envVal = std::getenv(envVar.c_str());
                if (envVal) {
                    resolvedValue = envVal;
                }
            }
        }
        envBlock += key + "=" + resolvedValue + '\0';
    }
    envBlock += '\0';

    STARTUPINFOA si = {};
    si.cb = sizeof(si);
    si.hStdInput = stdinRead;
    si.hStdOutput = stdoutWrite;
    si.hStdError = GetStdHandle(STD_ERROR_HANDLE);
    si.dwFlags |= STARTF_USESTDHANDLES;

    PROCESS_INFORMATION pi = {};

    BOOL success = CreateProcessA(
        nullptr,
        const_cast<char*>(cmdLine.c_str()),
        nullptr, nullptr,
        TRUE,
        0,
        envBlock.empty() ? nullptr : const_cast<char*>(envBlock.c_str()),
        nullptr,
        &si, &pi
    );

    CloseHandle(stdinRead);
    CloseHandle(stdoutWrite);

    if (!success) {
        CloseHandle(m_stdinWrite);
        CloseHandle(m_stdoutRead);
        return false;
    }

    m_processHandle = pi.hProcess;
    CloseHandle(pi.hThread);

    return true;

#else
    // Unix implementation
    int stdinPipe[2];
    int stdoutPipe[2];

    if (pipe(stdinPipe) < 0 || pipe(stdoutPipe) < 0) {
        return false;
    }

    m_pid = fork();

    if (m_pid < 0) {
        close(stdinPipe[0]);
        close(stdinPipe[1]);
        close(stdoutPipe[0]);
        close(stdoutPipe[1]);
        return false;
    }

    if (m_pid == 0) {
        // Child process
        close(stdinPipe[1]);  // Close write end of stdin
        close(stdoutPipe[0]); // Close read end of stdout

        dup2(stdinPipe[0], STDIN_FILENO);
        dup2(stdoutPipe[1], STDOUT_FILENO);

        close(stdinPipe[0]);
        close(stdoutPipe[1]);

        // Set environment variables
        for (const auto& [key, value] : m_config.env) {
            std::string resolvedValue = value;
            if (value.find("${") != std::string::npos) {
                size_t start = value.find("${");
                size_t end = value.find("}");
                if (end > start) {
                    std::string envVar = value.substr(start + 2, end - start - 2);
                    const char* envVal = std::getenv(envVar.c_str());
                    if (envVal) {
                        resolvedValue = envVal;
                    }
                }
            }
            setenv(key.c_str(), resolvedValue.c_str(), 1);
        }

        // Build argv
        std::vector<char*> argv;
        argv.push_back(const_cast<char*>(m_config.command.c_str()));
        for (const auto& arg : m_config.args) {
            argv.push_back(const_cast<char*>(arg.c_str()));
        }
        argv.push_back(nullptr);

        execvp(m_config.command.c_str(), argv.data());
        _exit(1);
    }

    // Parent process
    close(stdinPipe[0]);  // Close read end of stdin
    close(stdoutPipe[1]); // Close write end of stdout

    m_stdinFd = stdinPipe[1];
    m_stdoutFd = stdoutPipe[0];

    // Give the child a moment to fail if command doesn't exist
    // 100ms should be enough for execvp to fail on invalid command
    std::this_thread::sleep_for(std::chrono::milliseconds(100));

    // Check if child is still alive
    int status;
    pid_t result = waitpid(m_pid, &status, WNOHANG);
    if (result == m_pid) {
        // Child has exited - command likely failed (exit code 1 from _exit)
        m_logger->error("Child process exited immediately (command not found?)");
        close(m_stdinFd);
        close(m_stdoutFd);
        m_stdinFd = -1;
        m_stdoutFd = -1;
        m_pid = -1;
        return false;
    } else if (result == -1) {
        // Error in waitpid
        m_logger->error("waitpid error: {}", strerror(errno));
    }

    return true;
#endif
}

void StdioTransport::readerLoop() {
    m_logger->debug("Reader thread started");

    while (m_running) {
        std::string line = readLine();
        if (line.empty()) {
            if (m_running) {
                m_logger->warn("EOF from MCP server");
            }
            break;
        }

        try {
            json message = json::parse(line);
            handleMessage(message);
        } catch (const json::exception& e) {
            m_logger->error("Failed to parse MCP message: {}", e.what());
        }
    }

    m_logger->debug("Reader thread stopped");
}

void StdioTransport::handleMessage(const json& message) {
    // Check if this is a response (id must be a number, not null)
    if (message.contains("id") && message["id"].is_number() &&
        (message.contains("result") || message.contains("error"))) {
        int id = message["id"].get<int>();

        std::lock_guard<std::mutex> lock(m_mutex);
        auto it = m_pendingRequests.find(id);
        if (it != m_pendingRequests.end()) {
            auto pending = it->second;
            {
                std::lock_guard<std::mutex> reqLock(pending->mutex);
                pending->response = JsonRpcResponse::fromJson(message);
                pending->hasResponse = true;
            }
            pending->cv.notify_one();
        }
    }
    // TODO: Handle notifications from server
}

JsonRpcResponse StdioTransport::sendRequest(const JsonRpcRequest& request, int timeoutMs) {
    if (!m_running) {
        JsonRpcResponse error;
        error.error = json{{"code", -1}, {"message", "Transport not running"}};
        return error;
    }

    // Create a mutable copy, use provided ID if > 0, otherwise generate one
    JsonRpcRequest req = request;
    if (req.id <= 0) {
        req.id = m_nextRequestId++;
    }

    // Create pending request
    auto pending = std::make_shared<PendingRequest>();
    {
        std::lock_guard<std::mutex> lock(m_mutex);
        m_pendingRequests[req.id] = pending;
    }

    // Send request
    std::string line = req.toJson().dump();
    m_logger->debug("Sending: {}", line.substr(0, 200));
    writeLine(line);

    // Wait for response
    std::unique_lock<std::mutex> lock(pending->mutex);
    bool received = pending->cv.wait_for(lock, std::chrono::milliseconds(timeoutMs), [&] {
        return pending->hasResponse;
    });

    // Cleanup
    {
        std::lock_guard<std::mutex> globalLock(m_mutex);
        m_pendingRequests.erase(req.id);
    }

    if (!received) {
        m_logger->warn("Request {} timed out", req.method);
        JsonRpcResponse error;
        error.id = req.id;
        error.error = json{{"code", -2}, {"message", "Request timed out"}};
        return error;
    }

    return pending->response;
}

void StdioTransport::sendNotification(const std::string& method, const json& params) {
    if (!m_running) return;

    json notification = {
        {"jsonrpc", "2.0"},
        {"method", method}
    };
    if (!params.is_null()) {
        notification["params"] = params;
    }

    writeLine(notification.dump());
}

void StdioTransport::writeLine(const std::string& line) {
    std::string data = line + "\n";

#ifdef _WIN32
    DWORD written;
    WriteFile(m_stdinWrite, data.c_str(), static_cast<DWORD>(data.size()), &written, nullptr);
#else
    ssize_t result = write(m_stdinFd, data.c_str(), data.size());
    (void)result;  // Ignore return value for now
#endif
}

std::string StdioTransport::readLine() {
    std::string line;
    char c;

#ifdef _WIN32
    DWORD bytesRead;
    while (m_running && ReadFile(m_stdoutRead, &c, 1, &bytesRead, nullptr) && bytesRead > 0) {
        if (c == '\n') break;
        if (c != '\r') line += c;
    }
#else
    while (m_running) {
        ssize_t n = read(m_stdoutFd, &c, 1);
        if (n <= 0) break;
        if (c == '\n') break;
        if (c != '\r') line += c;
    }
#endif

    if (!line.empty()) {
        m_logger->debug("Received: {}", line.substr(0, 200));
    }

    return line;
}

} // namespace celuna::mcp
