#include "detail/transport_stdio.hpp"

#include <mw/exec.hpp>

namespace agent_smith::mcp::internal
{
namespace
{
mw::Error failure(std::string message)
{
    return McpError{std::move(message), {}, std::nullopt, nullptr, true};
}

class StdioTransport final : public Transport
{
public:
    StdioTransport(Runtime& runtime, mw::Pipe input, mw::Pipe output,
                   mw::Process process)
        : runtime_(runtime), input_(std::move(input)),
          output_(std::move(output)), process_(std::move(process))
    {
    }

    mw::Task<mw::E<void>> send(nlohmann::json frame,
                               std::chrono::steady_clock::time_point deadline,
                               CancellationToken cancel) override
    {
        std::string bytes = frame.dump() + "\n";
        std::function<mw::E<void>()> work =
            [this, bytes = std::move(bytes), deadline, cancel]
        {
            std::unique_lock lock(write_mutex_, std::defer_lock);
            while(!lock.try_lock_for(std::chrono::milliseconds(2)))
            {
                if(cancel.cancelled() || interrupted_ ||
                   std::chrono::steady_clock::now() >= deadline)
                {
                    return mw::E<void>(std::unexpected(failure(
                        "stdio write stopped before lock acquisition")));
                }
            }
            if(closed_)
            {
                return mw::E<void>(std::unexpected(failure("stdio closed")));
            }
            std::size_t offset = 0;
            while(offset < bytes.size())
            {
                const auto now = std::chrono::steady_clock::now();
                if(cancel.cancelled() || interrupted_ || now >= deadline)
                {
                    return mw::E<void>(std::unexpected(failure(
                        "stdio write stopped (request may be partial)")));
                }
                auto result = input_.writeUntil(
                    std::string_view(bytes).substr(offset),
                    std::min(deadline, now + std::chrono::milliseconds(2)));
                if(!result)
                {
                    return mw::E<void>(std::unexpected(failure(
                        "stdio write failed (request may be partial)")));
                }
                offset += result->bytes_written;
            }
            return mw::E<void>{};
        };
        co_return co_await runtime_.runOnWorker(std::move(work));
    }

    mw::Task<mw::E<nlohmann::json>> recv() override
    {
        while(true)
        {
            if(interrupted_)
            {
                co_return std::unexpected(failure("stdio interrupted"));
            }
            const auto newline = buffer_.find('\n');
            if(newline != std::string::npos)
            {
                auto frame = nlohmann::json::parse(buffer_.substr(0, newline),
                                                   nullptr, false);
                buffer_.erase(0, newline + 1);
                if(frame.is_discarded())
                {
                    co_return std::unexpected(
                        McpError{"invalid stdio JSON frame"});
                }
                co_return frame;
            }
            std::function<mw::E<mw::PipeReadResult>()> work = [this]
            {
                std::lock_guard lock(read_mutex_);
                if(closed_)
                {
                    return mw::E<mw::PipeReadResult>(
                        std::unexpected(failure("stdio closed")));
                }
                return output_.tryRead(8192);
            };
            auto bytes = co_await runtime_.runOnWorker(std::move(work));
            if(!bytes)
            {
                co_return std::unexpected(failure(bytes.error().msg()));
            }
            if(bytes->state == mw::PipeState::EOF_REACHED)
            {
                co_return std::unexpected(failure("MCP server closed stdout"));
            }
            buffer_ += bytes->data;
            if(buffer_.size() > 4 * 1024 * 1024)
            {
                co_return std::unexpected(McpError{"MCP frame too large"});
            }
            if(bytes->state == mw::PipeState::WOULD_BLOCK)
            {
                co_await runtime_.sleepFor(std::chrono::milliseconds(2));
            }
        }
    }

    mw::Task<mw::E<void>> close() override
    {
        interrupt();
        std::function<mw::E<void>()> work = [this]
        {
            std::scoped_lock lock(read_mutex_, write_mutex_);
            if(closed_)
            {
                return mw::E<void>{};
            }
            closed_ = true;
            std::ignore = input_.closeInput();
            std::ignore = output_.closeOutput();
            auto exited = process_.waitFor(std::chrono::milliseconds(50));
            if(!exited || exited->state != mw::ProcessWaitState::COMPLETE)
            {
                exited = process_.shutdown(std::chrono::milliseconds(50),
                                           std::chrono::steady_clock::now() +
                                               std::chrono::milliseconds(500));
            }
            if(!exited)
            {
                return mw::E<void>(
                    std::unexpected(failure(exited.error().msg())));
            }
            if(exited->state != mw::ProcessWaitState::COMPLETE)
            {
                return mw::E<void>(std::unexpected(
                    failure("MCP child shutdown deadline expired")));
            }
            return mw::E<void>{};
        };
        co_return co_await runtime_.runOnWorker(std::move(work));
    }

    void interrupt() noexcept override
    {
        interrupted_ = true;
    }

private:
    Runtime& runtime_;
    mw::Pipe input_;
    mw::Pipe output_;
    mw::Process process_;
    std::mutex read_mutex_;
    std::timed_mutex write_mutex_;
    bool closed_ = false;
    std::atomic<bool> interrupted_{false};
    std::string buffer_;
};
} // namespace

mw::E<std::unique_ptr<Transport>> openStdio(Runtime& runtime,
                                            const StdioServerConfig& config)
{
    auto input = mw::Pipe::create();
    auto output = mw::Pipe::create();
    if(!input)
    {
        return std::unexpected(failure(input.error().msg()));
    }
    if(!output)
    {
        return std::unexpected(failure(output.error().msg()));
    }
    std::vector<std::string_view> argv;
    for(const auto& arg : config.argv)
    {
        argv.push_back(arg);
    }
    auto process = mw::Process::execArgs(&*input, argv, &*output);
    if(!process)
    {
        return std::unexpected(failure(process.error().msg()));
    }
    return std::unique_ptr<Transport>(new StdioTransport(
        runtime, std::move(*input), std::move(*output), std::move(*process)));
}
} // namespace agent_smith::mcp::internal