CgiHandler::CgiHandler(const HttpRequest &request, HttpResponse &response)
    : AHandler(request, response), cgiProcess_(nullptr), cgiStdIn_(nullptr), cgiStdOut_(nullptr)
{
    Log::debug("CgiHandler constructed");
    bodyWriteOffset_ = 0;
    headersParsed_ = false;
    expectedBody_.reset();
}

void CgiHandler::setCgiSockets(std::unique_ptr<CgiSocket> cgiStdIn, std::unique_ptr<CgiSocket> cgiStdOut,
                               std::unique_ptr<CgiSocket> cgiStdErr)
{
    cgiStdIn->setCallback([this]() { write(); });
    cgiStdOut->setCallback([this]() { read(); });
    cgiStdErr->setCallback([this]() { error(); });
    
    cgiStdOut_ = std::move(cgiStdOut);
    cgiStdIn_ = std::move(cgiStdIn);
    cgiStdErr_ = std::move(cgiStdErr);

    request_.getClient().addSocket(cgiStdIn_.get());
    request_.getClient().addSocket(cgiStdOut_.get());
    request_.getClient().addSocket(cgiStdErr_.get());
    // Ensure stdout/stderr are set to READ interest; stdin to WRITE only if there's a body
    cgiStdOut_->setIOState(ASocket::IoState::READ);
    cgiStdOut_->markDirty();
    cgiStdErr_->setIOState(ASocket::IoState::READ);
    cgiStdErr_->markDirty();
    if (!request_.getBody().empty())
    {
        cgiStdIn_->setIOState(ASocket::IoState::WRITE);
        cgiStdIn_->markDirty();
    }
    else
    {
        request_.getClient().removeSocket(cgiStdIn_.get());
        cgiStdIn_.reset();
    }
}

void CgiHandler::write()
{
    Log::trace(LOCATION);
    if (!cgiStdIn_) return;
    const std::string &body = request_.getBody();
    if (bodyWriteOffset_ >= body.size())
    {
        request_.getClient().removeSocket(cgiStdIn_.get());
        cgiStdIn_.reset();
        return;
    }
    // Stream body until pipe stops accepting data; no errno checks needed (level-triggered epoll)
    size_t before = bodyWriteOffset_;
    while (bodyWriteOffset_ < body.size())
    {
        const char *data = body.data() + bodyWriteOffset_;
        size_t remain = body.size() - bodyWriteOffset_;
        size_t chunk = remain > 32768 ? 32768 : remain;
        ssize_t n = cgiStdIn_->write(data, chunk);
        if (n > 0)
        {
            bodyWriteOffset_ += static_cast<size_t>(n);
        }
        else
        {
            break; // would block or peer closed; try again on next EPOLLOUT
        }
    }
    if (bodyWriteOffset_ >= body.size())
    {
        Log::debug("CGI stdin sent " + std::to_string(body.size()) + " bytes, closing write end");
        request_.getClient().removeSocket(cgiStdIn_.get());
        cgiStdIn_.reset();
    }
    else
    {
        Log::debug("CGI stdin progress " + std::to_string(before) + "→" + std::to_string(bodyWriteOffset_));
        cgiStdIn_->setIOState(ASocket::IoState::WRITE);
        cgiStdIn_->markDirty();
    }
}

static inline bool findHeaderEnd(const std::string &s, size_t &pos, long &sepSize)
{
    size_t a = s.find("\r\n\r\n");
    size_t b = s.find("\n\n");
    size_t c = s.find("\r\r");
    size_t end = std::min({a, b, c});
    if (end == std::string::npos) return false;
    sepSize = (end == a) ? 4 : 2;
    pos = end;
    return true;
}

void CgiHandler::read()
{
    Log::trace(LOCATION);
    if (!cgiStdOut_) return;
    // Drain as much as available this callback
    for (;;)
    {
        char buffer[bufferSize_] = {};
        ssize_t n = cgiStdOut_->read(buffer, sizeof(buffer));
        if (n > 0)
        {
            appendToBuffer(buffer, static_cast<size_t>(n));
            // Parse headers once, as soon as we have them
            if (!headersParsed_)
            {
                size_t headerEnd = 0;
                long sepSize = 0;
                std::string snapshot(buffer_.begin(), buffer_.end());
                if (findHeaderEnd(snapshot, headerEnd, sepSize))
                {
                    std::string headers(snapshot.begin(), snapshot.begin() + static_cast<long>(headerEnd));
                    parseCgiHeaders(headers);
                    // After headers parsed, remove them from buffer_ so it contains only body
                    buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<long>(headerEnd) + sepSize);
                    headersParsed_ = true;
                    // Capture expected body size if provided
                    std::string cl = response_.getHeaders().get("Content-Length");
                    if (!cl.empty())
                    {
                        expectedBody_ = static_cast<size_t>(std::strtoul(cl.c_str(), nullptr, 10));
                    }
                    else
                    {
                        expectedBody_.reset();
                    }
                }
            }
            // If we know Content-Length and already have it, finalize
            if (headersParsed_ && expectedBody_.has_value() && buffer_.size() >= expectedBody_.value())
            {
                request_.getClient().removeSocket(cgiStdOut_.get());
                cgiStdOut_.reset();
                finalizeCgiResponse();
                return;
            }
            continue; // try to read more this tick
        }
        else if (n == 0)
        {
            // EOF → finalize with whatever body we have
            Log::info("CGI process closed stdout, fd: " + std::to_string(cgiStdOut_->getFd()));
            request_.getClient().removeSocket(cgiStdOut_.get());
            cgiStdOut_.reset();
            // If headers not parsed yet, try once more (some scripts may end with only headers)
            if (!headersParsed_)
            {
                size_t headerEnd = 0;
                long sep = 0;
                std::string snap(buffer_.begin(), buffer_.end());
                if (findHeaderEnd(snap, headerEnd, sep))
                {
                    std::string headers(snap.begin(), snap.begin() + static_cast<long>(headerEnd));
                    parseCgiHeaders(headers);
                    buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<long>(headerEnd) + sep);
                    headersParsed_ = true;
                }
            }
            finalizeCgiResponse();
            return;
        }
        else
        {
            // Would block or transient error; wait for next EPOLLIN
            break;
        }
    }
}

void CgiHandler::error()
{
    Log::trace(LOCATION);
    if (!cgiStdErr_) return;
    for (;;)
    {
        char buffer[bufferSize_] = {};
        ssize_t n = cgiStdErr_->read(buffer, sizeof(buffer));
        if (n > 0)
        {
            appendToBuffer(buffer, static_cast<size_t>(n));
            Log::error("CGI stderr output (fd: " + std::to_string(cgiStdErr_->getFd())
                       + "): " + std::string(buffer, static_cast<size_t>(n)));
            continue;
        }
        else if (n == 0)
        {
            Log::info("CGI process closed stderr, fd: " + std::to_string(cgiStdErr_->getFd()));
            request_.getClient().removeSocket(cgiStdErr_.get());
            cgiStdErr_.reset();
            break;
        }
        else
        {
            break;
        }
    }
}

void CgiHandler::parseCgiOutput()
{
    // Changed: only parse headers if present; do NOT finalize here
    Log::trace(LOCATION);
    if (headersParsed_) return;
    size_t headerEnd = 0;
    long sepSize = 0;
    std::string header(buffer_.begin(), buffer_.end());
    if (!findHeaderEnd(header, headerEnd, sepSize))
    {
        Log::debug("CGI output headers not complete yet");
        return;
    }
    std::string headers(header.begin(), header.begin() + static_cast<long>(headerEnd));
    Log::debug("CGI output headers: " + headers);
    parseCgiHeaders(headers);
    buffer_.erase(buffer_.begin(), buffer_.begin() + static_cast<long>(headerEnd) + sepSize);
    headersParsed_ = true;
    std::string cl = response_.getHeaders().get("Content-Length");
    if (!cl.empty())
        expectedBody_ = static_cast<size_t>(std::strtoul(cl.c_str(), nullptr, 10));
    else
        expectedBody_.reset();
}

void CgiHandler::parseCgiHeaders(std::string &headers)
{
    Log::trace(LOCATION);
    // ...existing code...
    // At end, capture Content-Length if present
    std::string cl = response_.getHeaders().get("Content-Length");
    if (!cl.empty()) expectedBody_ = static_cast<size_t>(std::strtoul(cl.c_str(), nullptr, 10));
    // ...existing code...
}

void CgiHandler::finalizeCgiResponse()
{
    Log::trace(LOCATION);
    auto status = response_.getHeaders().get("Status");
    wait();
    if (cgiProcess_ && cgiProcess_->getExitCode() > 0 && status.empty())
    {
        response_.setStatus(500);
    }
    else if (!status.empty())
    {
        response_.setStatus(std::atoi(status.c_str()));
    }
    // Append only the body (headers already stripped)
    response_.appendBody(buffer_);
    response_.setComplete();
    buffer_.clear();
}