#include <iostream>
#include <string>
#include <vector>
#include <thread>
#include <mutex>
#include <atomic>
#include <chrono>
#include <ctime>
#include <iomanip>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>

#pragma comment(lib, "ws2_32.lib")

// 全局变量
std::mutex cout_mutex;
std::atomic<int> total_sent(0);
std::atomic<bool> stop_flag(false);

enum class Protocol { UDP, TCP };

// 获取当前时间字符串
std::string get_current_time() {
    auto now = std::chrono::system_clock::now();
    std::time_t now_time = std::chrono::system_clock::to_time_t(now);
    std::tm now_tm;
    localtime_s(&now_tm, &now_time);

    char buffer[80];
    strftime(buffer, sizeof(buffer), "%H:%M:%S", &now_tm);
    return std::string(buffer);
}

std::string get_windows_error_message(int error_code) {
    LPSTR buffer = nullptr;
    FormatMessageA(
        FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
        nullptr, error_code, 0,
        (LPSTR)&buffer, 0, nullptr);
    std::string message(buffer ? buffer : "Unknown error");
    if (buffer) LocalFree(buffer);
    return message;
}

// 域名解析函数
std::string resolve_dns(const std::string& hostname) {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        throw std::runtime_error("WSAStartup失败");
    }

    addrinfo hints = { 0 };
    hints.ai_family = AF_INET; // IPv4
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = IPPROTO_TCP;

    addrinfo* result = nullptr;
    int error = getaddrinfo(hostname.c_str(), nullptr, &hints, &result);
    if (error != 0) {
        WSACleanup();
        throw std::runtime_error("域名解析失败: " + std::to_string(error));
    }

    char ip_str[INET_ADDRSTRLEN];
    sockaddr_in* addr = reinterpret_cast<sockaddr_in*>(result->ai_addr);
    inet_ntop(AF_INET, &addr->sin_addr, ip_str, INET_ADDRSTRLEN);

    std::string ip_address(ip_str);
    freeaddrinfo(result);
    WSACleanup();

    return ip_address;
}

// UDP发送线程
void udp_send_thread(const std::string& target_ip, int target_port,
    const std::string& packet_data, float interval,
    int thread_id, int packets_per_thread) {
    WSADATA wsaData;
    SOCKET sock;
    sockaddr_in serverAddr;

    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::lock_guard<std::mutex> lock(cout_mutex);
        std::cerr << "UDP线程" << thread_id << ": WSAStartup失败" << std::endl;
        return;
    }

    if ((sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) == INVALID_SOCKET) {
        std::lock_guard<std::mutex> lock(cout_mutex);
        std::cerr << "UDP线程" << thread_id << ": 创建套接字失败" << std::endl;
        WSACleanup();
        return;
    }

    memset(&serverAddr, 0, sizeof(serverAddr));
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(target_port);
    inet_pton(AF_INET, target_ip.c_str(), &serverAddr.sin_addr);

    int sent = 0;
    while (!stop_flag && (packets_per_thread == 0 || sent < packets_per_thread)) {
        if (sendto(sock, packet_data.c_str(), packet_data.size(), 0,
            (sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR) {
            std::lock_guard<std::mutex> lock(cout_mutex);
            std::cerr << "UDP线程" << thread_id << ": 发送失败" << std::endl;
            break;
        }

        sent++;
        total_sent++;

        if (sent % 100 == 0) {
            std::lock_guard<std::mutex> lock(cout_mutex);
            std::cout << "[" << get_current_time() << "] "
                << "UDP线程" << thread_id << ": 已发送 " << sent << " 个包" << std::endl;
        }

        if (interval > 0) {
            std::this_thread::sleep_for(
                std::chrono::milliseconds(static_cast<int>(interval * 1000)));
        }
    }

    closesocket(sock);
    WSACleanup();

    std::lock_guard<std::mutex> lock(cout_mutex);
    std::cout << "UDP线程" << thread_id << ": 已完成，共发送 " << sent << " 个包" << std::endl;
}

// TCP发送线程
void tcp_send_thread(const std::string& target_ip, int target_port,
    const std::string& packet_data, float interval,
    int thread_id, int packets_per_thread) {
    WSADATA wsaData;
    SOCKET sock = INVALID_SOCKET;
    int sent_count = 0;
    const int timeout_ms = 2000; // 2秒超时

    // 1. 初始化Winsock
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        std::lock_guard<std::mutex> lock(cout_mutex);
        std::cerr << "[" << get_current_time() << "] TCP线程" << thread_id
            << " WSAStartup失败: " << WSAGetLastError() << std::endl;
        return;
    }

    // 2. 创建socket
    if ((sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) == INVALID_SOCKET) {
        std::lock_guard<std::mutex> lock(cout_mutex);
        std::cerr << "[" << get_current_time() << "] TCP线程" << thread_id
            << " 创建socket失败: " << get_windows_error_message(WSAGetLastError()) << std::endl;
        WSACleanup();
        return;
    }

    // 3. 设置超时选项
    setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, (const char*)&timeout_ms, sizeof(timeout_ms));

    // 4. 准备目标地址
    sockaddr_in serverAddr;
    memset(&serverAddr, 0, sizeof(serverAddr));
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(target_port);
    if (inet_pton(AF_INET, target_ip.c_str(), &serverAddr.sin_addr) <= 0) {
        std::lock_guard<std::mutex> lock(cout_mutex);
        std::cerr << "[" << get_current_time() << "] TCP线程" << thread_id
            << " 无效的IP地址: " << target_ip << std::endl;
        closesocket(sock);
        WSACleanup();
        return;
    }

    // 5. 建立连接
    {
        std::lock_guard<std::mutex> lock(cout_mutex);
        std::cout << "[" << get_current_time() << "] TCP线程" << thread_id
            << " 正在连接 " << target_ip << ":" << target_port << "..." << std::endl;
    }

    if (connect(sock, (sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR) {
        int err = WSAGetLastError();
        std::lock_guard<std::mutex> lock(cout_mutex);
        std::cerr << "[" << get_current_time() << "] TCP线程" << thread_id
            << " 连接失败: " << err << " - " << get_windows_error_message(err) << std::endl;
        closesocket(sock);
        WSACleanup();
        return;
    }

    // 6. 连接成功，开始发送数据
    {
        std::lock_guard<std::mutex> lock(cout_mutex);
        std::cout << "[" << get_current_time() << "] TCP线程" << thread_id
            << " 连接成功，开始发送数据..." << std::endl;
    }

    const char* data_ptr = packet_data.c_str();
    const int data_len = static_cast<int>(packet_data.size());

    while (!stop_flag && (packets_per_thread == 0 || sent_count < packets_per_thread)) {
        // 7. 发送数据包
        int bytes_sent = send(sock, data_ptr, data_len, 0);

        if (bytes_sent == SOCKET_ERROR) {
            int err = WSAGetLastError();
            std::lock_guard<std::mutex> lock(cout_mutex);
            std::cerr << "[" << get_current_time() << "] TCP线程" << thread_id
                << " 发送失败: " << err << " - " << get_windows_error_message(err) << std::endl;
            break;
        }

        sent_count++;
        total_sent++;

        // 8. 输出进度
        if (sent_count % 10 == 0) {
            std::lock_guard<std::mutex> lock(cout_mutex);
            std::cout << "[" << get_current_time() << "] TCP线程" << thread_id
                << " 已发送 " << sent_count << " 个包 ("
                << sent_count * data_len << " 字节)" << std::endl;
        }

        // 9. 间隔控制
        if (interval > 0) {
            std::this_thread::sleep_for(
                std::chrono::milliseconds(static_cast<int>(interval * 1000)));
        }
    }

    // 10. 关闭连接
    {
        std::lock_guard<std::mutex> lock(cout_mutex);
        std::cout << "[" << get_current_time() << "] TCP线程" << thread_id
            << " 关闭连接，共发送 " << sent_count << " 个包" << std::endl;
    }

    shutdown(sock, SD_SEND);
    closesocket(sock);
    WSACleanup();
}

// 获取用户输入
void getUserInput(std::string& target_ip, int& target_port, std::string& packet_data,
    float& interval, int& total_count, int& thread_count, Protocol& protocol) {
    std::cout << "=== 多协议网络数据包发送工具 ===" << std::endl;

    // 选择协议
    while (true) {
        std::cout << "请选择协议 (1-UDP, 2-TCP): ";
        std::string protocol_input;
        std::getline(std::cin, protocol_input);

        if (protocol_input == "1") {
            protocol = Protocol::UDP;
            break;
        }
        else if (protocol_input == "2") {
            protocol = Protocol::TCP;
            break;
        }
        else {
            std::cout << "错误: 请输入1或2" << std::endl;
        }
    }

    // 获取目标地址（IP或域名）
    while (true) {
        std::cout << "请输入目标地址(IP或域名): ";
        std::getline(std::cin, target_ip);
        if (!target_ip.empty()) {
            // 尝试解析域名
            try {
                // 检查是否是有效的IP地址
                sockaddr_in sa;
                if (inet_pton(AF_INET, target_ip.c_str(), &(sa.sin_addr)) != 1) {
                    std::cout << "正在解析域名: " << target_ip << "..." << std::endl;
                    std::string resolved_ip = resolve_dns(target_ip);
                    std::cout << "解析成功: " << target_ip << " -> " << resolved_ip << std::endl;
                    target_ip = resolved_ip;
                }
                break;
            }
            catch (const std::exception& e) {
                std::cout << "错误: " << e.what() << std::endl;
            }
        }
        else {
            std::cout << "错误: 地址不能为空" << std::endl;
        }
    }

    // 获取端口号
    while (true) {
        std::cout << "请输入目标端口(1-65535): ";
        std::string port_input;
        std::getline(std::cin, port_input);

        try {
            target_port = std::stoi(port_input);
            if (target_port >= 1 && target_port <= 65535) {
                break;
            }
            std::cout << "错误: 端口必须在1-65535范围内" << std::endl;
        }
        catch (...) {
            std::cout << "错误: 请输入有效的端口号" << std::endl;
        }
    }

    // 获取发送内容
    std::cout << "请输入要发送的内容(默认为'ping'): ";
    std::getline(std::cin, packet_data);
    if (packet_data.empty()) {
        packet_data = "ping";
    }

    // 获取发送间隔
    while (true) {
        std::cout << "请输入发送间隔(秒，0表示无间隔，默认为0): ";
        std::string interval_input;
        std::getline(std::cin, interval_input);

        if (interval_input.empty()) {
            interval = 0.0f;
            break;
        }

        try {
            interval = std::stof(interval_input);
            if (interval >= 0) {
                break;
            }
            std::cout << "错误: 间隔时间不能为负数" << std::endl;
        }
        catch (...) {
            std::cout << "错误: 请输入有效的数字" << std::endl;
        }
    }

    // 获取总发送次数
    while (true) {
        std::cout << "请输入总发送次数(0表示无限，默认为0): ";
        std::string count_input;
        std::getline(std::cin, count_input);

        if (count_input.empty()) {
            total_count = 0;
            break;
        }

        try {
            total_count = std::stoi(count_input);
            if (total_count >= 0) {
                break;
            }
            std::cout << "错误: 次数不能为负数" << std::endl;
        }
        catch (...) {
            std::cout << "错误: 请输入有效的整数" << std::endl;
        }
    }

    // 获取线程数
    while (true) {
        std::cout << "请输入线程数(1-10000，默认为4): ";
        std::string thread_input;
        std::getline(std::cin, thread_input);

        if (thread_input.empty()) {
            thread_count = 4;
            break;
        }

        try {
            thread_count = std::stoi(thread_input);
            if (thread_count >= 1 && thread_count <= 10000) {
                break;
            }
            std::cout << "错误: 线程数必须在1-10000范围内" << std::endl;
        }
        catch (...) {
            std::cout << "错误: 请输入有效的整数" << std::endl;
        }
    }
}

int main() {
    std::string ip, data;
    int port, total_count, thread_count;
    float interval;
    Protocol protocol;

    // 获取用户输入
    getUserInput(ip, port, data, interval, total_count, thread_count, protocol);

    // 计算每个线程需要发送的包数
    int packets_per_thread = total_count == 0 ? 0 : total_count / thread_count;
    int remaining_packets = total_count == 0 ? 0 : total_count % thread_count;

    std::cout << "\n开始发送..." << std::endl;
    std::cout << "协议: " << (protocol == Protocol::UDP ? "UDP" : "TCP") << std::endl;
    std::cout << "目标: " << ip << ":" << port << std::endl;
    std::cout << "内容: " << data << std::endl;
    std::cout << "间隔: " << interval << "秒" << std::endl;
    std::cout << "线程数: " << thread_count << std::endl;
    std::cout << "总发送次数: " << (total_count == 0 ? "无限" : std::to_string(total_count)) << std::endl;
    std::cout << "按回车键停止发送...\n" << std::endl;

    // 启动发送线程
    std::lock_guard<std::mutex> lock(cout_mutex);
    std::vector<std::thread> threads;
    for (int i = 0; i < thread_count; ++i) {
        int this_thread_packets = packets_per_thread;
        if (i == thread_count - 1 && remaining_packets > 0) {
            this_thread_packets += remaining_packets;
        }

        if (protocol == Protocol::UDP) {
            threads.emplace_back(udp_send_thread, ip, port, data, interval,
                i + 1, this_thread_packets);
        }
        else {
            threads.emplace_back(tcp_send_thread, ip, port, data, interval,
                i + 1, this_thread_packets);
        }
    }


    // 等待用户输入停止
    std::cin.get();
    stop_flag = true;

    // 等待所有线程结束
    for (auto& t : threads) {
        if (t.joinable()) {
            t.join();  // 确保所有线程完成
        }
    }

    std::cout << "\n发送已停止，总共发送 " << total_sent << " 个包" << std::endl;
    std::cout << "按回车键退出...";
    std::cin.ignore();

    return 0;
}