#include "screencast/app/pipeline.h"

#include "state_store.h"

#include "screencast/network/discovery.h"
#include "screencast/network/h264_packetizer.h"
#include "screencast/network/signaling.h"

#include <unistd.h>

#include <unistd.h>

#include <algorithm>
#include <array>
#include <chrono>
#include <condition_variable>
#include <deque>
#include <iostream>
#include <memory>
#include <mutex>
#include <thread>

namespace sc {
namespace {

std::string receiver_service_name() {
    std::array<char, 256> hostname{};
    if (::gethostname(hostname.data(), hostname.size()) != 0 || hostname[0] == '\0') {
        return "Screencast receiver";
    }
    return std::string{"Screencast receiver on "} + hostname.data();
}

} // namespace

#ifdef SC_HAS_SENDER

class SenderPipeline::Impl {
  public:
    explicit Impl(SenderPipelineConfig config) : config_(std::move(config)) {}

    ~Impl() {
        stop();
    }

    bool start() {
        auto capture_result = CaptureFactory::create(config_.capture_target);
        if (is_capture_error(capture_result)) {
            std::cerr << std::format("screencast: capture failed: {}\n", capture_error(capture_result).message);
            return false;
        }
        capture_ = std::move(capture_value(capture_result));

        if (!transport_->start(config_.local_rtp_endpoint, [](RtpPacket) {})) {
            std::cerr << "screencast: failed to start the RTP transport\n";
            return false;
        }
        transport_->set_peer(config_.peer_rtp_endpoint);

        run_thread_ = std::jthread([this](std::stop_token stop_token) { run(std::move(stop_token)); });

        // Publish the session for status widgets and one-click restarts.
        // The RTP peer is what a restart needs; the session id identifies it.
        write_sender_state(SenderState{
            .session_id = config_.session_id,
            .receiver = std::format("{}:{}", config_.peer_rtp_endpoint.address, config_.peer_rtp_endpoint.port),
            .bitrate_kbps = config_.encoder.bitrate_kbps,
            .pid = ::getpid(),
            .started_epoch_ms = 0});
        // A restart targets the receiver's signaling endpoint; the RTP
        // endpoint is re-negotiated from it.
        std::string restart_peer = std::format("{}:5005", config_.peer_rtp_endpoint.address);
        if (config_.signaling_server.has_value()) {
            restart_peer = std::format("{}:{}", config_.signaling_server->address, config_.signaling_server->port);
        }
        write_last_session(LastSession{.peer = restart_peer, .bitrate_kbps = config_.encoder.bitrate_kbps});
        return true;
    }

    void stop() {
        remove_sender_state();
        if (capture_ != nullptr) {
            capture_->stop();
        }
        run_thread_ = std::jthread{};
        transport_->stop();
    }

    void request_keyframe() {
        keyframe_requested_.store(true);
    }

  private:
    void run(std::stop_token stop_token) {
        while (!stop_token.stop_requested()) {
            if (keyframe_requested_.exchange(false) && encoder_ != nullptr) {
                encoder_->request_keyframe();
            }

            const std::optional<CapturedFrame> frame = capture_->next_frame();
            if (!frame.has_value()) {
                break;
            }

            if (encoder_ == nullptr && !create_encoder(*frame)) {
                break;
            }

            auto encoded_result = encoder_->encode(*frame);
            if (is_codec_error(encoded_result)) {
                std::cerr << std::format("screencast: encode failed: {}\n", codec_error(encoded_result).message);
                break;
            }
            for (const EncodedFrame& encoded : codec_value(encoded_result)) {
                for (RtpPacket& packet : packetizer_.packetize(encoded)) {
                    (void)transport_->send(packet);
                }
            }
        }
    }

    bool create_encoder(const CapturedFrame& frame) {
        EncoderConfig config = config_.encoder;
        config.width = frame.width;
        config.height = frame.height;

        // Downscale to the receiver's display when the capture is larger,
        // preserving aspect ratio and rounding to even values (YUV420P
        // requires even dimensions for the chroma planes). This avoids
        // encoding pixels the display cannot show.
        if (config_.max_encode_width > 0 && config_.max_encode_height > 0 &&
            (frame.width > config_.max_encode_width || frame.height > config_.max_encode_height)) {
            const double scale = std::min(static_cast<double>(config_.max_encode_width) / frame.width,
                                          static_cast<double>(config_.max_encode_height) / frame.height);
            config.width = std::max(2, static_cast<int>(frame.width * scale) & ~1);
            config.height = std::max(2, static_cast<int>(frame.height * scale) & ~1);
            std::cerr << std::format("screencast: downscaling {}x{} to {}x{} for the receiver's display\n",
                                     frame.width,
                                     frame.height,
                                     config.width,
                                     config.height);
        }

        auto encoder_result = EncoderFactory::create(config);
        if (is_codec_error(encoder_result)) {
            std::cerr << std::format("screencast: encoder creation failed: {}\n", codec_error(encoder_result).message);
            return false;
        }
        encoder_ = std::move(codec_value(encoder_result));
        return true;
    }

    SenderPipelineConfig config_;
    std::unique_ptr<CaptureSession> capture_;
    std::unique_ptr<Encoder> encoder_;
    H264Packetizer packetizer_;
    std::unique_ptr<RtpTransport> transport_ = RtpTransportFactory::create();
    std::jthread run_thread_;
    std::atomic<bool> keyframe_requested_{false};
};

void SenderPipeline::request_keyframe() {
    impl_->request_keyframe();
}

SenderPipeline::SenderPipeline(SenderPipelineConfig config) : impl_(std::make_unique<Impl>(std::move(config))) {}

SenderPipeline::~SenderPipeline() = default;

bool SenderPipeline::start() {
    return impl_->start();
}

void SenderPipeline::stop() {
    impl_->stop();
}

#endif // SC_HAS_SENDER

class ReceiverPipeline::Impl {
  public:
    explicit Impl(ReceiverPipelineConfig config) : config_(std::move(config)) {}

    ~Impl() {
        stop();
    }

    bool start() {
        auto decoder_result = DecoderFactory::create(config_.decoder);
        if (is_codec_error(decoder_result)) {
            std::cerr << std::format("screencast: decoder creation failed: {}\n", codec_error(decoder_result).message);
            return false;
        }
        decoder_ = std::move(codec_value(decoder_result));

        if (!transport_->start(config_.local_rtp_endpoint,
                               [this](RtpPacket packet) { on_packet(std::move(packet)); })) {
            std::cerr << "screencast: failed to start the RTP transport\n";
            return false;
        }

        // Signaling server: answer session offers with our RTP port so a
        // sender can find the media endpoint without configuration.
        auto signaling_result = SignalingFactory::create_server(config_.signaling_port);
        if (is_network_error(signaling_result)) {
            std::cerr << std::format("screencast: {}\n", network_error(signaling_result).message);
            transport_->stop();
            return false;
        }
        signaling_ = std::move(network_value(signaling_result));
        signaling_->on_message([this](const SignalingMessage& message) { handle_signaling(message); });

        // mDNS announcement; best effort, since senders can still use --peer.
        auto discovery_result = DiscoveryFactory::create_avahi();
        if (is_network_error(discovery_result)) {
            std::cerr << std::format("screencast: discovery unavailable: {}\n",
                                     network_error(discovery_result).message);
        } else {
            discovery_ = std::move(network_value(discovery_result));
            if (!discovery_->announce(receiver_service_name(), config_.signaling_port)) {
                std::cerr << std::format("screencast: announcing the receiver failed: {}\n", discovery_->last_error());
            }
        }

        // Every SDL call — window creation, event pumping, presenting, and
        // destruction — happens on the render thread. The Wayland backend
        // does not tolerate cross-thread windows: created on another thread,
        // the surface never maps and no window appears.
        std::mutex init_mutex;
        std::condition_variable init_done_cv;
        bool init_done = false;
        bool init_ok = false;

        render_thread_ =
            std::jthread([this, &init_mutex, &init_done_cv, &init_done, &init_ok](std::stop_token stop_token) {
                auto renderer_result = RendererFactory::create(config_.renderer);
                if (is_renderer_error(renderer_result)) {
                    std::cerr << std::format("screencast: renderer creation failed: {}\n",
                                             renderer_error(renderer_result).message);
                } else {
                    renderer_ = std::move(renderer_value(renderer_result));
                }
                {
                    std::lock_guard lock(init_mutex);
                    init_ok = !is_renderer_error(renderer_result);
                    init_done = true;
                }
                init_done_cv.notify_all();
                if (!init_ok) {
                    return;
                }
                render_loop(std::move(stop_token));
                if (renderer_ != nullptr) {
                    renderer_->shutdown();
                    renderer_ = nullptr;
                }
            });

        std::unique_lock lock(init_mutex);
        if (!init_done_cv.wait_for(lock, std::chrono::seconds(10), [&init_done] { return init_done; })) {
            if (discovery_ != nullptr) {
                discovery_->stop();
            }
            signaling_ = nullptr;
            transport_->stop();
            return false;
        }
        if (!init_ok) {
            render_thread_ = std::jthread{};
            if (discovery_ != nullptr) {
                discovery_->stop();
            }
            signaling_ = nullptr;
            transport_->stop();
            return false;
        }
        return true;
    }

    void stop() {
        render_thread_ = std::jthread{};
        // Join the signaling threads before dropping the channel so no
        // callback races the destruction.
        if (signaling_ != nullptr) {
            signaling_->disconnect();
        }
        if (discovery_ != nullptr) {
            discovery_->stop();
        }
        signaling_ = nullptr;
        discovery_ = nullptr;
        transport_->stop();
        jitter_.clear();
        renderer_ = nullptr;
        decoder_ = nullptr;
        {
            std::lock_guard lock(queue_mutex_);
            queue_.clear();
        }
    }

  private:
    void on_packet(RtpPacket packet) {
        // Absorb reordering (Wi-Fi) before the in-order depacketizer, so a
        // late packet is not misread as loss.
        for (RtpPacket& ordered : jitter_.push(std::move(packet))) {
            deliver_packet(ordered);
        }
    }

    void deliver_packet(const RtpPacket& packet) {
        const DepacketizeResult result = depacketizer_.depacketize(packet);
        if (result.frame_dropped) {
            maybe_send_pli();
        }
        if (!result.access_unit.has_value()) {
            return;
        }

        EncodedFrame encoded;
        encoded.data = std::move(*result.access_unit);
        encoded.rtp_timestamp = packet.header.timestamp;
        encoded.is_keyframe = false;

        auto decoded_result = decoder_->decode(encoded);
        if (is_codec_error(decoded_result)) {
            std::cerr << std::format("screencast: decode failed: {}\n", codec_error(decoded_result).message);
            return;
        }
        for (DecodedFrame& decoded : codec_value(decoded_result)) {
            if (!stream_started_) {
                stream_started_ = true;
                std::cerr << std::format("screencast: stream started ({}x{})\n", decoded.width, decoded.height);
            }
            std::lock_guard lock(queue_mutex_);
            // Keep latency low: drop the oldest frame when the queue is full.
            if (queue_.size() >= kMaxQueuedFrames) {
                queue_.pop_front();
            }
            queue_.push_back(std::move(decoded));
        }
    }

    void handle_signaling(const SignalingMessage& message) {
        const SessionOffer* offer = std::get_if<SessionOffer>(&message);
        if (offer == nullptr) {
            return;
        }
        session_id_ = offer->session_id;
        SessionAnswer answer;
        answer.session_id = offer->session_id;
        // Empty address: the sender targets the address of its signaling
        // connection, which reaches this RTP port.
        answer.rtp_endpoint = Endpoint{"", config_.local_rtp_endpoint.port};
        // Tell the sender what display it is rendering to so it can
        // downscale instead of encoding pixels the display cannot show.
        if (renderer_ != nullptr) {
            answer.display_width = renderer_->display_width();
            answer.display_height = renderer_->display_height();
        }
        signaling_->send(answer);
    }

    // Ask the sender for a keyframe after a damaged frame, rate-limited so
    // sustained loss cannot flood the signaling channel.
    void maybe_send_pli() {
        if (signaling_ == nullptr || session_id_.empty()) {
            return;
        }
        const auto now = std::chrono::steady_clock::now();
        if (now - last_pli_time_ < kPliMinInterval) {
            return;
        }
        last_pli_time_ = now;
        SessionPli pli;
        pli.session_id = session_id_;
        signaling_->send(pli);
        std::cerr << std::format("screencast: frame damaged; requesting a keyframe\n");
    }

    void render_loop(std::stop_token stop_token) {
        while (!stop_token.stop_requested()) {
            if (!renderer_->poll_events()) {
                break; // window closed
            }
            std::optional<DecodedFrame> frame;
            {
                std::unique_lock lock(queue_mutex_);
                if (!queue_.empty()) {
                    frame = std::move(queue_.front());
                    queue_.pop_front();
                }
            }
            if (frame.has_value()) {
                if (!renderer_->present(*frame) && !present_error_logged_) {
                    present_error_logged_ = true;
                    std::cerr << "screencast: rendering a decoded frame failed\n";
                }
            } else {
                std::this_thread::sleep_for(std::chrono::milliseconds(5));
            }
        }
    }

    static constexpr std::size_t kMaxQueuedFrames = 3;
    static constexpr std::chrono::milliseconds kPliMinInterval{500};

    ReceiverPipelineConfig config_;
    std::unique_ptr<Renderer> renderer_;
    std::unique_ptr<Decoder> decoder_;
    H264Depacketizer depacketizer_;
    RtpJitterBuffer jitter_;
    std::unique_ptr<RtpTransport> transport_ = RtpTransportFactory::create();
    std::unique_ptr<SignalingChannel> signaling_;
    std::unique_ptr<DiscoveryService> discovery_;
    std::jthread render_thread_;
    std::mutex queue_mutex_;
    std::deque<DecodedFrame> queue_;
    std::string session_id_;
    std::chrono::steady_clock::time_point last_pli_time_{};
    bool stream_started_ = false;
    bool present_error_logged_ = false;
};

ReceiverPipeline::ReceiverPipeline(ReceiverPipelineConfig config) : impl_(std::make_unique<Impl>(std::move(config))) {}

ReceiverPipeline::~ReceiverPipeline() = default;

bool ReceiverPipeline::start() {
    return impl_->start();
}

void ReceiverPipeline::stop() {
    impl_->stop();
}

} // namespace sc