#include <atomic>
#include <csignal>
#include <cstdint>
#include <errno.h>
#include <iostream>
#include <memory>
#include <thread>
#include <unordered_map>
#include <vector>

#include <arpa/inet.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <sys/mman.h>
#include <sys/socket.h>
#include <unistd.h>

#include <turbojpeg.h>
#include <libcamera/libcamera.h>
#include <libcamera/control_ids.h>
#include <libcamera/formats.h>

#include "event_loop.hpp"

using namespace libcamera;

// --- knobs -------------------------------------------------------------------
// Goal: throughput >= latency >= quality. Tune here and rebuild.

constexpr int   kPort         = 5000;
constexpr int   kWidth        = 640;
constexpr int   kHeight       = 480;
constexpr int   kTargetFps    = 30;          // 0 = leave to camera defaults
constexpr int   kBufferCount  = 4;           // pipeline depth
constexpr int   kJpegQuality  = 80;
constexpr int   kJpegSubsamp  = TJSAMP_420;
constexpr int   kSendBufBytes = 1 << 20;     // 1 MiB SO_SNDBUF (throughput)

// Autofocus: true → continuous AF; false → fixed lens at kLensDiopters
//   diopters: 0 = infinity, ~10 = closest. Silently ignored on cameras
//   without AF (Pi cam v1/v2). Pi cam v3 supports it.
constexpr bool  kAutoFocus    = true;
constexpr float kLensDiopters = 0.0f;

// Pi viewfinder native format: XRGB8888 — DRM byte order, so memory layout
// is B,G,R,X. Feed turbojpeg TJPF_BGRX so it walks 4 B/pixel and ignores X.
// No software color conversion, no row copy.
constexpr int   kTjPixelFmt   = TJPF_BGRX;

// --- state -------------------------------------------------------------------

namespace {

std::shared_ptr<Camera> camera;
EventLoop loop;

std::atomic<int> client_fd{-1};

struct Mapped { uint8_t *data; size_t size; };
std::unordered_map<FrameBuffer *, Mapped> mapped;

tjhandle      tj_handle = nullptr;
uint8_t      *jpeg_buf  = nullptr;
unsigned long jpeg_cap  = 0;

// --- net ---------------------------------------------------------------------

bool sendAll(int fd, const uint8_t *p, size_t n)
{
    while (n) {
        ssize_t r = ::send(fd, p, n, MSG_NOSIGNAL);
        if (r <= 0) return false;
        p += r;
        n -= r;
    }
    return true;
}

void closeClient()
{
    int fd = client_fd.exchange(-1);
    if (fd >= 0) ::close(fd);
}

void serverThread()
{
    int srv = ::socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, 0);
    int yes = 1;
    ::setsockopt(srv, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof(yes));

    sockaddr_in addr{};
    addr.sin_family      = AF_INET;
    addr.sin_port        = htons(kPort);
    addr.sin_addr.s_addr = INADDR_ANY;

    if (::bind(srv, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)) < 0
     || ::listen(srv, 4) < 0) {
        perror("bind/listen");
        ::close(srv);
        return;
    }

    std::cerr << "tcp://0.0.0.0:" << kPort << " waiting...\n";

    for (;;) {
        int c = ::accept(srv, nullptr, nullptr);
        if (c < 0) {
            if (errno == EINTR) continue;
            break;
        }
        int one = 1, sndbuf = kSendBufBytes;
        ::setsockopt(c, IPPROTO_TCP, TCP_NODELAY, &one,    sizeof(one));
        ::setsockopt(c, SOL_SOCKET,  SO_SNDBUF,   &sndbuf, sizeof(sndbuf));
        int old = client_fd.exchange(c);
        if (old >= 0) ::close(old);
        std::cerr << "client connected\n";
    }
}

// --- camera ------------------------------------------------------------------

void processRequest(Request *request);

void requestComplete(Request *request)
{
    if (request->status() == Request::RequestCancelled)
        return;
    loop.callLater([request] { processRequest(request); });
}

void processRequest(Request *request)
{
    auto it = request->buffers().begin();
    const Stream *stream = it->first;
    FrameBuffer  *buffer = it->second;

    int fd = client_fd.load(std::memory_order_acquire);
    if (fd >= 0 && buffer->metadata().status == FrameMetadata::FrameSuccess) {
        const StreamConfiguration &cfg = stream->configuration();
        const Mapped &m = mapped[buffer];

        unsigned long size = jpeg_cap;
        if (tjCompress2(tj_handle,
                        m.data, cfg.size.width, cfg.stride, cfg.size.height,
                        kTjPixelFmt,
                        &jpeg_buf, &size,
                        kJpegSubsamp, kJpegQuality,
                        TJFLAG_FASTDCT | TJFLAG_NOREALLOC) == 0) {
            uint32_t len = htonl(static_cast<uint32_t>(size));
            if (!sendAll(fd, reinterpret_cast<uint8_t *>(&len), sizeof(len))
             || !sendAll(fd, jpeg_buf, size)) {
                std::cerr << "client gone\n";
                closeClient();
            }
        } else {
            std::cerr << "tjCompress2: " << tjGetErrorStr2(tj_handle) << "\n";
        }
    }

    request->reuse(Request::ReuseBuffers);
    camera->queueRequest(request);
}

} // namespace

// --- main --------------------------------------------------------------------

int main()
{
    std::signal(SIGPIPE, SIG_IGN);
    std::thread(serverThread).detach();

    auto cm = std::make_unique<CameraManager>();
    cm->start();
    if (cm->cameras().empty()) {
        std::cerr << "no cameras\n";
        cm->stop();
        return 1;
    }

    camera = cm->get(cm->cameras()[0]->id());
    camera->acquire();

    auto config = camera->generateConfiguration({ StreamRole::Viewfinder });
    StreamConfiguration &sc = config->at(0);
    sc.pixelFormat = formats::XRGB8888;
    sc.size        = { kWidth, kHeight };
    sc.bufferCount = kBufferCount;

    if (config->validate() == CameraConfiguration::Invalid
     || sc.pixelFormat != formats::XRGB8888) {
        std::cerr << "got " << sc.pixelFormat.toString()
                  << ", need XRGB8888\n";
        return 1;
    }

    std::cerr << "stream "   << sc.size.toString()
              << " "         << sc.pixelFormat.toString()
              << " stride "  << sc.stride
              << " buffers " << sc.bufferCount << "\n";

    if (camera->configure(config.get()) < 0) {
        std::cerr << "configure failed\n";
        return 1;
    }

    auto *allocator = new FrameBufferAllocator(camera);
    Stream *stream = sc.stream();
    if (allocator->allocate(stream) < 0) {
        std::cerr << "alloc failed\n";
        return 1;
    }

    for (const std::unique_ptr<FrameBuffer> &b : allocator->buffers(stream)) {
        const FrameBuffer::Plane &p = b->planes()[0];
        void *m = mmap(nullptr, p.length, PROT_READ, MAP_SHARED,
                       p.fd.get(), p.offset);
        if (m == MAP_FAILED) { perror("mmap"); return 1; }
        mapped[b.get()] = { static_cast<uint8_t *>(m), p.length };
    }

    tj_handle = tjInitCompress();
    jpeg_cap  = tjBufSize(sc.size.width, sc.size.height, kJpegSubsamp);
    jpeg_buf  = tjAlloc(jpeg_cap);

    std::vector<std::unique_ptr<Request>> requests;
    for (const std::unique_ptr<FrameBuffer> &b : allocator->buffers(stream)) {
        std::unique_ptr<Request> r = camera->createRequest();
        if (!r || r->addBuffer(stream, b.get()) < 0) {
            std::cerr << "request setup failed\n";
            return 1;
        }
        requests.push_back(std::move(r));
    }

    // Initial controls: AF + frame-rate cap.
    ControlList ctrls(camera->controls());
    const ControlInfoMap &caps = camera->controls();

    if (caps.count(&controls::AfMode)) {
        if (kAutoFocus) {
            ctrls.set(controls::AfMode, controls::AfModeContinuous);
        } else {
            ctrls.set(controls::AfMode,     controls::AfModeManual);
            ctrls.set(controls::LensPosition, kLensDiopters);
        }
    }
    if (kTargetFps > 0) {
        int64_t dur = 1000000 / kTargetFps;
        ctrls.set(controls::FrameDurationLimits,
                  Span<const int64_t, 2>({ dur, dur }));
    }

    camera->requestCompleted.connect(requestComplete);
    if (camera->start(&ctrls) < 0) {
        std::cerr << "camera start failed\n";
        return 1;
    }
    for (std::unique_ptr<Request> &r : requests)
        camera->queueRequest(r.get());

    std::cerr << "streaming, Ctrl+C to stop\n";
    loop.exec();

    camera->stop();
    camera->requestCompleted.disconnect();
    for (auto &kv : mapped) munmap(kv.second.data, kv.second.size);
    allocator->free(stream);
    delete allocator;
    camera->release();
    camera.reset();
    cm->stop();
    if (jpeg_buf)  tjFree(jpeg_buf);
    if (tj_handle) tjDestroy(tj_handle);
    closeClient();
    return 0;
}