#include <arpa/inet.h>
#include <chrono>
#include <cstring>
#include <filesystem>
#include <sys/select.h>
#include <fstream>
#include <iostream>
#include <netinet/in.h>
#include <pthread.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include <vector>

#ifdef __SWITCH__
#include <client.hpp>
#include <switch.h>
#endif

#include <protocol.hpp>
#include <TransferState.hpp>
#include <net/Socket.hpp>

namespace fs = std::filesystem;
using path = fs::path;

static TransferState g_client_state;

bool isClientTransferDone()        { return g_client_state.done.load(); }
bool isClientTransferCancelled()   { return g_client_state.cancelled.load(); }
bool isClientConnectionFailed()    { return g_client_state.connection_failed.load(); }
void cancelClientTransfer()        { g_client_state.cancelled.store(true); }
double getClientProgress()         { return g_client_state.progress(); }
std::string getClientStatusText()  { return g_client_state.getStatus(); }

static bool send_all(int sock, const void* buf, size_t len) {
    size_t sent = 0;
    while (sent < len) {
        ssize_t n = send(sock, static_cast<const char*>(buf) + sent, len - sent, 0);
        if (n <= 0) return false;
        sent += n;
    }
    return true;
}

static bool sendFile(int sock, const fs::path& filepath) {
    std::ifstream infile(filepath, std::ios::binary | std::ios::ate);
    if (!infile.is_open()) {
        std::cerr << "File not found: " << filepath << std::endl;
        return false;
    }

    uint32_t filename_len = (uint32_t)filepath.string().size();
    uint64_t file_size    = (uint64_t)infile.tellg();
    infile.seekg(0, std::ios::beg);

    std::cout << "Sending: " << filepath << " (" << file_size << " bytes)" << std::endl;

    if (!send_all(sock, &filename_len, sizeof(filename_len))) return false;
    if (!send_all(sock, filepath.c_str(), filename_len))      return false;
    if (!send_all(sock, &file_size, sizeof(file_size)))       return false;

    std::vector<char> buffer(proto::BUF_SIZE);
    uint64_t remaining = file_size;
    while (remaining > 0) {
        size_t to_read = (size_t)std::min(remaining, (uint64_t)proto::BUF_SIZE);
        infile.read(buffer.data(), (std::streamsize)to_read);
        std::streamsize count = infile.gcount();
        if (count <= 0) break;
        if (!send_all(sock, buffer.data(), (size_t)count)) {
            std::cerr << "Failed to send data for: " << filepath << std::endl;
            return false;
        }
        g_client_state.bytes_done.fetch_add((uint64_t)count);
        remaining -= (uint64_t)count;
    }

    return true;
}

static int find_server(char* server_ip);

static void fail_connect() {
    g_client_state.connection_failed.store(true);
    g_client_state.done.store(true);
}

static void* discovery_and_send_thread(void* arg) {
    fs::path directory = *static_cast<fs::path*>(arg);
    delete static_cast<fs::path*>(arg);

    char server_ip[INET_ADDRSTRLEN];
    if (find_server(server_ip) != 0) {
        if (!g_client_state.cancelled.load()) fail_connect();
        g_client_state.done.store(true);
        return nullptr;
    }
    if (g_client_state.cancelled.load()) { g_client_state.done.store(true); return nullptr; }

    g_client_state.setStatus("Connecting...");
    Socket tcp(socket(AF_INET, SOCK_STREAM, 0));
    if (!tcp.valid()) { fail_connect(); return nullptr; }

    sockaddr_in serv{};
    serv.sin_family = AF_INET;
    serv.sin_port   = htons(proto::TCP_PORT);
    if (inet_pton(AF_INET, server_ip, &serv.sin_addr) <= 0 ||
        connect(tcp, (sockaddr*)&serv, sizeof(serv)) < 0) {
        fail_connect();
        return nullptr;
    }

    uint64_t total = 0;
    for (const auto& entry : fs::recursive_directory_iterator(directory))
        if (fs::is_regular_file(entry.path()))
            total += fs::file_size(entry.path());
    g_client_state.bytes_total.store(total);

    for (const auto& entry : fs::recursive_directory_iterator(directory)) {
        if (g_client_state.cancelled.load()) break;
        const path& p = entry.path();
        if (fs::is_regular_file(p)) {
            g_client_state.setStatus(p.filename().string());
            if (!sendFile(tcp.fd, p)) break;
        }
    }

    uint32_t sentinel = proto::EOF_SENTINEL;
    send_all(tcp.fd, &sentinel, sizeof(sentinel));

    g_client_state.setStatus("");
    g_client_state.done.store(true);
    return nullptr;
}

static int find_server(char* server_ip) {
    int udp_fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (udp_fd < 0) return -1;

    sockaddr_in addr{};
    addr.sin_family      = AF_INET;
    addr.sin_port        = htons(proto::MULTICAST_PORT);
    addr.sin_addr.s_addr = inet_addr(proto::MULTICAST_GROUP);

    const char* msg = "DISCOVER_SERVER";
    if (sendto(udp_fd, msg, strlen(msg), 0, (sockaddr*)&addr, sizeof(addr)) < 0) {
        close(udp_fd);
        return -1;
    }

    // Poll in 100ms slices so we can react to cancel within 100ms
    auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(3);
    while (std::chrono::steady_clock::now() < deadline) {
        if (g_client_state.cancelled.load()) {
            close(udp_fd);
            return -1;
        }
        struct timeval tv{0, 100000};
        fd_set fds;
        FD_ZERO(&fds);
        FD_SET(udp_fd, &fds);
        if (select(udp_fd + 1, &fds, nullptr, nullptr, &tv) > 0) {
            sockaddr_in from{};
            socklen_t   fromlen = sizeof(from);
            char        buf[256];
            ssize_t     n = recvfrom(udp_fd, buf, sizeof(buf) - 1, 0, (sockaddr*)&from, &fromlen);
            if (n > 0) {
                buf[n] = '\0';
                if (strcmp(buf, "SERVER_HERE") == 0) {
                    inet_ntop(AF_INET, &from.sin_addr, server_ip, INET_ADDRSTRLEN);
                    close(udp_fd);
                    return 0;
                }
            }
        }
    }

    close(udp_fd);
    return -1;
}

int transfer_files(fs::path directory) {
    g_client_state.reset();
    g_client_state.setStatus("Searching for receiver...");

    fs::path* arg = new fs::path(directory);
    pthread_t thread;
    if (pthread_create(&thread, nullptr, discovery_and_send_thread, arg) != 0) {
        delete arg;
        return -1;
    }
    pthread_detach(thread);
    return 0;
}

#ifndef __SWITCH__
int main(int argc, char* argv[]) {
    if (argc < 2) {
        std::cerr << "Usage: " << argv[0] << " <path>" << std::endl;
        return 1;
    }
    if (transfer_files(fs::path(argv[1])) != 0) return 1;
    while (!isClientTransferDone()) usleep(16000);
    return 0;
}
#endif