#include "client/client.h"
#include "client/discovery.h"
#include "common/models.h"
#include "common/protocol.h"
#include <algorithm>
#include <cstdio>
#include <iostream>
#include <map>
#include <set>
#include <string>
#include <unistd.h>
#include <vector>

// ── help ──────────────────────────────────────────────────────────────────────
static void print_help() {
    std::cout <<
R"(NAME
       inventory-cli — device inventory client and hardware discovery agent

SYNOPSIS
       inventory-cli [--host HOST] [--port PORT] <command> [args...]
       inventory-cli --help

DESCRIPTION
       inventory-cli is a command-line client for the device-inventory server.
       It manages servers and hardware parts in the inventory, and can run
       local hardware discovery to populate the inventory automatically.

       Hardware discovery reads from dmidecode (CPU and memory), /sys/block
       (disks), and /sys/class/net with lspci (NICs).  Root privileges are
       required for dmidecode.

GLOBAL OPTIONS
       --host HOST, -H HOST
              Connect to inventory server at HOST.  Default: 127.0.0.1

       --port PORT, -p PORT
              Connect to inventory server on PORT.  Default: 9876

       --help, -h
              Print this help and exit.

COMMANDS
   Server Management
       add-server <name> <hostname> <location> <description>
              Register a new server in the inventory.

       list-servers
              List all registered servers.

       edit-server <id> <field> <value>
              Update a field on an existing server.  FIELD is one of:
                name | hostname | location | description

       remove-server <id>
              Permanently delete a server and all its parts.

   Part Type Management
       add-part-type <name> <description>
              Define a new part type label (e.g. "gpu", "psu").

       list-part-types
              List all part type labels.

       edit-part-type <id> <field> <value>
              Update a part type.  FIELD is one of:  name | description

       remove-part-type <id>
              Delete a part type label.

   Part Management
       add-part <type> <server_id> [key=value ...]
              Add a new hardware part to a server.

       upsert-part <type> <server_id> [key=value ...]
              Insert or update a part, matching on serial number or natural
              key (locator for memory, socket for cpu, device_name for disk).

       edit-part <part_id> [key=value ...]
              Update fields on an existing part.

       remove-part <part_id>
              Delete a part by ID.

       list-parts <server_id> [--type <type>]
              List all parts for a server, optionally filtered by type.

       get-part <part_id>
              Show all fields of a single part.

   Discovery
       discover <server_id> [--type <type>] [--dry-run]
              Detect hardware on this machine and upsert it into the server
              inventory under server_id.

              --type TYPE    Only discover one part type (see PART TYPES).
              --dry-run      Print discovered parts without contacting server.

       discover-only [--type <type>]
              Detect hardware on this machine and print a tree summary.
              Does not contact the inventory server.  Useful for inspection
              and troubleshooting.

              --type TYPE    Only discover one part type.

PART TYPES
       memory_stick    A physical RAM module installed in a slot.
       memory_slot     A DIMM slot on the motherboard.
       cpu             A physical processor installed in a socket.
       cpu_slot        A CPU socket on the motherboard.
       disk            A storage device (HDD, SSD, NVMe, virtual disk).
       nic             A physical network interface card.

FIELD KEYS
       serial                Hardware serial number.
       name                  Human-readable name or kernel interface name.
       manufacturer          Manufacturer or vendor name.
       locator               DIMM slot locator string (memory_stick/memory_slot).
       socket_designation    CPU socket label, e.g. "Proc 1" (cpu/cpu_slot).
       device_name           Kernel device name, e.g. sda, nvme0n1 (disk).
       speed_mhz             Memory speed in MHz.
       size_mb               Memory module size in megabytes.
       allowed_speed_mhz     Maximum supported speed for a memory slot.
       allowed_size_mb       Maximum module size supported by a slot.
       speed_ghz             CPU clock speed in GHz.
       cores                 Number of physical CPU cores.
       threads               Number of logical CPU threads.
       disk_size_gb          Disk capacity in gigabytes.
       disk_type             Disk technology: ssd | hdd | nvme | virtual | ...
       connection_type       Interface type: SATA | NVMe | SAS | ethernet | ...
       connection_speed_mbps Link speed in megabits per second.
       mac_address           MAC address of a NIC.
       ip_addresses          0x02-separated list of CIDR addresses.
       dhcp                  "true" if the interface uses DHCP.

EXAMPLES
       Register this machine:
              inventory-cli --host 10.0.0.1 add-server my-host my-host.local \
                "rack 4" "Primary web server"

       Run full hardware discovery and submit to server 3:
              inventory-cli --host 10.0.0.1 discover 3

       Preview discovery results locally without contacting the server:
              inventory-cli discover-only

       Discover only disks:
              inventory-cli discover-only --type disk

       List all parts for server 3:
              inventory-cli --host 10.0.0.1 list-parts 3

       List only memory sticks for server 3:
              inventory-cli --host 10.0.0.1 list-parts 3 --type memory_stick

NOTES
       Hardware discovery requires root privileges or the CAP_SYS_ADMIN
       capability for dmidecode access.  Run as root or via sudo.

       CPU "ID" fields from dmidecode (CPUID family/model/stepping) are not
       used as serial numbers because they are identical across processors of
       the same model.  Socket designation is used as the deduplication key.

       Memory slot blocks without a Locator field (sometimes emitted by HP
       ProLiant firmware) are silently skipped during discovery.

AUTHOR
       Part of the homelab device-inventory project.
)";
}

// ── parse key=value args into token pairs ─────────────────────────────────────
// Returns false and prints an error if any arg is malformed.
static bool parse_kv_args(int argc, char* argv[], int start,
                           std::vector<std::string>& tokens) {
    // Use ordered map to preserve insertion order (use vector of pairs)
    std::vector<std::pair<std::string,std::string>> kv_seen;
    std::map<std::string, size_t> kv_index;  // key → index in kv_seen

    for (int j = start; j < argc; ++j) {
        std::string arg = argv[j];
        size_t eq = arg.find('=');
        if (eq == std::string::npos) {
            std::cerr << "error: expected key=value, got: " << arg << "\n";
            return false;
        }
        std::string k = arg.substr(0, eq);
        std::string v = arg.substr(eq + 1);

        auto it = kv_index.find(k);
        if (it != kv_index.end()) {
            // Append to existing value with LS for list fields
            auto& existing = kv_seen[it->second].second;
            existing += LS;
            existing += v;
        } else {
            kv_index[k] = kv_seen.size();
            kv_seen.emplace_back(k, v);
        }
    }

    for (auto& [k, v] : kv_seen) {
        tokens.push_back(k);
        tokens.push_back(v);
    }
    return true;
}

// ── discover-only: tree-printing helpers ──────────────────────────────────────
static std::string kv_get(const std::map<std::string,std::string>& m,
                           const std::string& key,
                           const std::string& def = "") {
    auto it = m.find(key);
    return (it != m.end() && !it->second.empty() && it->second != "NULL")
        ? it->second : def;
}

// Replace LS (0x02) separators with ", " for display
static std::string ls_join(const std::string& s) {
    std::string out;
    out.reserve(s.size());
    for (char c : s) {
        if (c == LS) out += ", ";
        else         out += c;
    }
    return out;
}

static int cmd_discover_only(const std::string& filter_type) {
    Discovery disc;
    std::vector<DiscoveredPart> parts = filter_type.empty()
        ? disc.discover_all()
        : disc.discover(filter_type);

    // Bucket by type
    std::map<std::string, std::vector<const DiscoveredPart*>> by_type;
    for (auto& p : parts) by_type[p.type_name].push_back(&p);

    // Hostname as root of the tree
    char hbuf[256] = {};
    gethostname(hbuf, sizeof(hbuf) - 1);
    std::cout << hbuf << "\n";

    // Logical display order
    const std::vector<std::string> order = {
        PTYPE_CPU, PTYPE_CPU_SLOT,
        PTYPE_MEMORY_STICK, PTYPE_MEMORY_SLOT,
        PTYPE_DISK, PTYPE_NIC
    };
    static const std::map<std::string, std::string> labels = {
        {PTYPE_CPU,          "CPUs"},
        {PTYPE_CPU_SLOT,     "CPU Slots"},
        {PTYPE_MEMORY_STICK, "Memory Sticks"},
        {PTYPE_MEMORY_SLOT,  "Memory Slots"},
        {PTYPE_DISK,         "Disks"},
        {PTYPE_NIC,          "NICs"},
    };

    // Collect sections that actually have data
    std::vector<std::string> present;
    for (auto& t : order)
        if (by_type.count(t) && !by_type.at(t).empty())
            present.push_back(t);

    for (size_t si = 0; si < present.size(); ++si) {
        bool sec_last    = (si + 1 == present.size());
        const auto& type = present[si];
        auto& items      = by_type.at(type);

        std::string sec_pfx   = sec_last ? "└── " : "├── ";
        std::string child_pfx = sec_last ? "    " : "│   ";

        std::cout << sec_pfx << labels.at(type)
                  << " (" << items.size() << ")\n";

        for (size_t ii = 0; ii < items.size(); ++ii) {
            bool item_last = (ii + 1 == items.size());
            auto& m = items[ii]->kv;
            std::string line;

            if (type == PTYPE_CPU) {
                std::string sock    = kv_get(m, K_SOCKET);
                std::string name    = kv_get(m, K_NAME);
                std::string speed   = kv_get(m, K_SPEED_GHZ);
                std::string cores   = kv_get(m, K_CORES);
                std::string threads = kv_get(m, K_THREADS);
                if (!sock.empty()) line += "[" + sock + "]  ";
                line += name.empty() ? "(unknown)" : name;
                // Only append speed if the name doesn't already embed it (e.g. "@ 2.27GHz")
                if (!speed.empty() && name.find('@') == std::string::npos)
                    line += "  @  " + speed + " GHz";
                if (!cores.empty() || !threads.empty())
                    line += "  ·  " + (cores.empty() ? "?" : cores)
                         + " cores / " + (threads.empty() ? "?" : threads) + " threads";

            } else if (type == PTYPE_CPU_SLOT) {
                std::string sock = kv_get(m, K_SOCKET);
                std::string name = kv_get(m, K_NAME);
                if (!sock.empty()) line += "[" + sock + "]";
                if (!name.empty()) line += "  " + name;

            } else if (type == PTYPE_MEMORY_STICK) {
                std::string loc   = kv_get(m, K_LOCATOR);
                std::string size  = kv_get(m, K_SIZE_MB);
                std::string mfr   = kv_get(m, K_MANUFACTURER);
                std::string speed = kv_get(m, K_SPEED_MHZ);
                if (!loc.empty()) line += "[" + loc + "]  ";
                if (!size.empty()) {
                    uint64_t mb = std::stoull(size);
                    if (mb >= 1024)
                        line += std::to_string(mb / 1024) + " GB";
                    else
                        line += std::to_string(mb) + " MB";
                }
                if (!mfr.empty())   line += "  " + mfr;
                if (!speed.empty()) line += "  @ " + speed + " MHz";

            } else if (type == PTYPE_MEMORY_SLOT) {
                std::string loc      = kv_get(m, K_LOCATOR);
                std::string max_size = kv_get(m, K_ALLOWED_SIZE);
                std::string max_spd  = kv_get(m, K_ALLOWED_SPEED);
                if (!loc.empty()) line += "[" + loc + "]";
                if (!max_size.empty()) {
                    uint64_t mb = std::stoull(max_size);
                    line += "  up to ";
                    if (mb >= 1024)
                        line += std::to_string(mb / 1024) + " GB";
                    else
                        line += std::to_string(mb) + " MB";
                }
                if (!max_spd.empty()) line += "  @ " + max_spd + " MHz";

            } else if (type == PTYPE_DISK) {
                std::string dev    = kv_get(m, K_DEVICE_NAME);
                std::string size   = kv_get(m, K_DISK_SIZE);
                std::string dtype  = kv_get(m, K_DISK_TYPE);
                std::string model  = kv_get(m, K_MODEL);
                std::string serial = kv_get(m, K_SERIAL);
                if (!dev.empty())  line += dev + "  ";
                if (!size.empty()) {
                    char buf[32];
                    snprintf(buf, sizeof(buf), "%.1f GB", std::stod(size));
                    line += buf;
                }
                if (!dtype.empty())  line += "  (" + dtype + ")";
                if (!model.empty())  line += "  " + model;
                if (!serial.empty()) line += "  s/n: " + serial;

            } else if (type == PTYPE_NIC) {
                std::string ifname = kv_get(m, K_NAME);
                std::string model  = kv_get(m, K_MODEL);
                std::string mfr    = kv_get(m, K_MANUFACTURER);
                std::string mac    = kv_get(m, K_MAC);
                std::string ips    = kv_get(m, K_IPS);
                std::string speed  = kv_get(m, K_CONN_SPEED);
                if (!ifname.empty()) line += ifname + "  ";
                if (!model.empty())  line += model + "  ";
                else if (!mfr.empty()) line += mfr + "  ";
                if (!mac.empty())    line += mac + "  ";
                if (!ips.empty())    line += ls_join(ips) + "  ";
                if (!speed.empty() && speed != "0")
                    line += speed + " Mbps";
            }

            // Trim trailing spaces
            while (!line.empty() && line.back() == ' ') line.pop_back();

            std::cout << child_pfx << (item_last ? "└── " : "├── ") << line << "\n";
        }
    }

    return 0;
}

// ── discover subcommand ───────────────────────────────────────────────────────
static int cmd_discover(Client& client, const std::string& server_id,
                        const std::string& filter_type, bool dry_run) {
    Discovery disc;
    std::vector<DiscoveredPart> parts;

    if (filter_type.empty()) {
        parts = disc.discover_all();
    } else {
        parts = disc.discover(filter_type);
    }

    int submitted = 0;

    for (auto& p : parts) {
        if (dry_run) {
            std::cout << "[dry-run] type: " << p.type_name;
            auto sit = p.kv.find(K_SERIAL);
            if (sit != p.kv.end() && !sit->second.empty() && sit->second != "NULL")
                std::cout << "  serial: " << sit->second;
            std::cout << "\n";
            for (auto& [k, v] : p.kv) {
                std::cout << "  " << k << ": " << v << "\n";
            }
        } else {
            std::vector<std::string> tokens;
            tokens.push_back(CMD_UPSERT_PART);
            tokens.push_back(p.type_name);
            tokens.push_back(server_id);
            for (auto& [k, v] : p.kv) {
                tokens.push_back(k);
                tokens.push_back(v);
            }
            int rc = client.send(tokens);
            if (rc == 0) ++submitted;
        }
    }

    std::cout << "Discovered " << parts.size() << " parts";
    if (!dry_run)
        std::cout << ", submitted " << submitted << " to server";
    std::cout << "\n";

    return 0;
}

// ── main ──────────────────────────────────────────────────────────────────────
int main(int argc, char* argv[]) {
    std::string host = "127.0.0.1";
    uint16_t    port = DEFAULT_PORT;

    int i = 1;
    for (; i < argc; ++i) {
        std::string arg = argv[i];
        if ((arg == "--host" || arg == "-H") && i + 1 < argc) {
            host = argv[++i];
        } else if ((arg == "--port" || arg == "-p") && i + 1 < argc) {
            port = static_cast<uint16_t>(std::stoul(argv[++i]));
        } else if (arg == "--help" || arg == "-h") {
            print_help(); return 0;
        } else {
            break;
        }
    }

    if (i >= argc) { print_help(); return 1; }

    std::string sub = argv[i++];

    // ── Server management commands ────────────────────────────────────────────
    if (sub == "add-server") {
        if (i + 3 >= argc) {
            std::cerr << "usage: add-server <name> <hostname> <location> <description>\n";
            return 1;
        }
        std::vector<std::string> tokens = {
            CMD_ADD_SERVER, argv[i], argv[i+1], argv[i+2], argv[i+3]
        };
        return Client(host, port).send(tokens);
    }

    if (sub == "list-servers") {
        return Client(host, port).send({CMD_LIST_SERVERS});
    }

    if (sub == "edit-server") {
        if (i + 2 >= argc) {
            std::cerr << "usage: edit-server <id> <field> <value>\n"; return 1;
        }
        return Client(host, port).send({CMD_EDIT_SERVER, argv[i], argv[i+1], argv[i+2]});
    }

    if (sub == "remove-server") {
        if (i >= argc) { std::cerr << "usage: remove-server <id>\n"; return 1; }
        return Client(host, port).send({CMD_REMOVE_SERVER, argv[i]});
    }

    // ── Part type commands ────────────────────────────────────────────────────
    if (sub == "add-part-type") {
        if (i + 1 >= argc) {
            std::cerr << "usage: add-part-type <name> <description>\n"; return 1;
        }
        return Client(host, port).send({CMD_ADD_PART_TYPE, argv[i], argv[i+1]});
    }

    if (sub == "list-part-types") {
        return Client(host, port).send({CMD_LIST_PART_TYPES});
    }

    if (sub == "edit-part-type") {
        if (i + 2 >= argc) {
            std::cerr << "usage: edit-part-type <id> <field> <value>\n"; return 1;
        }
        return Client(host, port).send({CMD_EDIT_PART_TYPE, argv[i], argv[i+1], argv[i+2]});
    }

    if (sub == "remove-part-type") {
        if (i >= argc) { std::cerr << "usage: remove-part-type <id>\n"; return 1; }
        return Client(host, port).send({CMD_REMOVE_PART_TYPE, argv[i]});
    }

    // ── Part commands ─────────────────────────────────────────────────────────
    if (sub == "add-part" || sub == "upsert-part") {
        const std::string proto_cmd = (sub == "add-part") ? CMD_ADD_PART : CMD_UPSERT_PART;
        if (i + 1 >= argc) {
            std::cerr << "usage: " << sub << " <type> <server_id> [key=value ...]\n";
            return 1;
        }
        std::vector<std::string> tokens = {proto_cmd, argv[i], argv[i+1]};
        if (!parse_kv_args(argc, argv, i + 2, tokens)) return 1;
        return Client(host, port).send(tokens);
    }

    if (sub == "edit-part") {
        if (i >= argc) {
            std::cerr << "usage: edit-part <part_id> [key=value ...]\n"; return 1;
        }
        std::vector<std::string> tokens = {CMD_EDIT_PART, argv[i]};
        if (!parse_kv_args(argc, argv, i + 1, tokens)) return 1;
        return Client(host, port).send(tokens);
    }

    if (sub == "remove-part") {
        if (i >= argc) { std::cerr << "usage: remove-part <part_id>\n"; return 1; }
        return Client(host, port).send({CMD_REMOVE_PART, argv[i]});
    }

    if (sub == "list-parts") {
        if (i >= argc) { std::cerr << "usage: list-parts <server_id> [--type <type>]\n"; return 1; }
        std::string server_id = argv[i++];
        std::string type_filter;
        while (i < argc) {
            std::string a = argv[i++];
            if (a == "--type" && i < argc) type_filter = argv[i++];
        }
        std::vector<std::string> tokens = {CMD_LIST_PARTS, server_id};
        if (!type_filter.empty()) tokens.push_back(type_filter);
        return Client(host, port).send(tokens);
    }

    if (sub == "get-part") {
        if (i >= argc) { std::cerr << "usage: get-part <part_id>\n"; return 1; }
        return Client(host, port).send({CMD_GET_PART, argv[i]});
    }

    // ── Discovery ─────────────────────────────────────────────────────────────
    if (sub == "discover-only") {
        std::string filter_type;
        while (i < argc) {
            std::string a = argv[i++];
            if (a == "--type" && i < argc) {
                filter_type = argv[i++];
            } else {
                std::cerr << "error: unknown discover-only option: " << a << "\n";
                return 1;
            }
        }
        return cmd_discover_only(filter_type);
    }

    if (sub == "discover") {
        if (i >= argc) {
            std::cerr << "usage: discover <server_id> [--type <type>] [--dry-run]\n";
            return 1;
        }
        std::string server_id = argv[i++];
        std::string filter_type;
        bool dry_run = false;

        while (i < argc) {
            std::string a = argv[i++];
            if (a == "--type" && i < argc) {
                filter_type = argv[i++];
            } else if (a == "--dry-run") {
                dry_run = true;
            } else {
                std::cerr << "error: unknown discover option: " << a << "\n";
                return 1;
            }
        }

        Client client(host, port);
        return cmd_discover(client, server_id, filter_type, dry_run);
    }

    std::cerr << "error: unknown command '" << sub << "'\n\n";
    print_help();
    return 1;
}