<?php
/**
 * MMDBWriter - a pure-PHP library for generating MaxMind DB (.mmdb) files.
 *
 * Implements the MaxMind DB binary format from scratch:
 *   https://maxmind.github.io/MaxMind-DB/
 *
 * Supports IPv4 and IPv6 (dual-stack) databases, arbitrary nested
 * map/array data structures, and all MMDB scalar types.
 *
 * Usage:
 *   $w = new MMDBWriter(6, 'GeoIP-Custom', ['en']);
 *   $w->addRecord('1.2.3.0/24', ['country' => 'US', 'city' => 'Ashburn']);
 *   $w->addRecord('2001:db8::/32', ['country' => 'DE']);
 *   $w->write('/path/to/out.mmdb');
 */

// ---------------------------------------------------------------------
// Explicit value typing (use when auto-detection of PHP types isn't
// precise enough, e.g. you need a uint16 instead of the default uint32).
// ---------------------------------------------------------------------
final class MMDBType
{
    const POINTER = 1;
    const STRING  = 2;
    const DOUBLE  = 3;
    const BYTES   = 4;
    const UINT16  = 5;
    const UINT32  = 6;
    const MAP     = 7;
    const INT32   = 8;
    const UINT64  = 9;
    const UINT128 = 10;
    const ARRAY_T = 11;
    const BOOLEAN = 14;
    const FLOAT   = 15;
}

final class MMDBValue
{
    public $type;
    public $value;

    private function __construct(int $type, $value)
    {
        $this->type = $type;
        $this->value = $value;
    }

    public static function string(string $v): self  { return new self(MMDBType::STRING, $v); }
    public static function bytes(string $v): self    { return new self(MMDBType::BYTES, $v); }
    public static function double(float $v): self    { return new self(MMDBType::DOUBLE, $v); }
    public static function float(float $v): self     { return new self(MMDBType::FLOAT, $v); }
    public static function uint16(int $v): self       { return new self(MMDBType::UINT16, $v); }
    public static function uint32(int $v): self       { return new self(MMDBType::UINT32, $v); }
    public static function uint64(int $v): self       { return new self(MMDBType::UINT64, $v); }
    // $v may be a PHP int (<= 64 bit) or a hex string like "0x1ff...".
    public static function uint128($v): self          { return new self(MMDBType::UINT128, $v); }
    public static function int32(int $v): self        { return new self(MMDBType::INT32, $v); }
    public static function bool(bool $v): self        { return new self(MMDBType::BOOLEAN, $v); }
    public static function map(array $v): self        { return new self(MMDBType::MAP, $v); }
    public static function arr(array $v): self        { return new self(MMDBType::ARRAY_T, $v); }
}

// ---------------------------------------------------------------------
// Binary encoder for the MaxMind DB "data format"
// ---------------------------------------------------------------------
final class MMDBEncoder
{
    /** Encode any supported PHP value (or MMDBValue wrapper) to raw bytes. */
    public static function encode($value): string
    {
        if ($value instanceof MMDBValue) {
            return self::encodeTyped($value->type, $value->value);
        }

        if (is_bool($value)) {
            return self::encodeTyped(MMDBType::BOOLEAN, $value);
        }
        if (is_int($value)) {
            if ($value >= 0 && $value <= 0xFFFFFFFF) {
                return self::encodeTyped(MMDBType::UINT32, $value);
            }
            if ($value < 0 && $value >= -2147483648) {
                return self::encodeTyped(MMDBType::INT32, $value);
            }
            return self::encodeTyped(MMDBType::UINT64, $value);
        }
        if (is_float($value)) {
            return self::encodeTyped(MMDBType::DOUBLE, $value);
        }
        if (is_string($value)) {
            return self::encodeTyped(MMDBType::STRING, $value);
        }
        if (is_array($value)) {
            if (self::isList($value)) {
                return self::encodeTyped(MMDBType::ARRAY_T, $value);
            }
            return self::encodeTyped(MMDBType::MAP, $value);
        }

        throw new InvalidArgumentException('Unsupported value type: ' . gettype($value));
    }

    private static function isList(array $a): bool
    {
        $i = 0;
        foreach ($a as $k => $_) {
            if ($k !== $i++) {
                return false;
            }
        }
        return true;
    }

    private static function encodeTyped(int $type, $value): string
    {
        if ($type !== MMDBType::MAP && $type !== MMDBType::ARRAY_T) {
            return self::encodeScalar($type, $value);
        }
        switch ($type) {
            case MMDBType::MAP:
                $out = self::controlAndSize($type, count($value));
                foreach ($value as $k => $v) {
                    $out .= self::encodeTyped(MMDBType::STRING, (string) $k);
                    $out .= self::encode($v);
                }
                return $out;

            case MMDBType::ARRAY_T:
                $out = self::controlAndSize($type, count($value));
                foreach ($value as $v) {
                    $out .= self::encode($v);
                }
                return $out;
        }
    }

    /** Encode a single non-composite (leaf) MMDB value to bytes. Public so
     *  MMDBDataSection can reuse it for the leaves of its deduped structures. */
    public static function encodeScalar(int $type, $value): string
    {
        switch ($type) {
            case MMDBType::STRING:
                $payload = $value;
                return self::controlAndSize($type, strlen($payload)) . $payload;

            case MMDBType::BYTES:
                $payload = $value;
                return self::controlAndSize($type, strlen($payload)) . $payload;

            case MMDBType::DOUBLE:
                $payload = pack('E', $value); // 8-byte big-endian double
                return self::controlAndSize($type, 8) . $payload;

            case MMDBType::FLOAT:
                $payload = pack('G', $value); // 4-byte big-endian float
                return self::controlAndSize($type, 4) . $payload;

            case MMDBType::UINT16:
                $payload = pack('n', $value & 0xFFFF);
                return self::controlAndSize($type, 2) . $payload;

            case MMDBType::UINT32:
                $payload = pack('N', $value & 0xFFFFFFFF);
                return self::controlAndSize($type, 4) . $payload;

            case MMDBType::INT32:
                $payload = pack('N', $value & 0xFFFFFFFF);
                return self::controlAndSize($type, 4) . $payload;

            case MMDBType::UINT64:
                $payload = pack('J', $value); // 8-byte big-endian uint64
                return self::controlAndSize($type, 8) . $payload;

            case MMDBType::UINT128:
                $payload = self::to16Bytes($value);
                return self::controlAndSize($type, 16) . $payload;

            case MMDBType::BOOLEAN:
                // Boolean has no payload; the value lives in the size field.
                return self::controlAndSize($type, $value ? 1 : 0);

            default:
                throw new InvalidArgumentException("Unsupported MMDB type: $type");
        }
    }

    private static function to16Bytes($value): string
    {
        if (is_string($value) && preg_match('/^0x[0-9a-fA-F]+$/', $value)) {
            $hex = str_pad(substr($value, 2), 32, '0', STR_PAD_LEFT);
            return hex2bin($hex);
        }
        // Treat as a (<=64-bit) PHP int, zero-padded to 128 bits.
        return str_repeat("\x00", 8) . pack('J', (int) $value);
    }

    /**
     * Build the control byte (+ extra size bytes, + extended-type byte)
     * for a given MMDB type and payload size, per the spec's control
     * byte format.
     */
    public static function controlAndSize(int $type, int $size): string
    {
        $typeBits = $type <= 7 ? $type : 0; // 0 = "look at next byte for real type"
        $out = '';

        if ($size < 29) {
            $out .= chr(($typeBits << 5) | $size);
        } elseif ($size < 285) {
            $out .= chr(($typeBits << 5) | 29);
            $out .= chr($size - 29);
        } elseif ($size < 65821) {
            $out .= chr(($typeBits << 5) | 30);
            $out .= substr(pack('N', $size - 285), 2, 2); // 2-byte BE
        } else {
            $out .= chr(($typeBits << 5) | 31);
            $out .= substr(pack('N', $size - 65821), 1, 3); // 3-byte BE
        }

        if ($type > 7) {
            $out .= chr($type - 7); // extended type byte
        }

        return $out;
    }
}

// ---------------------------------------------------------------------
// Deduplicating data section builder. Any value (string, number, or
// whole map/array structure) that has already been written once is
// replaced everywhere else with a 2-5 byte MMDB "pointer" record instead
// of being re-encoded, which is where most of the size savings on
// real-world data (repeated keys like "country", repeated country codes,
// repeated sub-structures) comes from.
// ---------------------------------------------------------------------
final class MMDBDataSection
{
    private $buffer = '';
    /** @var array<string,int> cache key => byte offset of its real (non-pointer) encoding */
    private $cache = [];

    public function bytes(): string
    {
        return $this->buffer;
    }

    /**
     * Ensure $value is present in the data section (writing it if this is
     * the first time it's been seen) and return the offset of its real
     * encoding. Used for top-level records referenced directly by tree
     * leaves, so this always returns a real offset, never a pointer.
     */
    public function internTop($value): int
    {
        [$type, $raw, $cacheKey] = $this->classify($value);
        if ($cacheKey !== null && isset($this->cache[$cacheKey])) {
            return $this->cache[$cacheKey];
        }
        $offset = strlen($this->buffer);
        if ($cacheKey !== null) {
            $this->cache[$cacheKey] = $offset;
        }
        $this->appendEncoding($type, $raw);
        return $offset;
    }

    /** Append bytes for a value nested inside a map/array: a pointer if
     *  it's a repeat, otherwise its full encoding (registering it for
     *  future reuse). */
    private function writeRef($value): void
    {
        [$type, $raw, $cacheKey] = $this->classify($value);
        if ($cacheKey !== null && isset($this->cache[$cacheKey])) {
            $this->buffer .= $this->pointerBytes($this->cache[$cacheKey]);
            return;
        }
        $offset = strlen($this->buffer);
        if ($cacheKey !== null) {
            $this->cache[$cacheKey] = $offset;
        }
        $this->appendEncoding($type, $raw);
    }

    private function appendEncoding(int $type, $raw): void
    {
        if ($type === MMDBType::MAP) {
            $this->buffer .= MMDBEncoder::controlAndSize(MMDBType::MAP, count($raw));
            foreach ($raw as $k => $v) {
                $this->writeRef((string) $k);
                $this->writeRef($v);
            }
            return;
        }
        if ($type === MMDBType::ARRAY_T) {
            $this->buffer .= MMDBEncoder::controlAndSize(MMDBType::ARRAY_T, count($raw));
            foreach ($raw as $v) {
                $this->writeRef($v);
            }
            return;
        }
        $this->buffer .= MMDBEncoder::encodeScalar($type, $raw);
    }

    /** @return array{0:int,1:mixed,2:?string} [type, raw value, cache key or null if not worth caching] */
    private function classify($value): array
    {
        if ($value instanceof MMDBValue) {
            $type = $value->type;
            $raw = $value->value;
        } elseif (is_bool($value)) {
            $type = MMDBType::BOOLEAN;
            $raw = $value;
        } elseif (is_int($value)) {
            if ($value >= 0 && $value <= 0xFFFFFFFF) {
                $type = MMDBType::UINT32;
            } elseif ($value < 0 && $value >= -2147483648) {
                $type = MMDBType::INT32;
            } else {
                $type = MMDBType::UINT64;
            }
            $raw = $value;
        } elseif (is_float($value)) {
            $type = MMDBType::DOUBLE;
            $raw = $value;
        } elseif (is_string($value)) {
            $type = MMDBType::STRING;
            $raw = $value;
        } elseif (is_array($value)) {
            $type = $this->isList($value) ? MMDBType::ARRAY_T : MMDBType::MAP;
            $raw = $value;
        } else {
            throw new InvalidArgumentException('Unsupported value type: ' . gettype($value));
        }

        // A pointer costs 2-5 bytes; booleans already cost 1 byte with no
        // payload, so caching them can only ever waste space -- skip.
        $cacheKey = $type === MMDBType::BOOLEAN ? null : ($type . ':' . serialize($raw));

        return [$type, $raw, $cacheKey];
    }

    private function isList(array $a): bool
    {
        $i = 0;
        foreach ($a as $k => $_) {
            if ($k !== $i++) {
                return false;
            }
        }
        return true;
    }

    /** Encode an MMDB pointer (type 1) record per the spec's 4 size classes. */
    private function pointerBytes(int $offset): string
    {
        if ($offset <= 0x7FF) { // 11-bit value, 1 extra byte
            $b0 = (MMDBType::POINTER << 5) | (0 << 3) | (($offset >> 8) & 0x07);
            return chr($b0) . chr($offset & 0xFF);
        }

        $v = $offset - 2048;
        if ($v <= 0x7FFFF) { // 19-bit value, 2 extra bytes
            $b0 = (MMDBType::POINTER << 5) | (1 << 3) | (($v >> 16) & 0x07);
            return chr($b0) . chr(($v >> 8) & 0xFF) . chr($v & 0xFF);
        }

        $v2 = $offset - 2048 - 524288;
        if ($v2 <= 0x7FFFFFF) { // 27-bit value, 3 extra bytes
            $b0 = (MMDBType::POINTER << 5) | (2 << 3) | (($v2 >> 24) & 0x07);
            return chr($b0) . chr(($v2 >> 16) & 0xFF) . chr(($v2 >> 8) & 0xFF) . chr($v2 & 0xFF);
        }

        // Size class 3: full absolute 32-bit offset, 4 extra bytes.
        $b0 = (MMDBType::POINTER << 5) | (3 << 3);
        return chr($b0) . pack('N', $offset);
    }
}

// ---------------------------------------------------------------------
// Binary trie used to build the MMDB search tree
// ---------------------------------------------------------------------
final class MMDBTree
{
    /** @var array<int, array{0:mixed,1:mixed}> node index => [left, right] */
    private $nodes = [[null, null]];

    public function nodeCount(): int
    {
        return count($this->nodes);
    }

    /**
     * @param int[] $bits       0/1 array, the path from the root
     * @param int   $prefixLen  how many bits of $bits to actually use
     * @param int   $dataOffset byte offset of this record's value in the data section
     */
    public function insert(array $bits, int $prefixLen, int $dataOffset): void
    {
        $current = 0;
        for ($i = 0; $i < $prefixLen; $i++) {
            $bit = $bits[$i];
            $isLast = ($i === $prefixLen - 1);
            $rec = $this->nodes[$current][$bit];

            if ($isLast) {
                $this->nodes[$current][$bit] = ['data' => $dataOffset];
                continue;
            }

            if ($rec === null) {
                $newIdx = count($this->nodes);
                $this->nodes[] = [null, null];
                $this->nodes[$current][$bit] = $newIdx;
                $current = $newIdx;
            } elseif (is_array($rec)) {
                // A less-specific network already terminates here; split it
                // so both children inherit its data, then keep descending.
                $newIdx = count($this->nodes);
                $this->nodes[] = [$rec, $rec];
                $this->nodes[$current][$bit] = $newIdx;
                $current = $newIdx;
            } else {
                $current = $rec;
            }
        }
    }

    /**
     * Look up the data offset of whatever already-inserted network covers
     * this bit path (the "parent" network), without mutating the tree.
     * Returns null if no covering network has been inserted yet.
     */
    public function lookupCoveringData(array $bits, int $prefixLen): ?int
    {
        $current = 0;
        for ($i = 0; $i < $prefixLen; $i++) {
            $bit = $bits[$i];
            $rec = $this->nodes[$current][$bit];
            if ($rec === null) {
                return null;
            }
            if (is_array($rec)) {
                return $rec['data'];
            }
            $current = $rec;
        }
        return null;
    }

    /**
     * Serialize the tree to bytes using the smallest record size (24, 28,
     * or 32 bits) that fits every record value.
     *
     * @return array{0:string,1:int} [tree bytes, record size used]
     */
    public function serialize(int $dataSectionLength): array
    {
        $nodeCount = count($this->nodes);
        $maxDataPointer = $dataSectionLength + $nodeCount + 16;
        $maxValue = max($nodeCount, $maxDataPointer);

        if ($maxValue <= 0xFFFFFF) {
            $recordSize = 24;
        } elseif ($maxValue <= 0xFFFFFFF) {
            $recordSize = 28;
        } elseif ($maxValue <= 0xFFFFFFFF) {
            $recordSize = 32;
        } else {
            throw new RuntimeException('Database too large for a 32-bit record size.');
        }

        $out = '';
        foreach ($this->nodes as $node) {
            $left = $this->recordValue($node[0], $nodeCount);
            $right = $this->recordValue($node[1], $nodeCount);
            $out .= $this->packPair($left, $right, $recordSize);
        }

        return [$out, $recordSize];
    }

    private function recordValue($rec, int $nodeCount): int
    {
        if ($rec === null) {
            return $nodeCount; // "not found" sentinel
        }
        if (is_array($rec)) {
            return $rec['data'] + $nodeCount + 16;
        }
        return $rec; // pointer to another node
    }

    private function packPair(int $left, int $right, int $recordSize): string
    {
        if ($recordSize === 24) {
            return substr(pack('N', $left), 1, 3) . substr(pack('N', $right), 1, 3);
        }
        if ($recordSize === 32) {
            return pack('N', $left) . pack('N', $right);
        }
        // 28-bit: 7 bytes total, per the MMDB spec's packing scheme.
        $b0 = ($left >> 20) & 0xFF;
        $b1 = ($left >> 12) & 0xFF;
        $b2 = ($left >> 4) & 0xFF;
        $b3 = (($left & 0xF) << 4) | (($right >> 24) & 0xF);
        $b4 = ($right >> 16) & 0xFF;
        $b5 = ($right >> 8) & 0xFF;
        $b6 = $right & 0xFF;
        return chr($b0) . chr($b1) . chr($b2) . chr($b3) . chr($b4) . chr($b5) . chr($b6);
    }
}

// ---------------------------------------------------------------------
// Top-level writer: ties the tree, data section and metadata together
// ---------------------------------------------------------------------
final class MMDBWriter
{
    private $ipVersion;
    private $totalDepth;
    private $tree;
    private $data;
    private $rawDataByOffset = [];
    private $pending = [];
    private $databaseType;
    private $languages;
    private $description;

    /**
     * @param int      $ipVersion    4 (IPv4-only tree) or 6 (dual-stack tree)
     * @param string   $databaseType free-form name, e.g. "GeoIP-Custom"
     * @param string[] $languages    language codes present in $description
     * @param array    $description  language code => human description
     */
    public function __construct(
        int $ipVersion = 6,
        string $databaseType = 'Custom',
        array $languages = ['en'],
        array $description = ['en' => 'Custom MMDB database']
    ) {
        if ($ipVersion !== 4 && $ipVersion !== 6) {
            throw new InvalidArgumentException('ipVersion must be 4 or 6');
        }
        $this->ipVersion = $ipVersion;
        $this->totalDepth = $ipVersion === 4 ? 32 : 128;
        $this->tree = new MMDBTree();
        $this->data = new MMDBDataSection();
        $this->databaseType = $databaseType;
        $this->languages = $languages;
        $this->description = $description;
    }

    /**
     * Queue a network to be written. Records can be added in ANY order --
     * broader and narrower networks alike -- regardless of which one you
     * call this with first. At write() time, records are resolved from
     * broadest to narrowest so inheritance (see $inherit) always works
     * correctly no matter the call order.
     *
     * @param string $cidr    e.g. "1.2.3.0/24", "2001:db8::/32", or a bare IP
     * @param mixed  $data    scalar/array/MMDBValue data to associate with the network
     * @param bool   $inherit if true (default), and this network falls inside
     *                        a broader network also added to this writer,
     *                        fields from that broader record are merged in
     *                        first (this record's own fields still win on
     *                        conflicts).
     */
    public function addRecord(string $cidr, $data, bool $inherit = true): void
    {
        [$ip, $prefixLen] = $this->parseCidr($cidr);
        [$bits, $treePrefixLen] = $this->ipToTreeBits($ip, $prefixLen);
        $this->pending[] = [$bits, $treePrefixLen, $data, $inherit, count($this->pending)];
    }

    /**
     * Resolve all queued records (broadest network first, so inheritance
     * is correct regardless of the order addRecord() was called in) and
     * insert them into the tree.
     */
    private function resolveRecords(): void
    {
        $records = $this->pending;
        usort($records, function ($a, $b) {
            return $a[1] <=> $b[1] ?: $a[4] <=> $b[4]; // prefix length asc, then original call order
        });

        foreach ($records as [$bits, $treePrefixLen, $data, $inherit]) {
            if ($inherit) {
                $parentOffset = $this->tree->lookupCoveringData($bits, $treePrefixLen);
                if ($parentOffset !== null && isset($this->rawDataByOffset[$parentOffset])) {
                    $data = $this->mergeWithParent($this->rawDataByOffset[$parentOffset], $data);
                }
            }

            $offset = $this->data->internTop($data);
            $this->rawDataByOffset[$offset] = $data;
            $this->tree->insert($bits, $treePrefixLen, $offset);
        }
    }

    /**
     * Shallow-merge a child record over a parent record: parent fields are
     * inherited, and any key the child also sets overrides the parent's
     * value for that key. Non-map values (or MMDBValue-wrapped non-maps)
     * are left untouched -- the child value simply replaces the parent.
     */
    private function mergeWithParent($parentData, $data)
    {
        $parentArr = $parentData instanceof MMDBValue && $parentData->type === MMDBType::MAP
            ? $parentData->value
            : $parentData;
        $childArr = $data instanceof MMDBValue && $data->type === MMDBType::MAP
            ? $data->value
            : $data;

        if (!$this->isMap($parentArr) || !$this->isMap($childArr)) {
            return $data; // nothing sensible to merge; child value wins outright
        }

        return array_replace($parentArr, $childArr);
    }

    private function isMap($v): bool
    {
        if (!is_array($v) || $v === []) {
            return false;
        }
        $i = 0;
        foreach ($v as $k => $_) {
            if ($k !== $i++) {
                return true; // has at least one non-sequential/string key -> it's a map
            }
        }
        return false; // it's a plain list
    }

    public function write(string $path): void
    {
        $this->resolveRecords();

        [$treeBytes, $recordSize] = $this->tree->serialize(strlen($this->data->bytes()));

        $metadata = [
            'node_count'                  => MMDBValue::uint32($this->tree->nodeCount()),
            'record_size'                 => MMDBValue::uint16($recordSize),
            'ip_version'                  => MMDBValue::uint16($this->ipVersion),
            'database_type'               => MMDBValue::string($this->databaseType),
            'languages'                   => array_values($this->languages),
            'binary_format_major_version' => MMDBValue::uint16(2),
            'binary_format_minor_version' => MMDBValue::uint16(0),
            'build_epoch'                 => MMDBValue::uint64(time()),
            'description'                 => $this->description,
        ];

        $metadataBytes = MMDBEncoder::encode(MMDBValue::map($metadata));

        $separator = str_repeat("\x00", 16);
        $marker = "\xAB\xCD\xEFMaxMind.com";

        $blob = $treeBytes . $separator . $this->data->bytes() . $marker . $metadataBytes;

        if (file_put_contents($path, $blob) === false) {
            throw new RuntimeException("Failed to write $path");
        }
    }

    private function parseCidr(string $cidr): array
    {
        if (strpos($cidr, '/') !== false) {
            [$ip, $prefix] = explode('/', $cidr, 2);
            return [$ip, (int) $prefix];
        }
        $isV4 = strpos($cidr, ':') === false;
        return [$cidr, $isV4 ? 32 : 128];
    }

    private function ipToTreeBits(string $ip, int $prefixLen): array
    {
        $packed = @inet_pton($ip);
        if ($packed === false) {
            throw new InvalidArgumentException("Invalid IP address: $ip");
        }
        $isV4 = strlen($packed) === 4;

        if ($isV4 && $this->totalDepth === 128) {
            // Embed IPv4 networks under ::/96, as MaxMind's own dual-stack
            // databases do.
            $prefix = array_fill(0, 96, 0);
            $treePrefixLen = 96 + $prefixLen;
        } else {
            $prefix = [];
            $treePrefixLen = $prefixLen;
        }

        $bits = $prefix;
        foreach (str_split($packed) as $byte) {
            $b = ord($byte);
            for ($i = 7; $i >= 0; $i--) {
                $bits[] = ($b >> $i) & 1;
            }
        }

        return [$bits, $treePrefixLen];
    }
}