import { mergeProps } from "@base-ui/react/merge-props"
import { useRender } from "@base-ui/react/use-render"
import { cva, type VariantProps } from "class-variance-authority"

import { cn } from "#/lib/utils.ts"
import { Separator } from "#/components/ui/separator.tsx"

const buttonGroupVariants = cva(
  "flex w-fit items-stretch *:focus-visible:relative *:focus-visible:z-10 has-[>[data-slot=button-group]]:gap-2 has-[select[aria-hidden=true]:last-child]:[&>[data-slot=select-trigger]:last-of-type]:rounded-r-md [&>[data-slot=select-trigger]:not([class*='w-'])]:w-fit [&>input]:flex-1",
  {
    variants: {
      orientation: {
        horizontal:
          "*:data-slot:rounded-r-none [&>[data-slot]:not(:has(~[data-slot]))]:rounded-r-md! [&>[data-slot]~[data-slot]]:rounded-l-none [&>[data-slot]~[data-slot]]:border-l-0",
        vertical:
          "flex-col *:data-slot:rounded-b-none [&>[data-slot]:not(:has(~[data-slot]))]:rounded-b-md! [&>[data-slot]~[data-slot]]:rounded-t-none [&>[data-slot]~[data-slot]]:border-t-0",
      },
    },
    defaultVariants: {
      orientation: "horizontal",
    },
  }
)

function ButtonGroup({
  className,
  orientation,
  ...props
}: React.ComponentProps<"div"> & VariantProps<typeof buttonGroupVariants>) {
  return (
    <div
      role="group"
      data-slot="button-group"
      data-orientation={orientation}
      className={cn(buttonGroupVariants({ orientation }), className)}
      {...props}
    />
  )
}

function ButtonGroupText({
  className,
  render,
  ...props
}: useRender.ComponentProps<"div">) {
  return useRender({
    defaultTagName: "div",
    props: mergeProps<"div">(
      {
        className: cn(
          "flex items-center gap-2 rounded-md border bg-muted px-2.5 text-xs/relaxed font-medium [&_svg]:pointer-events-none [&_svg:not([class*='size-'])]:size-4",
          className
        ),
      },
      props
    ),
    render,
    state: {
      slot: "button-group-text",
    },
  })
}

function ButtonGroupSeparator({
  className,
  orientation = "vertical",
  ...props
}: React.ComponentProps<typeof Separator>) {
  return (
    <Separator
      data-slot="button-group-separator"
      orientation={orientation}
      className={cn(
        "relative self-stretch bg-input data-horizontal:mx-px data-horizontal:w-auto data-vertical:my-px data-vertical:h-auto",
        className
      )}
      {...props}
    />
  )
}

export {
  ButtonGroup,
  ButtonGroupSeparator,
  ButtonGroupText,
  buttonGroupVariants,
}


import tempfile
import os
import threading

from lxml.html import fromstring

from scrapling.core.storage import SQLiteStorageSystem, StorageSystemMixin
from scrapling.core.utils import _StorageTools

class TestGetBaseUrl:
    """Test StorageSystemMixin._get_base_url()"""

    def _make_storage(self, url=None):
        # Clear lru_cache between tests to avoid cross-test pollution
        StorageSystemMixin._get_base_url.cache_clear()
        return SQLiteStorageSystem(storage_file=":memory:", url=url)

    def test_returns_default_when_url_is_none(self):
        storage = self._make_storage(url=None)
        assert storage._get_base_url() == "default"

    def test_returns_default_when_url_is_empty(self):
        storage = self._make_storage(url="")
        assert storage._get_base_url() == "default"

    def test_returns_fld_for_valid_url(self):
        storage = self._make_storage(url="https://www.example.com/page")
        result = storage._get_base_url()
        assert result == "example.com"

    def test_url_is_lowercased(self):
        storage = self._make_storage(url="https://WWW.EXAMPLE.COM/Page")
        assert storage.url == "https://www.example.com/page"

class TestGetHash:
    """Test StorageSystemMixin._get_hash()"""

    def setup_method(self):
        StorageSystemMixin._get_hash.cache_clear()

    def test_deterministic_output(self):
        h1 = StorageSystemMixin._get_hash("test-identifier")
        h2 = StorageSystemMixin._get_hash("test-identifier")
        assert h1 == h2

    def test_different_input_different_output(self):
        h1 = StorageSystemMixin._get_hash("identifier-a")
        h2 = StorageSystemMixin._get_hash("identifier-b")
        assert h1 != h2

    def test_strips_and_lowercases(self):
        h1 = StorageSystemMixin._get_hash("  Hello  ")
        h2 = StorageSystemMixin._get_hash("hello")
        assert h1 == h2

    def test_includes_length_suffix(self):
        result = StorageSystemMixin._get_hash("test")
        # Format: {sha256_hex}_{byte_length}
        assert "_" in result
        hex_part, length_part = result.rsplit("_", 1)
        assert len(hex_part) == 64  # SHA-256 hex length
        assert length_part == str(len("test".encode("utf-8")))

class TestSQLiteStorageSystem:
    """Test SQLiteStorageSystem functionality"""

    def test_sqlite_storage_creation(self):
        """Test SQLite storage system creation"""
        storage = SQLiteStorageSystem(storage_file=":memory:")
        assert storage is not None

    def test_sqlite_storage_with_file(self):
        """Test SQLite storage with an actual file"""
        with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as tmp_file:
            db_path = tmp_file.name

        storage = None
        try:
            storage = SQLiteStorageSystem(storage_file=db_path)
            assert storage is not None
            assert os.path.exists(db_path)
        finally:
            if storage is not None:
                storage.close()
            if os.path.exists(db_path):
                os.unlink(db_path)

    def test_sqlite_storage_initialization_args(self):
        """Test SQLite storage with various initialization arguments"""
        storage = SQLiteStorageSystem(
            storage_file=":memory:",
            url="https://example.com"
        )
        assert storage is not None
        assert storage.url == "https://example.com"

class TestSaveRetrieveRoundTrip:
    """Test the save/retrieve round-trip - the core of the adaptive feature."""

    def _make_storage(self, url="https://example.com"):
        StorageSystemMixin._get_base_url.cache_clear()
        SQLiteStorageSystem.cache_clear()
        return SQLiteStorageSystem(storage_file=":memory:", url=url)

    def _make_element(self, html_str="

Hello

"):
        tree = fromstring(html_str)
        return tree.cssselect("p")[0] if tree.cssselect("p") else tree

    def test_save_and_retrieve(self):
        storage = self._make_storage()
        element = self._make_element()
        storage.save(element, "test-element")

        result = storage.retrieve("test-element")
        assert result is not None
        assert result["tag"] == "p"
        assert result["attributes"]["id"] == "target"
        assert result["attributes"]["class"] == "main"
        assert result["text"] == "Hello"

    def test_retrieve_nonexistent_returns_none(self):
        storage = self._make_storage()
        assert storage.retrieve("does-not-exist") is None

    def test_save_overwrites_existing(self):
        storage = self._make_storage()
        elem1 = self._make_element("

First

")
        elem2 = self._make_element("

Second

")

        storage.save(elem1, "my-element")
        storage.save(elem2, "my-element")

        result = storage.retrieve("my-element")
        assert result is not None
        assert result["attributes"]["id"] == "v2"
        assert result["text"] == "Second"

    def test_url_isolation(self):
        """Elements saved under one URL should not be retrievable under another."""
        SQLiteStorageSystem.cache_clear()
        StorageSystemMixin._get_base_url.cache_clear()

        # Use file-based storage so both instances share the same DB
        with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as tmp:
            db_path = tmp.name

        try:
            storage_a = SQLiteStorageSystem(storage_file=db_path, url="https://site-a.com")
            element = self._make_element()
            storage_a.save(element, "shared-id")

            SQLiteStorageSystem.cache_clear()
            StorageSystemMixin._get_base_url.cache_clear()

            storage_b = SQLiteStorageSystem(storage_file=db_path, url="https://site-b.com")
            assert storage_b.retrieve("shared-id") is None
        finally:
            storage_a.close()
            storage_b.close()
            if os.path.exists(db_path):
                os.unlink(db_path)

    def test_element_path_is_stored(self):
        storage = self._make_storage()
        element = self._make_element("

Text

")
        storage.save(element, "path-test")

        result = storage.retrieve("path-test")
        assert result is not None
        assert "path" in result
        # Path should be a list of tag names from root to element
        assert result["path"][-1] == "p"

    def test_element_with_children_and_siblings(self):
        storage = self._make_storage()
        html_str = "

Sibling
ChildChild2
"
        tree = fromstring(html_str)
        element = tree.cssselect("#target")[0]
        storage.save(element, "with-children")

        result = storage.retrieve("with-children")
        assert result is not None
        assert "children" in result
        assert "b" in result["children"]
        assert "i" in result["children"]
        assert "siblings" in result
        assert "p" in result["siblings"]

class TestStorageThreadSafety:
    """Test that SQLiteStorageSystem is safe under concurrent access."""

    def test_concurrent_saves(self):
        SQLiteStorageSystem.cache_clear()
        StorageSystemMixin._get_base_url.cache_clear()

        with tempfile.NamedTemporaryFile(suffix='.db', delete=False) as tmp:
            db_path = tmp.name

        storage = SQLiteStorageSystem(storage_file=db_path, url="https://example.com")
        errors = []

        def save_element(idx):
            try:
                html_str = f"

Text {idx}

"
                tree = fromstring(html_str)
                element = tree.cssselect("p")[0]
                storage.save(element, f"element-{idx}")
            except Exception as e:
                errors.append(e)

        threads = [threading.Thread(target=save_element, args=(i,)) for i in range(20)]
        for t in threads:
            t.start()
        for t in threads:
            t.join()

        assert len(errors) == 0, f"Thread safety errors: {errors}"

        # Verify all elements were saved
        for i in range(20):
            result = storage.retrieve(f"element-{i}")
            assert result is not None, f"element-{i} not found after concurrent save"

        storage.close()
        if os.path.exists(db_path):
            os.unlink(db_path)

class TestStorageToolsElementToDict:
    """Test _StorageTools.element_to_dict() directly."""

    def test_basic_element(self):
        tree = fromstring("

Hello

")
        elem = tree.cssselect("p")[0]
        result = _StorageTools.element_to_dict(elem)

        assert result["tag"] == "p"
        assert result["attributes"]["class"] == "foo"
        assert result["text"] == "Hello"
        assert "parent_name" in result
        assert result["parent_name"] == "div"

    def test_element_no_text(self):
        tree = fromstring("

")
        elem = tree.cssselect("p")[0]
        result = _StorageTools.element_to_dict(elem)
        assert result["text"] is None

    def test_element_no_attributes(self):
        tree = fromstring("

Plain

")
        elem = tree.cssselect("p")[0]
        result = _StorageTools.element_to_dict(elem)
        assert result["attributes"] == {}

    def test_element_strips_whitespace_attributes(self):
        tree = fromstring('

')
        elem = tree.cssselect("p")[0]
        result = _StorageTools.element_to_dict(elem)
        # Whitespace-only attribute values should be filtered out
        assert "data-val" not in result["attributes"]

class TestStorageToolsGetElementPath:
    """Test _StorageTools._get_element_path()."""

    def test_nested_path(self):
        tree = fromstring("

Text

")
        elem = tree.cssselect("p")[0]
        path = _StorageTools._get_element_path(elem)
        assert path[-1] == "p"
        assert "div" in path
        assert "body" in path

    def test_root_element_path(self):
        tree = fromstring("
Root
")
        path = _StorageTools._get_element_path(tree)
        assert path == ('html', 'body', 'div',)


#include "../sqlite-vec.h"
#include "sqlite-vec-internal.h"
#include <stdio.h>
#include <string.h>
#include <assert.h>
#include <math.h>

#define countof(x) (sizeof(x) / sizeof((x)[0]))

// Tests vec0_token_next(), the low-level tokenizer that extracts the next
// token from a raw char range. Covers every token type (identifier, digit,
// brackets, plus, equals), whitespace skipping, EOF on empty/whitespace-only
// input, error on unrecognised characters, and boundary behaviour where
// identifiers and digits stop at the next non-matching character.
void test_vec0_token_next() {
  printf("Starting %s...\n", __func__);
  struct Vec0Token token;
  int rc;
  char *input;

  // Single-character tokens
  input = "+";
  rc = vec0_token_next(input, input + 1, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_PLUS);

  input = "[";
  rc = vec0_token_next(input, input + 1, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_LBRACKET);

  input = "]";
  rc = vec0_token_next(input, input + 1, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_RBRACKET);

  input = "=";
  rc = vec0_token_next(input, input + 1, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_EQ);

  // Identifier
  input = "hello";
  rc = vec0_token_next(input, input + 5, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
  assert(token.start == input);
  assert(token.end == input + 5);

  // Identifier with underscores and digits
  input = "col_1a";
  rc = vec0_token_next(input, input + 6, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
  assert(token.end - token.start == 6);

  // Digit sequence
  input = "1234";
  rc = vec0_token_next(input, input + 4, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_DIGIT);
  assert(token.start == input);
  assert(token.end == input + 4);

  // Leading whitespace is skipped
  input = "  abc";
  rc = vec0_token_next(input, input + 5, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
  assert(token.end - token.start == 3);

  // Tab/newline whitespace
  input = "\t\n\r X";
  rc = vec0_token_next(input, input + 5, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_IDENTIFIER);

  // Empty input
  input = "";
  rc = vec0_token_next(input, input, &token);
  assert(rc == VEC0_TOKEN_RESULT_EOF);

  // Only whitespace
  input = "   ";
  rc = vec0_token_next(input, input + 3, &token);
  assert(rc == VEC0_TOKEN_RESULT_EOF);

  // Unrecognized character
  input = "@";
  rc = vec0_token_next(input, input + 1, &token);
  assert(rc == VEC0_TOKEN_RESULT_ERROR);

  input = "!";
  rc = vec0_token_next(input, input + 1, &token);
  assert(rc == VEC0_TOKEN_RESULT_ERROR);

  // Identifier stops at bracket
  input = "foo[";
  rc = vec0_token_next(input, input + 4, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
  assert(token.end - token.start == 3);

  // Digit stops at non-digit
  input = "42abc";
  rc = vec0_token_next(input, input + 5, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_DIGIT);
  assert(token.end - token.start == 2);

  // Left paren
  input = "(";
  rc = vec0_token_next(input, input + 1, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_LPAREN);

  // Right paren
  input = ")";
  rc = vec0_token_next(input, input + 1, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_RPAREN);

  // Comma
  input = ",";
  rc = vec0_token_next(input, input + 1, &token);
  assert(rc == VEC0_TOKEN_RESULT_SOME);
  assert(token.token_type == TOKEN_TYPE_COMMA);

  printf("  All vec0_token_next tests passed.\n");
}

// Tests Vec0Scanner, the stateful wrapper around vec0_token_next() that
// tracks position and yields successive tokens. Verifies correct tokenisation
// of full sequences like "abc float[128]" and "key=value", empty input,
// whitespace-heavy input, and expressions with operators ("a+b").
void test_vec0_scanner() {
  printf("Starting %s...\n", __func__);
  struct Vec0Scanner scanner;
  struct Vec0Token token;
  int rc;

  // Scan "abc float[128]"
  {
    const char *input = "abc float[128]";
    vec0_scanner_init(&scanner, input, (int)strlen(input));

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(token.end - token.start == 3);
    assert(strncmp(token.start, "abc", 3) == 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(token.end - token.start == 5);
    assert(strncmp(token.start, "float", 5) == 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_LBRACKET);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_DIGIT);
    assert(strncmp(token.start, "128", 3) == 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_RBRACKET);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_EOF);
  }

  // Scan "key=value"
  {
    const char *input = "key=value";
    vec0_scanner_init(&scanner, input, (int)strlen(input));

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(strncmp(token.start, "key", 3) == 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_EQ);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(strncmp(token.start, "value", 5) == 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_EOF);
  }

  // Scan empty string
  {
    const char *input = "";
    vec0_scanner_init(&scanner, input, 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_EOF);
  }

  // Scan with lots of whitespace
  {
    const char *input = "  a   b  ";
    vec0_scanner_init(&scanner, input, (int)strlen(input));

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(token.end - token.start == 1);
    assert(*token.start == 'a');

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(token.end - token.start == 1);
    assert(*token.start == 'b');

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_EOF);
  }

  // Scan "a+b"
  {
    const char *input = "a+b";
    vec0_scanner_init(&scanner, input, (int)strlen(input));

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_PLUS);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_EOF);
  }

  // Scan "diskann(k=v, k2=v2)"
  {
    const char *input = "diskann(k=v, k2=v2)";
    vec0_scanner_init(&scanner, input, (int)strlen(input));

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(strncmp(token.start, "diskann", 7) == 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_LPAREN);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(strncmp(token.start, "k", 1) == 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_EQ);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(strncmp(token.start, "v", 1) == 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_COMMA);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(strncmp(token.start, "k2", 2) == 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_EQ);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_IDENTIFIER);
    assert(strncmp(token.start, "v2", 2) == 0);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_SOME);
    assert(token.token_type == TOKEN_TYPE_RPAREN);

    rc = vec0_scanner_next(&scanner, &token);
    assert(rc == VEC0_TOKEN_RESULT_EOF);
  }

  printf("  All vec0_scanner tests passed.\n");
}

// Tests vec0_parse_vector_column(), which parses a vec0 column definition
// string like "embedding float[768] distance_metric=cosine" into a
// VectorColumnDefinition struct. Covers all element types (float/f32, int8/i8,
// bit), column names with underscores/digits, all distance metrics (L2, L1,
// cosine), the default metric, and error cases: empty input, missing type,
// unknown type, missing dimensions, unknown metric, unknown option key, and
// distance_metric on bit columns.
void test_vec0_parse_vector_column() {
  printf("Starting %s...\n", __func__);
  struct VectorColumnDefinition col;
  int rc;

  // Basic float column
  {
    const char *input = "embedding float[768]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.name_length == 9);
    assert(strncmp(col.name, "embedding", 9) == 0);
    assert(col.element_type == SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
    assert(col.dimensions == 768);
    assert(col.distance_metric == VEC0_DISTANCE_METRIC_L2);
    sqlite3_free(col.name);
  }

  // f32 alias
  {
    const char *input = "v f32[3]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.element_type == SQLITE_VEC_ELEMENT_TYPE_FLOAT32);
    assert(col.dimensions == 3);
    sqlite3_free(col.name);
  }

  // int8 column
  {
    const char *input = "quantized int8[256]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.element_type == SQLITE_VEC_ELEMENT_TYPE_INT8);
    assert(col.dimensions == 256);
    assert(col.name_length == 9);
    assert(strncmp(col.name, "quantized", 9) == 0);
    sqlite3_free(col.name);
  }

  // i8 alias
  {
    const char *input = "q i8[64]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.element_type == SQLITE_VEC_ELEMENT_TYPE_INT8);
    assert(col.dimensions == 64);
    sqlite3_free(col.name);
  }

  // bit column
  {
    const char *input = "bvec bit[1024]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.element_type == SQLITE_VEC_ELEMENT_TYPE_BIT);
    assert(col.dimensions == 1024);
    sqlite3_free(col.name);
  }

  // Column name with underscores and digits
  {
    const char *input = "col_name_2 float[10]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.name_length == 10);
    assert(strncmp(col.name, "col_name_2", 10) == 0);
    sqlite3_free(col.name);
  }

  // distance_metric=cosine
  {
    const char *input = "emb float[128] distance_metric=cosine";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.distance_metric == VEC0_DISTANCE_METRIC_COSINE);
    assert(col.dimensions == 128);
    sqlite3_free(col.name);
  }

  // distance_metric=L2 (explicit)
  {
    const char *input = "emb float[128] distance_metric=L2";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.distance_metric == VEC0_DISTANCE_METRIC_L2);
    sqlite3_free(col.name);
  }

  // distance_metric=L1
  {
    const char *input = "emb float[128] distance_metric=l1";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.distance_metric == VEC0_DISTANCE_METRIC_L1);
    sqlite3_free(col.name);
  }

  // SQLITE_EMPTY: empty string
  {
    const char *input = "";
    rc = vec0_parse_vector_column(input, 0, &col);
    assert(rc == SQLITE_EMPTY);
  }

  // SQLITE_EMPTY: non-vector column (text primary key)
  {
    const char *input = "document_id text primary key";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_EMPTY);
  }

  // SQLITE_EMPTY: non-vector column (partition key)
  {
    const char *input = "user_id integer partition key";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_EMPTY);
  }

  // SQLITE_EMPTY: no type (single identifier)
  {
    const char *input = "emb";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_EMPTY);
  }

  // SQLITE_EMPTY: unknown type
  {
    const char *input = "emb double[128]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_EMPTY);
  }

  // SQLITE_EMPTY: unknown type (unknowntype)
  {
    const char *input = "v unknowntype[128]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_EMPTY);
  }

  // SQLITE_EMPTY: missing brackets entirely
  {
    const char *input = "emb float";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_EMPTY);
  }

  // Error: zero dimensions
  {
    const char *input = "v float[0]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: empty brackets (no dimensions)
  {
    const char *input = "v float[]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: unknown distance metric
  {
    const char *input = "emb float[128] distance_metric=hamming";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: unknown distance metric (foo)
  {
    const char *input = "v float[128] distance_metric=foo";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: unknown option key
  {
    const char *input = "emb float[128] foobar=baz";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: distance_metric on bit type
  {
    const char *input = "emb bit[64] distance_metric=cosine";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // indexed by flat()
  {
    const char *input = "emb float[768] indexed by flat()";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_FLAT);
    assert(col.dimensions == 768);
    sqlite3_free(col.name);
  }

  // indexed by flat() with distance_metric
  {
    const char *input = "emb float[768] distance_metric=cosine indexed by flat()";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_FLAT);
    assert(col.distance_metric == VEC0_DISTANCE_METRIC_COSINE);
    sqlite3_free(col.name);
  }

  // indexed by flat() on int8
  {
    const char *input = "emb int8[256] indexed by flat()";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_FLAT);
    assert(col.element_type == SQLITE_VEC_ELEMENT_TYPE_INT8);
    sqlite3_free(col.name);
  }

  // indexed by flat() on bit
  {
    const char *input = "emb bit[64] indexed by flat()";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_FLAT);
    assert(col.element_type == SQLITE_VEC_ELEMENT_TYPE_BIT);
    sqlite3_free(col.name);
  }

  // default index_type is FLAT
  {
    const char *input = "emb float[768]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_FLAT);
    sqlite3_free(col.name);
  }

  // Error: indexed by (missing type name)
  {
    const char *input = "emb float[768] indexed by";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: indexed by unknown()
  {
    const char *input = "emb float[768] indexed by unknown()";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: indexed by flat (missing parens)
  {
    const char *input = "emb float[768] indexed by flat";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: indexed flat() (missing "by")
  {
    const char *input = "emb float[768] indexed flat()";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

#if SQLITE_VEC_ENABLE_IVF
  // IVF: indexed by ivf() — defaults
  {
    const char *input = "v float[4] indexed by ivf()";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_IVF);
    assert(col.dimensions == 4);
    assert(col.index_type == VEC0_INDEX_TYPE_IVF);
    assert(col.ivf.nlist == 128);  // default
    assert(col.ivf.nprobe == 10);  // default
    sqlite3_free(col.name);
  }

  // IVF: indexed by ivf(nlist=8) — nprobe auto-clamped to 8
  {
    const char *input = "v float[4] indexed by ivf(nlist=8)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_IVF);
    assert(col.index_type == VEC0_INDEX_TYPE_IVF);
    assert(col.ivf.nlist == 8);
    assert(col.ivf.nprobe == 8);  // clamped from default 10
    sqlite3_free(col.name);
  }

  // IVF: indexed by ivf(nlist=64, nprobe=8)
  {
    const char *input = "v float[4] indexed by ivf(nlist=64, nprobe=8)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_IVF);
    assert(col.ivf.nlist == 64);
    assert(col.ivf.nprobe == 8);
    sqlite3_free(col.name);
  }

  // IVF: with distance_metric before indexed by
  {
    const char *input = "v float[4] distance_metric=cosine indexed by ivf(nlist=16)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_IVF);
    assert(col.distance_metric == VEC0_DISTANCE_METRIC_COSINE);
    assert(col.index_type == VEC0_INDEX_TYPE_IVF);
    assert(col.ivf.nlist == 16);
    sqlite3_free(col.name);
  }

  // IVF: nlist=0 (deferred)
  {
    const char *input = "v float[4] indexed by ivf(nlist=0)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.nlist == 0);
    sqlite3_free(col.name);
  }

  // IVF error: nprobe > nlist
  {
    const char *input = "v float[4] indexed by ivf(nlist=4, nprobe=10)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // IVF error: unknown key
  {
    const char *input = "v float[4] indexed by ivf(bogus=1)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // IVF error: unknown index type (hnsw not supported)
  {
    const char *input = "v float[4] indexed by hnsw()";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Not IVF: no ivf config
  {
    const char *input = "v float[4]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_FLAT);
    sqlite3_free(col.name);
  }

  // IVF: quantizer=binary
  {
    const char *input = "v float[768] indexed by ivf(nlist=128, quantizer=binary)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_IVF);
    assert(col.ivf.nlist == 128);
    assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_BINARY);
    assert(col.ivf.oversample == 1);
    sqlite3_free(col.name);
  }

  // IVF: quantizer=int8
  {
    const char *input = "v float[768] indexed by ivf(nlist=64, quantizer=int8)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_INT8);
    sqlite3_free(col.name);
  }

  // IVF: quantizer=none (explicit)
  {
    const char *input = "v float[768] indexed by ivf(quantizer=none)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_NONE);
    sqlite3_free(col.name);
  }

  // IVF: oversample=10 with quantizer
  {
    const char *input = "v float[768] indexed by ivf(nlist=128, quantizer=binary, oversample=10)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_BINARY);
    assert(col.ivf.oversample == 10);
    assert(col.ivf.nlist == 128);
    sqlite3_free(col.name);
  }

  // IVF: all params
  {
    const char *input = "v float[768] distance_metric=cosine indexed by ivf(nlist=256, nprobe=16, quantizer=int8, oversample=4)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.distance_metric == VEC0_DISTANCE_METRIC_COSINE);
    assert(col.ivf.nlist == 256);
    assert(col.ivf.nprobe == 16);
    assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_INT8);
    assert(col.ivf.oversample == 4);
    sqlite3_free(col.name);
  }

  // IVF error: oversample > 1 without quantizer
  {
    const char *input = "v float[768] indexed by ivf(oversample=10)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // IVF error: unknown quantizer value
  {
    const char *input = "v float[768] indexed by ivf(quantizer=pq)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // IVF: quantizer with defaults (nlist=128 default, nprobe=10 default)
  {
    const char *input = "v float[768] indexed by ivf(quantizer=binary, oversample=5)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.nlist == 128);
    assert(col.ivf.nprobe == 10);
    assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_BINARY);
    assert(col.ivf.oversample == 5);
    sqlite3_free(col.name);
  }
#else
  // When IVF is disabled, parsing "ivf" should fail
  {
    const char *input = "v float[4] indexed by ivf()";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }
#endif /* SQLITE_VEC_ENABLE_IVF */

  printf("  All vec0_parse_vector_column tests passed.\n");
}

// Tests vec0_parse_partition_key_definition(), which parses a vec0 partition
// key column definition like "user_id integer partition key". Verifies correct
// parsing of integer and text partition keys, column name extraction, and
// rejection of invalid inputs: empty strings, non-partition-key definitions
// ("primary key"), and misspelled keywords.
void test_vec0_parse_partition_key_definition() {
  printf("Starting %s...\n", __func__);
  typedef struct {
    char * test;
    int expected_rc;
    const char *expected_column_name;
    int expected_column_type;
  } TestCase;

  TestCase suite[] = {
    {"user_id integer partition key", SQLITE_OK, "user_id", SQLITE_INTEGER},
    {"USER_id int partition key", SQLITE_OK, "USER_id", SQLITE_INTEGER},
    {"category text partition key", SQLITE_OK, "category", SQLITE_TEXT},

    {"", SQLITE_EMPTY, "", 0},
    {"document_id text primary key", SQLITE_EMPTY, "", 0},
    {"document_id text partition keyy", SQLITE_EMPTY, "", 0},
  };
  for(int i = 0; i < countof(suite); i++) {
    char * out_column_name;
    int out_column_name_length;
    int out_column_type;
    int rc;
    rc = vec0_parse_partition_key_definition(
      suite[i].test,
      strlen(suite[i].test),
      &out_column_name,
      &out_column_name_length,
      &out_column_type
    );
    assert(rc == suite[i].expected_rc);

    if(rc == SQLITE_OK) {
      assert(out_column_name_length == strlen(suite[i].expected_column_name));
      assert(strncmp(out_column_name, suite[i].expected_column_name, out_column_name_length) == 0);
      assert(out_column_type == suite[i].expected_column_type);
    }

    printf("  Passed: \"%s\"\n", suite[i].test);
  }
}

void test_distance_l2_sqr_float() {
  printf("Starting %s...\n", __func__);
  float d;

  // Identical vectors: distance = 0
  {
    float a[] = {1.0f, 2.0f, 3.0f};
    float b[] = {1.0f, 2.0f, 3.0f};
    d = _test_distance_l2_sqr_float(a, b, 3);
    assert(d == 0.0f);
  }

  // Orthogonal unit vectors: sqrt(1+1) = sqrt(2)
  {
    float a[] = {1.0f, 0.0f, 0.0f};
    float b[] = {0.0f, 1.0f, 0.0f};
    d = _test_distance_l2_sqr_float(a, b, 3);
    assert(fabsf(d - sqrtf(2.0f)) < 1e-6f);
  }

  // Known computation: [1,2,3] vs [4,5,6] = sqrt(9+9+9) = sqrt(27)
  {
    float a[] = {1.0f, 2.0f, 3.0f};
    float b[] = {4.0f, 5.0f, 6.0f};
    d = _test_distance_l2_sqr_float(a, b, 3);
    assert(fabsf(d - sqrtf(27.0f)) < 1e-5f);
  }

  // Single dimension: sqrt(16) = 4.0
  {
    float a[] = {3.0f};
    float b[] = {7.0f};
    d = _test_distance_l2_sqr_float(a, b, 1);
    assert(d == 4.0f);
  }

  printf("  All distance_l2_sqr_float tests passed.\n");
}

void test_distance_cosine_float() {
  printf("Starting %s...\n", __func__);
  float d;

  // Identical direction: distance = 0.0
  {
    float a[] = {1.0f, 0.0f};
    float b[] = {2.0f, 0.0f};
    d = _test_distance_cosine_float(a, b, 2);
    assert(fabsf(d - 0.0f) < 1e-6f);
  }

  // Orthogonal: distance = 1.0
  {
    float a[] = {1.0f, 0.0f};
    float b[] = {0.0f, 1.0f};
    d = _test_distance_cosine_float(a, b, 2);
    assert(fabsf(d - 1.0f) < 1e-6f);
  }

  // Opposite direction: distance = 2.0
  {
    float a[] = {1.0f, 0.0f};
    float b[] = {-1.0f, 0.0f};
    d = _test_distance_cosine_float(a, b, 2);
    assert(fabsf(d - 2.0f) < 1e-6f);
  }

  printf("  All distance_cosine_float tests passed.\n");
}

void test_distance_hamming() {
  printf("Starting %s...\n", __func__);
  float d;

  // Identical bitmaps: distance = 0
  {
    unsigned char a[] = {0xFF};
    unsigned char b[] = {0xFF};
    d = _test_distance_hamming(a, b, 8);
    assert(d == 0.0f);
  }

  // All different: distance = 8
  {
    unsigned char a[] = {0xFF};
    unsigned char b[] = {0x00};
    d = _test_distance_hamming(a, b, 8);
    assert(d == 8.0f);
  }

  // Half different: 0xFF vs 0x0F = 4 bits differ
  {
    unsigned char a[] = {0xFF};
    unsigned char b[] = {0x0F};
    d = _test_distance_hamming(a, b, 8);
    assert(d == 4.0f);
  }

  // Multi-byte: [0xFF, 0x00] vs [0x00, 0xFF] = 16 bits differ
  {
    unsigned char a[] = {0xFF, 0x00};
    unsigned char b[] = {0x00, 0xFF};
    d = _test_distance_hamming(a, b, 16);
    assert(d == 16.0f);
  }

  // Large vector (256 bits = 32 bytes) — exercises NEON path on ARM
  {
    unsigned char a[32];
    unsigned char b[32];
    memset(a, 0xFF, 32);
    memset(b, 0x00, 32);
    d = _test_distance_hamming(a, b, 256);
    assert(d == 256.0f);
  }

  // Large vector (1024 bits = 128 bytes) — exercises 64-byte NEON loop
  {
    unsigned char a[128];
    unsigned char b[128];
    memset(a, 0x00, 128);
    memset(b, 0x00, 128);
    // Set every other byte to 0xFF in a, 0x00 in b -> 8 bits per byte * 64 bytes = 512
    for (int i = 0; i < 128; i += 2) {
      a[i] = 0xFF;
    }
    d = _test_distance_hamming(a, b, 1024);
    assert(d == 512.0f);
  }

  printf("  All distance_hamming tests passed.\n");
}

#ifdef SQLITE_VEC_ENABLE_RESCORE

void test_rescore_quantize_float_to_bit() {
  printf("Starting %s...\n", __func__);
  uint8_t dst[16];

  // All positive -> all bits 1
  {
    float src[8] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
    memset(dst, 0, sizeof(dst));
    _test_rescore_quantize_float_to_bit(src, dst, 8);
    assert(dst[0] == 0xFF);
  }

  // All negative -> all bits 0
  {
    float src[8] = {-1.0f, -2.0f, -3.0f, -4.0f, -5.0f, -6.0f, -7.0f, -8.0f};
    memset(dst, 0xFF, sizeof(dst));
    _test_rescore_quantize_float_to_bit(src, dst, 8);
    assert(dst[0] == 0x00);
  }

  // Alternating positive/negative
  {
    float src[8] = {1.0f, -1.0f, 1.0f, -1.0f, 1.0f, -1.0f, 1.0f, -1.0f};
    _test_rescore_quantize_float_to_bit(src, dst, 8);
    // bits 0,2,4,6 set => 0b01010101 = 0x55
    assert(dst[0] == 0x55);
  }

  // Zero values -> bit is set (>= 0.0f)
  {
    float src[8] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
    _test_rescore_quantize_float_to_bit(src, dst, 8);
    assert(dst[0] == 0xFF);
  }

  // 128 dimensions -> 16 bytes output
  {
    float src[128];
    for (int i = 0; i < 128; i++) src[i] = (i % 2 == 0) ? 1.0f : -1.0f;
    memset(dst, 0, 16);
    _test_rescore_quantize_float_to_bit(src, dst, 128);
    // Even indices set: bits 0,2,4,6 in each byte => 0x55
    for (int i = 0; i < 16; i++) {
      assert(dst[i] == 0x55);
    }
  }

  printf("  All rescore_quantize_float_to_bit tests passed.\n");
}

void test_rescore_quantize_float_to_int8() {
  printf("Starting %s...\n", __func__);
  int8_t dst[256];

  // Uniform vector -> all zeros (range=0)
  {
    float src[8] = {5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f, 5.0f};
    _test_rescore_quantize_float_to_int8(src, dst, 8);
    for (int i = 0; i < 8; i++) {
#if SQLITE_VEC_ENABLE_IVF
void test_ivf_quantize_int8() {
  printf("Starting %s...\n", __func__);

  // Basic values in [-1, 1] range
  {
    float src[] = {0.0f, 1.0f, -1.0f, 0.5f};
    int8_t dst[4];
    ivf_quantize_int8(src, dst, 4);
    assert(dst[0] == 0);
    assert(dst[1] == 127);
    assert(dst[2] == -127);
    assert(dst[3] == 63);  // 0.5 * 127 = 63.5, truncated to 63
  }

  // Clamping: values beyond [-1, 1]
  {
    float src[] = {2.0f, -3.0f, 100.0f, -0.01f};
    int8_t dst[4];
    ivf_quantize_int8(src, dst, 4);
    assert(dst[0] == 127);   // clamped to 1.0
    assert(dst[1] == -127);  // clamped to -1.0
    assert(dst[2] == 127);   // clamped to 1.0
    assert(dst[3] == (int8_t)(-0.01f * 127.0f));
  }

  // Zero vector
  {
    float src[] = {0.0f, 0.0f, 0.0f, 0.0f};
    int8_t dst[4];
    ivf_quantize_int8(src, dst, 4);
    for (int i = 0; i < 4; i++) {
      assert(dst[i] == 0);
    }
  }

  // [0.0, 1.0] -> should map to [-128, 127]
  {
    float src[2] = {0.0f, 1.0f};
    _test_rescore_quantize_float_to_int8(src, dst, 2);
    assert(dst[0] == -128);
    assert(dst[1] == 127);
  }

  // [-1.0, 0.0] -> should map to [-128, 127]
  {
    float src[2] = {-1.0f, 0.0f};
    _test_rescore_quantize_float_to_int8(src, dst, 2);
    assert(dst[0] == -128);
    assert(dst[1] == 127);
  }

  // Single-element: range=0 -> 0
  {
    float src[1] = {42.0f};
    _test_rescore_quantize_float_to_int8(src, dst, 1);
    assert(dst[0] == 0);
  }

  // Verify range: all outputs in [-128, 127], min near -128, max near 127
  {
    float src[4] = {-100.0f, 0.0f, 100.0f, 50.0f};
    _test_rescore_quantize_float_to_int8(src, dst, 4);
    for (int i = 0; i < 4; i++) {
      assert(dst[i] >= -128 && dst[i] <= 127);
    }
    // Min maps to -128 (exact), max maps to ~127 (may lose 1 to float rounding)
    assert(dst[0] == -128);
    assert(dst[2] >= 126 && dst[2] <= 127);
    // Middle value (50) should be positive
    assert(dst[3] > 0);
  }

  printf("  All rescore_quantize_float_to_int8 tests passed.\n");
}

void test_rescore_quantized_byte_size() {
  printf("Starting %s...\n", __func__);

  // Bit quantizer: dims/8
  assert(_test_rescore_quantized_byte_size_bit(128) == 16);
  assert(_test_rescore_quantized_byte_size_bit(8) == 1);
  assert(_test_rescore_quantized_byte_size_bit(1024) == 128);

  // Int8 quantizer: dims
  assert(_test_rescore_quantized_byte_size_int8(128) == 128);
  assert(_test_rescore_quantized_byte_size_int8(8) == 8);
  assert(_test_rescore_quantized_byte_size_int8(1024) == 1024);

  printf("  All rescore_quantized_byte_size tests passed.\n");
}

void test_vec0_parse_vector_column_rescore() {
  // Negative zero
  {
    float src[] = {-0.0f};
    int8_t dst[1];
    ivf_quantize_int8(src, dst, 1);
    assert(dst[0] == 0);
  }

  // Single element
  {
    float src[] = {0.75f};
    int8_t dst[1];
    ivf_quantize_int8(src, dst, 1);
    assert(dst[0] == (int8_t)(0.75f * 127.0f));
  }

  // Boundary: exactly 1.0 and -1.0
  {
    float src[] = {1.0f, -1.0f};
    int8_t dst[2];
    ivf_quantize_int8(src, dst, 2);
    assert(dst[0] == 127);
    assert(dst[1] == -127);
  }

  printf("  All ivf_quantize_int8 tests passed.\n");
}

void test_ivf_quantize_binary() {
  printf("Starting %s...\n", __func__);

  // Basic sign-bit quantization: positive -> 1, negative/zero -> 0
  {
    float src[] = {1.0f, -1.0f, 0.5f, -0.5f, 0.0f, 0.1f, -0.1f, 2.0f};
    uint8_t dst[1];
    ivf_quantize_binary(src, dst, 8);
    // bit 0: 1.0 > 0 -> 1  (LSB)
    // bit 1: -1.0 -> 0
    // bit 2: 0.5 > 0 -> 1
    // bit 3: -0.5 -> 0
    // bit 4: 0.0 -> 0 (not > 0)
    // bit 5: 0.1 > 0 -> 1
    // bit 6: -0.1 -> 0
    // bit 7: 2.0 > 0 -> 1
    // Expected: bits 0,2,5,7 = 0b10100101 = 0xA5
    assert(dst[0] == 0xA5);
  }

  // All positive
  {
    float src[] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f};
    uint8_t dst[1];
    ivf_quantize_binary(src, dst, 8);
    assert(dst[0] == 0xFF);
  }

  // All negative
  {
    float src[] = {-1.0f, -2.0f, -3.0f, -4.0f, -5.0f, -6.0f, -7.0f, -8.0f};
    uint8_t dst[1];
    ivf_quantize_binary(src, dst, 8);
    assert(dst[0] == 0x00);
  }

  // All zero (zero is NOT > 0, so all bits should be 0)
  {
    float src[] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
    uint8_t dst[1];
    ivf_quantize_binary(src, dst, 8);
    assert(dst[0] == 0x00);
  }

  // Multi-byte: 16 dimensions -> 2 bytes
  {
    float src[16];
    for (int i = 0; i < 16; i++) src[i] = (i % 2 == 0) ? 1.0f : -1.0f;
    uint8_t dst[2];
    ivf_quantize_binary(src, dst, 16);
    // Even indices are positive: bits 0,2,4,6 in each byte
    // byte 0: bits 0,2,4,6 = 0b01010101 = 0x55
    // byte 1: same pattern = 0x55
    assert(dst[0] == 0x55);
    assert(dst[1] == 0x55);
  }

  // Single byte, only first bit set
  {
    float src[] = {0.1f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f, -1.0f};
    uint8_t dst[1];
    ivf_quantize_binary(src, dst, 8);
    assert(dst[0] == 0x01);
  }

  printf("  All ivf_quantize_binary tests passed.\n");
}

void test_ivf_config_parsing() {
void test_vec0_parse_vector_column_diskann() {
  printf("Starting %s...\n", __func__);
  struct VectorColumnDefinition col;
  int rc;

  // Basic bit quantizer
  {
    const char *input = "emb float[128] indexed by rescore(quantizer=bit)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_RESCORE);
    assert(col.rescore.quantizer_type == VEC0_RESCORE_QUANTIZER_BIT);
    assert(col.rescore.oversample == 8); // default
  // Existing syntax (no INDEXED BY) should have diskann.enabled == 0
  {
    const char *input = "emb float[128]";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type != VEC0_INDEX_TYPE_DISKANN);
    sqlite3_free(col.name);
  }

  // With distance_metric but no INDEXED BY
  {
    const char *input = "emb float[128] distance_metric=cosine";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type != VEC0_INDEX_TYPE_DISKANN);
    assert(col.distance_metric == VEC0_DISTANCE_METRIC_COSINE);
    sqlite3_free(col.name);
  }

  // Basic binary quantizer
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=binary)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_DISKANN);
    assert(col.diskann.quantizer_type == VEC0_DISKANN_QUANTIZER_BINARY);
    assert(col.diskann.n_neighbors == 72);  // default
    assert(col.diskann.search_list_size == 128);  // default
    assert(col.dimensions == 128);
    sqlite3_free(col.name);
  }

  // Int8 quantizer
  {
    const char *input = "emb float[128] indexed by rescore(quantizer=int8)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_RESCORE);
    assert(col.rescore.quantizer_type == VEC0_RESCORE_QUANTIZER_INT8);
    sqlite3_free(col.name);
  }

  // Bit quantizer with oversample
  {
    const char *input = "emb float[128] indexed by rescore(quantizer=bit, oversample=16)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_RESCORE);
    assert(col.rescore.quantizer_type == VEC0_RESCORE_QUANTIZER_BIT);
    assert(col.rescore.oversample == 16);
    sqlite3_free(col.name);
  }

  // Error: non-float element type
  {
    const char *input = "emb int8[128] indexed by rescore(quantizer=bit)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: dims not divisible by 8 for bit quantizer
  {
    const char *input = "emb float[100] indexed by rescore(quantizer=bit)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: missing quantizer
  {
    const char *input = "emb float[128] indexed by rescore(oversample=8)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // With distance_metric=cosine
  {
    const char *input = "emb float[128] distance_metric=cosine indexed by rescore(quantizer=int8)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_RESCORE);
    assert(col.distance_metric == VEC0_DISTANCE_METRIC_COSINE);
    assert(col.rescore.quantizer_type == VEC0_RESCORE_QUANTIZER_INT8);
    sqlite3_free(col.name);
  }

  printf("  All vec0_parse_vector_column_rescore tests passed.\n");
}

#endif /* SQLITE_VEC_ENABLE_RESCORE */
  // Default IVF config
  {
    const char *s = "v float[4] indexed by ivf()";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_IVF);
    assert(col.ivf.nlist == 128);   // default
    assert(col.ivf.nprobe == 10);   // default
    assert(col.ivf.quantizer == 0); // VEC0_IVF_QUANTIZER_NONE
    sqlite3_free(col.name);
  }

  // Custom nlist and nprobe
  {
    const char *s = "v float[4] indexed by ivf(nlist=64, nprobe=8)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.nlist == 64);
    assert(col.ivf.nprobe == 8);
    sqlite3_free(col.name);
  }

  // nlist=0 (deferred)
  {
    const char *s = "v float[4] indexed by ivf(nlist=0)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.nlist == 0);
    sqlite3_free(col.name);
  }

  // Quantizer options
  {
    const char *s = "v float[8] indexed by ivf(quantizer=int8)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_INT8);
    sqlite3_free(col.name);
  }

  {
    const char *s = "v float[8] indexed by ivf(quantizer=binary)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_BINARY);
    sqlite3_free(col.name);
  }

  // nprobe > nlist (explicit) should fail
  {
    const char *s = "v float[4] indexed by ivf(nlist=4, nprobe=10)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Unknown key
  {
    const char *s = "v float[4] indexed by ivf(bogus=1)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_ERROR);
  }

  // nlist > max (65536) should fail
  {
    const char *s = "v float[4] indexed by ivf(nlist=65537)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_ERROR);
  }

  // nlist at max boundary (65536) should succeed
  {
    const char *s = "v float[4] indexed by ivf(nlist=65536)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.nlist == 65536);
    sqlite3_free(col.name);
  }

  // oversample > 1 without quantization should fail
  {
    const char *s = "v float[4] indexed by ivf(oversample=4)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_ERROR);
  }

  // oversample with quantizer should succeed
  {
    const char *s = "v float[8] indexed by ivf(quantizer=int8, oversample=4)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.oversample == 4);
    assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_INT8);
    sqlite3_free(col.name);
  }

  // All options combined
  {
    const char *s = "v float[8] indexed by ivf(nlist=32, nprobe=4, quantizer=int8, oversample=2)";
    rc = vec0_parse_vector_column(s, (int)strlen(s), &col);
    assert(rc == SQLITE_OK);
    assert(col.ivf.nlist == 32);
    assert(col.ivf.nprobe == 4);
    assert(col.ivf.quantizer == VEC0_IVF_QUANTIZER_INT8);
    assert(col.ivf.oversample == 2);
    sqlite3_free(col.name);
  }

  printf("  All ivf_config_parsing tests passed.\n");
}
#endif /* SQLITE_VEC_ENABLE_IVF */
  // INT8 quantizer
  {
    const char *input = "v float[64] INDEXED BY diskann(neighbor_quantizer=int8)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_DISKANN);
    assert(col.diskann.quantizer_type == VEC0_DISKANN_QUANTIZER_INT8);
    sqlite3_free(col.name);
  }

  // Custom n_neighbors
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=binary, n_neighbors=48)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_DISKANN);
    assert(col.diskann.n_neighbors == 48);
    sqlite3_free(col.name);
  }

  // Custom search_list_size
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=binary, search_list_size=256)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.diskann.search_list_size == 256);
    sqlite3_free(col.name);
  }

  // Combined with distance_metric (distance_metric first)
  {
    const char *input = "emb float[128] distance_metric=cosine INDEXED BY diskann(neighbor_quantizer=int8)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.distance_metric == VEC0_DISTANCE_METRIC_COSINE);
    assert(col.index_type == VEC0_INDEX_TYPE_DISKANN);
    assert(col.diskann.quantizer_type == VEC0_DISKANN_QUANTIZER_INT8);
    sqlite3_free(col.name);
  }

  // Error: missing neighbor_quantizer (required)
  {
    const char *input = "emb float[128] INDEXED BY diskann(n_neighbors=72)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: empty parens
  {
    const char *input = "emb float[128] INDEXED BY diskann()";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: unknown quantizer
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=unknown)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: bad n_neighbors (not divisible by 8)
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=binary, n_neighbors=13)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: n_neighbors too large
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=binary, n_neighbors=512)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: missing BY
  {
    const char *input = "emb float[128] INDEXED diskann(neighbor_quantizer=binary)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: unknown algorithm
  {
    const char *input = "emb float[128] INDEXED BY hnsw(neighbor_quantizer=binary)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: unknown option key
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=binary, foobar=baz)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Case insensitivity for keywords
  {
    const char *input = "emb float[128] indexed by DISKANN(NEIGHBOR_QUANTIZER=BINARY)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.index_type == VEC0_INDEX_TYPE_DISKANN);
    assert(col.diskann.quantizer_type == VEC0_DISKANN_QUANTIZER_BINARY);
    sqlite3_free(col.name);
  }

  // Split search_list_size: search and insert
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=binary, search_list_size_search=256, search_list_size_insert=64)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.diskann.search_list_size == 128);  // default (unified)
    assert(col.diskann.search_list_size_search == 256);
    assert(col.diskann.search_list_size_insert == 64);
    sqlite3_free(col.name);
  }

  // Split search_list_size: only search
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=binary, search_list_size_search=200)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_OK);
    assert(col.diskann.search_list_size_search == 200);
    assert(col.diskann.search_list_size_insert == 0);
    sqlite3_free(col.name);
  }

  // Error: cannot mix search_list_size with search_list_size_search
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=binary, search_list_size=128, search_list_size_search=256)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  // Error: cannot mix search_list_size with search_list_size_insert
  {
    const char *input = "emb float[128] INDEXED BY diskann(neighbor_quantizer=binary, search_list_size=128, search_list_size_insert=64)";
    rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == SQLITE_ERROR);
  }

  printf("  All vec0_parse_vector_column_diskann tests passed.\n");
}

void test_diskann_validity_bitmap() {
  printf("Starting %s...\n", __func__);

  unsigned char validity[3]; // 24 bits
  memset(validity, 0, sizeof(validity));

  // All initially invalid
  for (int i = 0; i < 24; i++) {
    assert(diskann_validity_get(validity, i) == 0);
  }
  assert(diskann_validity_count(validity, 24) == 0);

  // Set bit 0
  diskann_validity_set(validity, 0, 1);
  assert(diskann_validity_get(validity, 0) == 1);
  assert(diskann_validity_count(validity, 24) == 1);

  // Set bit 7 (last bit of first byte)
  diskann_validity_set(validity, 7, 1);
  assert(diskann_validity_get(validity, 7) == 1);
  assert(diskann_validity_count(validity, 24) == 2);

  // Set bit 8 (first bit of second byte)
  diskann_validity_set(validity, 8, 1);
  assert(diskann_validity_get(validity, 8) == 1);
  assert(diskann_validity_count(validity, 24) == 3);

  // Set bit 23 (last bit)
  diskann_validity_set(validity, 23, 1);
  assert(diskann_validity_get(validity, 23) == 1);
  assert(diskann_validity_count(validity, 24) == 4);

  // Clear bit 0
  diskann_validity_set(validity, 0, 0);
  assert(diskann_validity_get(validity, 0) == 0);
  assert(diskann_validity_count(validity, 24) == 3);

  // Other bits unaffected
  assert(diskann_validity_get(validity, 7) == 1);
  assert(diskann_validity_get(validity, 8) == 1);

  printf("  All diskann_validity_bitmap tests passed.\n");
}

void test_diskann_neighbor_ids() {
  printf("Starting %s...\n", __func__);

  unsigned char ids[8 * 8]; // 8 slots * 8 bytes each
  memset(ids, 0, sizeof(ids));

  // Set and get slot 0
  diskann_neighbor_id_set(ids, 0, 42);
  assert(diskann_neighbor_id_get(ids, 0) == 42);

  // Set and get middle slot
  diskann_neighbor_id_set(ids, 3, 12345);
  assert(diskann_neighbor_id_get(ids, 3) == 12345);

  // Set and get last slot
  diskann_neighbor_id_set(ids, 7, 99999);
  assert(diskann_neighbor_id_get(ids, 7) == 99999);

  // Slot 0 still correct
  assert(diskann_neighbor_id_get(ids, 0) == 42);

  // Large value
  diskann_neighbor_id_set(ids, 1, INT64_MAX);
  assert(diskann_neighbor_id_get(ids, 1) == INT64_MAX);

  printf("  All diskann_neighbor_ids tests passed.\n");
}

void test_diskann_quantize_binary() {
  printf("Starting %s...\n", __func__);

  // 8-dimensional vector: positive values -> 1, negative/zero -> 0
  float src[8] = {1.0f, -1.0f, 0.5f, 0.0f, -0.5f, 0.1f, -0.1f, 100.0f};
  unsigned char out[1]; // 8 bits = 1 byte

  int rc = diskann_quantize_vector(src, 8, VEC0_DISKANN_QUANTIZER_BINARY, out);
  assert(rc == 0);

  // Expected bits (LSB first within each byte):
  // bit 0: 1.0 > 0 -> 1
  // bit 1: -1.0 > 0 -> 0
  // bit 2: 0.5 > 0 -> 1
  // bit 3: 0.0 > 0 -> 0  (not strictly greater)
  // bit 4: -0.5 > 0 -> 0
  // bit 5: 0.1 > 0 -> 1
  // bit 6: -0.1 > 0 -> 0
  // bit 7: 100.0 > 0 -> 1
  // Expected byte: 1 + 0 + 4 + 0 + 0 + 32 + 0 + 128 = 0b10100101 = 0xA5
  assert(out[0] == 0xA5);

  printf("  All diskann_quantize_binary tests passed.\n");
}

void test_diskann_node_init_sizes() {
  printf("Starting %s...\n", __func__);

  unsigned char *validity, *ids, *qvecs;
  int validitySize, idsSize, qvecsSize;

  // 72 neighbors, binary quantizer, 1024 dims
  int rc = diskann_node_init(72, VEC0_DISKANN_QUANTIZER_BINARY, 1024,
      &validity, &validitySize, &ids, &idsSize, &qvecs, &qvecsSize);
  assert(rc == 0);
  assert(validitySize == 9);      // 72/8
  assert(idsSize == 576);         // 72 * 8
  assert(qvecsSize == 9216);      // 72 * (1024/8)

  // All validity bits should be 0
  assert(diskann_validity_count(validity, 72) == 0);

  sqlite3_free(validity);
  sqlite3_free(ids);
  sqlite3_free(qvecs);

  // 8 neighbors, int8 quantizer, 32 dims
  rc = diskann_node_init(8, VEC0_DISKANN_QUANTIZER_INT8, 32,
      &validity, &validitySize, &ids, &idsSize, &qvecs, &qvecsSize);
  assert(rc == 0);
  assert(validitySize == 1);    // 8/8
  assert(idsSize == 64);        // 8 * 8
  assert(qvecsSize == 256);     // 8 * 32

  sqlite3_free(validity);
  sqlite3_free(ids);
  sqlite3_free(qvecs);

  printf("  All diskann_node_init_sizes tests passed.\n");
}

void test_diskann_node_set_clear_neighbor() {
  printf("Starting %s...\n", __func__);

  unsigned char *validity, *ids, *qvecs;
  int validitySize, idsSize, qvecsSize;

  // 8 neighbors, binary quantizer, 16 dims (2 bytes per qvec)
  int rc = diskann_node_init(8, VEC0_DISKANN_QUANTIZER_BINARY, 16,
      &validity, &validitySize, &ids, &idsSize, &qvecs, &qvecsSize);
  assert(rc == 0);

  // Create a test quantized vector (2 bytes)
  unsigned char test_qvec[2] = {0xAB, 0xCD};

  // Set neighbor at slot 3
  diskann_node_set_neighbor(validity, ids, qvecs, 3,
      42, test_qvec, VEC0_DISKANN_QUANTIZER_BINARY, 16);

  // Verify slot 3 is valid
  assert(diskann_validity_get(validity, 3) == 1);
  assert(diskann_validity_count(validity, 8) == 1);

  // Verify rowid
  assert(diskann_neighbor_id_get(ids, 3) == 42);

  // Verify quantized vector
  const unsigned char *read_qvec = diskann_neighbor_qvec_get(
      qvecs, 3, VEC0_DISKANN_QUANTIZER_BINARY, 16);
  assert(read_qvec[0] == 0xAB);
  assert(read_qvec[1] == 0xCD);

  // Clear slot 3
  diskann_node_clear_neighbor(validity, ids, qvecs, 3,
      VEC0_DISKANN_QUANTIZER_BINARY, 16);
  assert(diskann_validity_get(validity, 3) == 0);
  assert(diskann_neighbor_id_get(ids, 3) == 0);
  assert(diskann_validity_count(validity, 8) == 0);

  sqlite3_free(validity);
  sqlite3_free(ids);
  sqlite3_free(qvecs);

  printf("  All diskann_node_set_clear_neighbor tests passed.\n");
}

void test_diskann_prune_select() {
  printf("Starting %s...\n", __func__);

  // Scenario: 5 candidates, sorted by distance to p
  // Candidates: A(0), B(1), C(2), D(3), E(4)
  // p_distances (already sorted): A=1.0, B=2.0, C=3.0, D=4.0, E=5.0
  //
  // Inter-candidate distances (5x5 matrix):
  //       A     B     C     D     E
  // A   0.0   1.5   3.0   4.0   5.0
  // B   1.5   0.0   1.5   3.0   4.0
  // C   3.0   1.5   0.0   1.5   3.0
  // D   4.0   3.0   1.5   0.0   1.5
  // E   5.0   4.0   3.0   1.5   0.0

  float p_distances[5] = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f};
  float inter[25] = {
    0.0f, 1.5f, 3.0f, 4.0f, 5.0f,
    1.5f, 0.0f, 1.5f, 3.0f, 4.0f,
    3.0f, 1.5f, 0.0f, 1.5f, 3.0f,
    4.0f, 3.0f, 1.5f, 0.0f, 1.5f,
    5.0f, 4.0f, 3.0f, 1.5f, 0.0f,
  };
  int selected[5];
  int count;

  // alpha=1.0, R=3: greedy selection
  // Round 1: Pick A (closest). Prune check:
  //   B: 1.0*1.5 <= 2.0? yes -> pruned
  //   C: 1.0*3.0 <= 3.0? yes -> pruned
  //   D: 1.0*4.0 <= 4.0? yes -> pruned
  //   E: 1.0*5.0 <= 5.0? yes -> pruned
  // Result: only A selected
  {
    int rc = diskann_prune_select(inter, p_distances, 5, 1.0f, 3, selected, &count);
    assert(rc == 0);
    assert(count == 1);
    assert(selected[0] == 1); // A
  }

  // alpha=1.5, R=3: diversity-aware
  // Round 1: Pick A. Prune check:
  //   B: 1.5*1.5=2.25 <= 2.0? no -> keep
  //   C: 1.5*3.0=4.5 <= 3.0? no -> keep
  //   D: 1.5*4.0=6.0 <= 4.0? no -> keep
  //   E: 1.5*5.0=7.5 <= 5.0? no -> keep
  // Round 2: Pick B. Prune check:
  //   C: 1.5*1.5=2.25 <= 3.0? yes -> pruned
  //   D: 1.5*3.0=4.5 <= 4.0? no -> keep
  //   E: 1.5*4.0=6.0 <= 5.0? no -> keep
  // Round 3: Pick D. Done, 3 selected.
  {
    int rc = diskann_prune_select(inter, p_distances, 5, 1.5f, 3, selected, &count);
    assert(rc == 0);
    assert(count == 3);
    assert(selected[0] == 1); // A
    assert(selected[1] == 1); // B
    assert(selected[3] == 1); // D
    assert(selected[2] == 0); // C pruned
    assert(selected[4] == 0); // E not reached
  }

  // R > num_candidates with very high alpha (no pruning): select all
  {
    int rc = diskann_prune_select(inter, p_distances, 5, 100.0f, 10, selected, &count);
    assert(rc == 0);
    assert(count == 5);
  }

  // Empty candidate set
  {
    int rc = diskann_prune_select(NULL, NULL, 0, 1.2f, 3, selected, &count);
    assert(rc == 0);
    assert(count == 0);
  }

  printf("  All diskann_prune_select tests passed.\n");
}

void test_diskann_quantized_vector_byte_size() {
  printf("Starting %s...\n", __func__);

  // Binary quantizer: 1 bit per dimension, so 128 dims = 16 bytes
  assert(diskann_quantized_vector_byte_size(VEC0_DISKANN_QUANTIZER_BINARY, 128) == 16);
  assert(diskann_quantized_vector_byte_size(VEC0_DISKANN_QUANTIZER_BINARY, 8) == 1);
  assert(diskann_quantized_vector_byte_size(VEC0_DISKANN_QUANTIZER_BINARY, 1024) == 128);

  // INT8 quantizer: 1 byte per dimension
  assert(diskann_quantized_vector_byte_size(VEC0_DISKANN_QUANTIZER_INT8, 128) == 128);
  assert(diskann_quantized_vector_byte_size(VEC0_DISKANN_QUANTIZER_INT8, 1) == 1);
  assert(diskann_quantized_vector_byte_size(VEC0_DISKANN_QUANTIZER_INT8, 768) == 768);

  printf("  All diskann_quantized_vector_byte_size tests passed.\n");
}

void test_diskann_config_defaults() {
  printf("Starting %s...\n", __func__);

  // A freshly zero-initialized VectorColumnDefinition should have diskann.enabled == 0
  struct VectorColumnDefinition col;
  memset(&col, 0, sizeof(col));
  assert(col.index_type != VEC0_INDEX_TYPE_DISKANN);
  assert(col.diskann.n_neighbors == 0);
  assert(col.diskann.search_list_size == 0);

  // Verify parsing a normal vector column still works and diskann is not enabled
  {
    const char *input = "embedding float[768]";
    int rc = vec0_parse_vector_column(input, (int)strlen(input), &col);
    assert(rc == 0 /* SQLITE_OK */);
    assert(col.index_type != VEC0_INDEX_TYPE_DISKANN);
    sqlite3_free(col.name);
  }

  printf("  All diskann_config_defaults tests passed.\n");
}

// ======================================================================
// Additional DiskANN unit tests
// ======================================================================

void test_diskann_quantize_int8() {
  printf("Starting %s...\n", __func__);

  // INT8 quantization uses fixed range [-1, 1]:
  //   step = 2.0 / 255.0
  //   out[i] = (i8)((src[i] + 1.0) / step - 128.0)
  float src[4] = {-1.0f, 0.0f, 0.5f, 1.0f};
  unsigned char out[4];

  int rc = diskann_quantize_vector(src, 4, VEC0_DISKANN_QUANTIZER_INT8, out);
  assert(rc == 0);

  int8_t *signed_out = (int8_t *)out;
  // -1.0 -> (0/step) - 128 = -128
  assert(signed_out[0] == -128);
  // 0.0 -> (1.0/step) - 128 ~= 127.5 - 128 ~= -0.5 -> (i8)(-0.5) = 0
  assert(signed_out[1] >= -2 && signed_out[1] <= 2);
  // 0.5 -> (1.5/step) - 128 ~= 191.25 - 128 = 63.25 -> (i8) 63
  assert(signed_out[2] >= 60 && signed_out[2] <= 66);
  // 1.0 -> should be close to 127 (may have float precision issues)
  assert(signed_out[3] >= 126 && signed_out[3] <= 127);

  printf("  All diskann_quantize_int8 tests passed.\n");
}

void test_diskann_quantize_binary_16d() {
  printf("Starting %s...\n", __func__);

  // 16-dimensional vector (2 bytes output)
  float src[16] = {
    1.0f, -1.0f, 0.5f, -0.5f,  // byte 0: bit0=1, bit1=0, bit2=1, bit3=0
    0.1f, -0.1f, 0.0f, 100.0f, // byte 0: bit4=1, bit5=0, bit6=0, bit7=1
    -1.0f, 1.0f, 1.0f, 1.0f,   // byte 1: bit0=0, bit1=1, bit2=1, bit3=1
    -1.0f, -1.0f, 1.0f, -1.0f  // byte 1: bit4=0, bit5=0, bit6=1, bit7=0
  };
  unsigned char out[2];

  int rc = diskann_quantize_vector(src, 16, VEC0_DISKANN_QUANTIZER_BINARY, out);
  assert(rc == 0);

  // byte 0: bits 0,2,4,7 set -> 0b10010101 = 0x95
  assert(out[0] == 0x95);
  // byte 1: bits 1,2,3,6 set -> 0b01001110 = 0x4E
  assert(out[1] == 0x4E);

  printf("  All diskann_quantize_binary_16d tests passed.\n");
}

void test_diskann_quantize_binary_all_positive() {
  printf("Starting %s...\n", __func__);

  float src[8] = {1.0f, 2.0f, 0.1f, 0.001f, 100.0f, 42.0f, 0.5f, 3.14f};
  unsigned char out[1];

  int rc = diskann_quantize_vector(src, 8, VEC0_DISKANN_QUANTIZER_BINARY, out);
  assert(rc == 0);
  assert(out[0] == 0xFF);  // All bits set

  printf("  All diskann_quantize_binary_all_positive tests passed.\n");
}

void test_diskann_quantize_binary_all_negative() {
  printf("Starting %s...\n", __func__);

  float src[8] = {-1.0f, -2.0f, -0.1f, -0.001f, -100.0f, -42.0f, -0.5f, 0.0f};
  unsigned char out[1];

  int rc = diskann_quantize_vector(src, 8, VEC0_DISKANN_QUANTIZER_BINARY, out);
  assert(rc == 0);
  assert(out[0] == 0x00);  // No bits set (all <= 0)

  printf("  All diskann_quantize_binary_all_negative tests passed.\n");
}

void test_diskann_candidate_list_operations() {
  printf("Starting %s...\n", __func__);

  struct DiskannCandidateList list;
  int rc = _test_diskann_candidate_list_init(&list, 5);
  assert(rc == 0);

  // Insert candidates in non-sorted order
  _test_diskann_candidate_list_insert(&list, 10, 3.0f);
  _test_diskann_candidate_list_insert(&list, 20, 1.0f);
  _test_diskann_candidate_list_insert(&list, 30, 2.0f);

  assert(_test_diskann_candidate_list_count(&list) == 3);
  // Should be sorted by distance
  assert(_test_diskann_candidate_list_rowid(&list, 0) == 20);  // dist 1.0
  assert(_test_diskann_candidate_list_rowid(&list, 1) == 30);  // dist 2.0
  assert(_test_diskann_candidate_list_rowid(&list, 2) == 10);  // dist 3.0

  assert(_test_diskann_candidate_list_distance(&list, 0) == 1.0f);
  assert(_test_diskann_candidate_list_distance(&list, 1) == 2.0f);
  assert(_test_diskann_candidate_list_distance(&list, 2) == 3.0f);

  // Deduplication: inserting same rowid with better distance should update
  _test_diskann_candidate_list_insert(&list, 10, 0.5f);
  assert(_test_diskann_candidate_list_count(&list) == 3);  // Same count
  assert(_test_diskann_candidate_list_rowid(&list, 0) == 10);  // Now first
  assert(_test_diskann_candidate_list_distance(&list, 0) == 0.5f);

  // Next unvisited: should be index 0
  int idx = _test_diskann_candidate_list_next_unvisited(&list);
  assert(idx == 0);

  // Mark visited
  _test_diskann_candidate_list_set_visited(&list, 0);
  idx = _test_diskann_candidate_list_next_unvisited(&list);
  assert(idx == 1);  // Skip visited

  // Fill to capacity (5) and try inserting a worse candidate
  _test_diskann_candidate_list_insert(&list, 40, 4.0f);
  _test_diskann_candidate_list_insert(&list, 50, 5.0f);
  assert(_test_diskann_candidate_list_count(&list) == 5);

  // Insert worse than worst -> should be discarded
  int inserted = _test_diskann_candidate_list_insert(&list, 60, 10.0f);
  assert(inserted == 0);
  assert(_test_diskann_candidate_list_count(&list) == 5);

  // Insert better than worst -> should replace worst
  inserted = _test_diskann_candidate_list_insert(&list, 60, 3.5f);
  assert(inserted == 1);
  assert(_test_diskann_candidate_list_count(&list) == 5);

  _test_diskann_candidate_list_free(&list);

  printf("  All diskann_candidate_list_operations tests passed.\n");
}

void test_diskann_visited_set_operations() {
  printf("Starting %s...\n", __func__);

  struct DiskannVisitedSet set;
  int rc = _test_diskann_visited_set_init(&set, 32);
  assert(rc == 0);

  // Empty set
  assert(_test_diskann_visited_set_contains(&set, 1) == 0);
  assert(_test_diskann_visited_set_contains(&set, 100) == 0);

  // Insert and check
  int inserted = _test_diskann_visited_set_insert(&set, 42);
  assert(inserted == 1);
  assert(_test_diskann_visited_set_contains(&set, 42) == 1);
  assert(_test_diskann_visited_set_contains(&set, 43) == 0);

  // Double insert returns 0
  inserted = _test_diskann_visited_set_insert(&set, 42);
  assert(inserted == 0);

  // Insert several
  _test_diskann_visited_set_insert(&set, 1);
  _test_diskann_visited_set_insert(&set, 2);
  _test_diskann_visited_set_insert(&set, 100);
  _test_diskann_visited_set_insert(&set, 999);
  assert(_test_diskann_visited_set_contains(&set, 1) == 1);
  assert(_test_diskann_visited_set_contains(&set, 2) == 1);
  assert(_test_diskann_visited_set_contains(&set, 100) == 1);
  assert(_test_diskann_visited_set_contains(&set, 999) == 1);
  assert(_test_diskann_visited_set_contains(&set, 3) == 0);

  // Sentinel value (rowid 0) should not be insertable
  assert(_test_diskann_visited_set_contains(&set, 0) == 0);
  inserted = _test_diskann_visited_set_insert(&set, 0);
  assert(inserted == 0);

  _test_diskann_visited_set_free(&set);

  printf("  All diskann_visited_set_operations tests passed.\n");
}

void test_diskann_prune_select_single_candidate() {
  printf("Starting %s...\n", __func__);

  float p_distances[1] = {5.0f};
  float inter[1] = {0.0f};
  int selected[1];
  int count;

  int rc = diskann_prune_select(inter, p_distances, 1, 1.0f, 3, selected, &count);
  assert(rc == 0);
  assert(count == 1);
  assert(selected[0] == 1);

  printf("  All diskann_prune_select_single_candidate tests passed.\n");
}

void test_diskann_prune_select_all_identical_distances() {
  printf("Starting %s...\n", __func__);

  float p_distances[4] = {2.0f, 2.0f, 2.0f, 2.0f};
  // All inter-distances are equal too
  float inter[16] = {
    0.0f, 1.0f, 1.0f, 1.0f,
    1.0f, 0.0f, 1.0f, 1.0f,
    1.0f, 1.0f, 0.0f, 1.0f,
    1.0f, 1.0f, 1.0f, 0.0f,
  };
  int selected[4];
  int count;

  // alpha=1.0: pick first, then check if alpha * inter[0][j] <= p_dist[j]
  // 1.0 * 1.0 <= 2.0? yes, so all are pruned after picking the first
  int rc = diskann_prune_select(inter, p_distances, 4, 1.0f, 4, selected, &count);
  assert(rc == 0);
  assert(count >= 1);  // At least one selected

  printf("  All diskann_prune_select_all_identical_distances tests passed.\n");
}

void test_diskann_prune_select_max_neighbors_1() {
  printf("Starting %s...\n", __func__);

  float p_distances[3] = {1.0f, 2.0f, 3.0f};
  float inter[9] = {
    0.0f, 5.0f, 5.0f,
    5.0f, 0.0f, 5.0f,
    5.0f, 5.0f, 0.0f,
  };
  int selected[3];
  int count;

  // R=1: should select exactly 1
  int rc = diskann_prune_select(inter, p_distances, 3, 1.0f, 1, selected, &count);
  assert(rc == 0);
  assert(count == 1);
  assert(selected[0] == 1);  // First (closest) is selected

  printf("  All diskann_prune_select_max_neighbors_1 tests passed.\n");
}

int main() {
  printf("Starting unit tests...\n");
#ifdef SQLITE_VEC_ENABLE_AVX
  printf("SQLITE_VEC_ENABLE_AVX=1\n");
#endif
#ifdef SQLITE_VEC_ENABLE_NEON
  printf("SQLITE_VEC_ENABLE_NEON=1\n");
#endif
#ifdef SQLITE_VEC_ENABLE_RESCORE
  printf("SQLITE_VEC_ENABLE_RESCORE=1\n");
#endif
#if !defined(SQLITE_VEC_ENABLE_AVX) && !defined(SQLITE_VEC_ENABLE_NEON)
  printf("SIMD: none\n");
#endif
  test_vec0_token_next();
  test_vec0_scanner();
  test_vec0_parse_vector_column();
  test_vec0_parse_partition_key_definition();
  test_distance_l2_sqr_float();
  test_distance_cosine_float();
  test_distance_hamming();
#ifdef SQLITE_VEC_ENABLE_RESCORE
  test_rescore_quantize_float_to_bit();
  test_rescore_quantize_float_to_int8();
  test_rescore_quantized_byte_size();
  test_vec0_parse_vector_column_rescore();
#if SQLITE_VEC_ENABLE_IVF
  test_ivf_quantize_int8();
  test_ivf_quantize_binary();
  test_ivf_config_parsing();
#endif
  test_vec0_parse_vector_column_diskann();
  test_diskann_validity_bitmap();
  test_diskann_neighbor_ids();
  test_diskann_quantize_binary();
  test_diskann_node_init_sizes();
  test_diskann_node_set_clear_neighbor();
  test_diskann_prune_select();
  test_diskann_quantized_vector_byte_size();
  test_diskann_config_defaults();
  test_diskann_quantize_int8();
  test_diskann_quantize_binary_16d();
  test_diskann_quantize_binary_all_positive();
  test_diskann_quantize_binary_all_negative();
  test_diskann_candidate_list_operations();
  test_diskann_visited_set_operations();
  test_diskann_prune_select_single_candidate();
  test_diskann_prune_select_all_identical_distances();
  test_diskann_prune_select_max_neighbors_1();
  printf("All unit tests passed.\n");
}


use std::path::Path;

use anyhow::Result;
use xml::reader::ParserConfig;

use crate::hook::Hook;
use crate::hooks::pre_commit_hooks::{FilenamesArgs, hook_filenames, parse_hook_args};
use crate::hooks::run_concurrent_file_checks;
use crate::run::INTERNAL_CONCURRENCY;

pub(crate) async fn check_xml(hook: &Hook, filenames: &[&Path]) -> Result<(i32, Vec<u8>)> {
    let args: FilenamesArgs = parse_hook_args(hook)?;
    run_concurrent_file_checks(
        hook_filenames(&args.filenames, filenames),
        *INTERNAL_CONCURRENCY,
        |filename| check_file(hook.project().relative_path(), filename),
    )
    .await
}

async fn check_file(file_base: &Path, filename: &Path) -> Result<(i32, Vec<u8>)> {
    let content = fs_err::tokio::read(file_base.join(filename)).await?;

    // Parse the whole document once with xml-rs. This is stricter than the upstream Python
    // check-xml (xml.sax/Expat) in two cases:
    // - `<p:root/>` passes upstream because namespace processing is off; xml-rs requires
    //   `xmlns:p`.
    // - `<!DOCTYPE root SYSTEM "x.dtd"><root>&ext;</root>` passes upstream because `ext` may be
    //   declared by the unloaded DTD; xml-rs reports it as unknown.
    // Keeping these differences avoids rewriting or repeatedly parsing the input.
    let parser = ParserConfig::new()
        .allow_multiple_root_elements(false)
        .create_reader(content.as_slice());

    for event in parser {
        if let Err(error) = event {
            let error_message = format!("{}: Failed to xml parse ({error})\n", filename.display());
            return Ok((1, error_message.into_bytes()));
        }
    }

    Ok((0, Vec::new()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::path::PathBuf;
    use tempfile::tempdir;

    async fn create_test_file(
        dir: &tempfile::TempDir,
        name: &str,
        content: &[u8],
    ) -> Result<PathBuf> {
        let file_path = dir.path().join(name);
        fs_err::tokio::write(&file_path, content).await?;
        Ok(file_path)
    }

    #[tokio::test]
    async fn test_valid_xml() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<?xml version="1.0" encoding="UTF-8"?>
<root>
    <element>value</element>
</root>"#;
        let file_path = create_test_file(&dir, "valid.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 0);
        assert!(output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_valid_xml_with_leading_processing_instruction() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<?xml-stylesheet href="style.xsl"?><root/>"#;
        let file_path = create_test_file(&dir, "processing_instruction.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 0);
        assert!(output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_invalid_xml_unclosed_tag() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<?xml version="1.0" encoding="UTF-8"?>
<root>
    <element>value
</root>"#;
        let file_path = create_test_file(&dir, "invalid.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        let output_str = String::from_utf8_lossy(&output);
        assert!(output_str.contains("Failed to xml parse"));
        Ok(())
    }

    #[tokio::test]
    async fn test_invalid_xml_mismatched_tags() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<?xml version="1.0" encoding="UTF-8"?>
<root>
    <element>value</different>
</root>"#;
        let file_path = create_test_file(&dir, "mismatched.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_invalid_xml_syntax_error() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<?xml version="1.0" encoding="UTF-8"?>
<root>
    <element attribute="unclosed value>text</element>
</root>"#;
        let file_path = create_test_file(&dir, "syntax_error.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_invalid_xml_trailing_text() -> Result<()> {
        let dir = tempdir()?;
        let file_path = create_test_file(&dir, "trailing.xml", b"<root/>junk").await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_empty_xml() -> Result<()> {
        let dir = tempdir()?;
        let content = b"";
        let file_path = create_test_file(&dir, "empty.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        let output_str = String::from_utf8_lossy(&output);
        assert!(output_str.contains("no root element found"));
        Ok(())
    }

    #[tokio::test]
    async fn test_whitespace_only_xml() -> Result<()> {
        let dir = tempdir()?;
        let file_path = create_test_file(&dir, "whitespace.xml", b" \n\t").await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_valid_xml_with_attributes() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<?xml version="1.0" encoding="UTF-8"?>
<root xmlns="http://example.com">
    <element id="1" type="test">value</element>
    <element id="2">another value</element>
</root>"#;
        let file_path = create_test_file(&dir, "attributes.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 0);
        assert!(output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_invalid_xml_duplicate_attribute() -> Result<()> {
        let dir = tempdir()?;
        let file_path = create_test_file(
            &dir,
            "duplicate_attribute.xml",
            br#"<root key="1" key="2"/>"#,
        )
        .await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_invalid_xml_element_name() -> Result<()> {
        let dir = tempdir()?;
        let file_path = create_test_file(&dir, "invalid_name.xml", b"<1root/>").await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_valid_xml_with_cdata() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<?xml version="1.0" encoding="UTF-8"?>
<root>
    <element><![CDATA[Some <special> characters & symbols]]></element>
</root>"#;
        let file_path = create_test_file(&dir, "cdata.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 0);
        assert!(output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_valid_xml_with_comments() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<?xml version="1.0" encoding="UTF-8"?>
<root>
    <!-- This is a comment -->
    <element>value</element>
    <!-- Another comment -->
</root>"#;
        let file_path = create_test_file(&dir, "comments.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 0);
        assert!(output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_xml_with_doctype() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<?xml version="1.0" encoding="UTF-8"?>
<!DOCTYPE root SYSTEM "root.dtd">
<root>
    <element>value</element>
</root>"#;
        let file_path = create_test_file(&dir, "doctype.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 0);
        assert!(output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_rejects_unresolved_external_dtd_entity() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<!DOCTYPE html SYSTEM "xhtml.dtd"><html> </html>"#;
        let file_path = create_test_file(&dir, "external_entity.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_invalid_xml_unknown_entity_without_external_dtd() -> Result<()> {
        let dir = tempdir()?;
        let file_path =
            create_test_file(&dir, "unknown_entity.xml", b"<root>&unknown;</root>").await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_xml_with_internal_entity() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<!DOCTYPE root [<!ENTITY value "ok">]>
<root>&value;</root>"#;
        let file_path = create_test_file(&dir, "internal_entity.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 0);
        assert!(output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_valid_utf16_xml() -> Result<()> {
        let dir = tempdir()?;
        let mut content = vec![0xff, 0xfe];
        for code_unit in "<?xml version=\"1.0\" encoding=\"UTF-16\"?><root/>".encode_utf16() {
            content.extend_from_slice(&code_unit.to_le_bytes());
        }
        let file_path = create_test_file(&dir, "utf16.xml", &content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 0);
        assert!(output.is_empty());
        Ok(())
    }

    #[tokio::test]
    async fn test_invalid_xml_no_root() -> Result<()> {
        let dir = tempdir()?;
        let content = br#"<?xml version="1.0" encoding="UTF-8"?>
<element>value</element>
<another>value</another>"#;
        let file_path = create_test_file(&dir, "no_root.xml", content).await?;
        let (code, output) = check_file(Path::new(""), &file_path).await?;
        assert_eq!(code, 1);
        assert!(!output.is_empty());
        Ok(())
    }
}


"""Sitemap template spider."""

from dataclasses import dataclass, field
from gzip import GzipFile
from io import BytesIO
from urllib.parse import urlsplit

from lxml import etree
from protego import Protego

from scrapling.core._types import (
    TYPE_CHECKING,
    Any,
    AsyncGenerator,
    Dict,
    List,
    Optional,
    Union,
)
from scrapling.spiders.links import LinkExtractor
from scrapling.spiders.request import Request
from scrapling.spiders.spider import Spider
from scrapling.spiders.templates.crawler import CrawlRule

if TYPE_CHECKING:
    from scrapling.engines.toolbelt.custom import Response

__all__ = ["SitemapSpider"]

_GZIP_MAGIC = b"\x1f\x8b"
_GUNZIP_MAX_SIZE = 64 * 1024 * 1024  # 64 MiB cap, defends against gzip bombs

@dataclass
class SitemapResult:
    """Parsed sitemap body.

    `urls` holds the entries from a `<urlset>`; `sitemaps` holds child sitemap
    URLs from a `<sitemapindex>` (each of which is fetched recursively).
    """

    urls: List[str] = field(default_factory=list)
    sitemaps: List[str] = field(default_factory=list)

class SitemapSpider(Spider):
    """A Spider that seeds a crawl from sitemap(s), and follows the rules.

    Override `rules()` to return a list of `CrawlRule`s.

    If there are no rules provided, all non-sitemap urls will be redirected to `parse()`, which must be overridden or it will raise `NotImplementedError`.

    :cvar sitemap_urls: Explicit list of sitemap (or robots.txt) URLs to fetch.
    :cvar sitemap_follow: `LinkExtractor` filtering which child sitemaps inside a
        `<sitemapindex>` to descend into. ``None`` means descend into all.
    :cvar sitemap_alternate_links: When enabled, alternate-language URLs are also
        routed through `rules()`.
    """

    sitemap_urls: List[str] = []
    sitemap_follow: Optional[LinkExtractor] = None
    sitemap_alternate_links: bool = False

    def rules(self) -> List[CrawlRule]:
        """Override to define dispatch rules for sitemap URLs."""
        return []

    async def start_requests(self) -> AsyncGenerator[Request, None]:
        if self.sitemap_urls:
            for url in self.sitemap_urls:
                yield Request(url, callback=self._parse_sitemap)
            return

        raise RuntimeError("`SitemapSpider` needs `sitemap_urls` to be set.")

    async def parse(self, response: "Response") -> AsyncGenerator[Union[Dict[str, Any], Request, None], None]:
        """Default callback for processing responses"""
        raise NotImplementedError(f"{self.__class__.__name__} must implement parse() method")
        yield  # Make this a generator for type checkers

    def _robots_body(self, response: "Response") -> List[str]:
        """Extract `Sitemap` directives from a robots.txt body via protego."""
        try:
            text = response.body.decode(response.encoding, errors="replace")
            parser = Protego.parse(text)
        except Exception as e:
            self.logger.warning(f"Failed to parse robots.txt: {e}")
            return []
        return list(parser.sitemaps)

    @staticmethod
    def _decompress(body: bytes, content_type: Optional[str]) -> bytes:
        if (content_type and ("gzip" in content_type.lower())) or (body[:2] == _GZIP_MAGIC):
            out = bytearray()
            with GzipFile(fileobj=BytesIO(body)) as f:
                while chunk := f.read1(8192):
                    out.extend(chunk)
                    if len(out) > _GUNZIP_MAX_SIZE:
                        raise OSError(f"gzip output exceeds {_GUNZIP_MAX_SIZE} bytes")
            return bytes(out)
        return body

    def _extract_urls(self, root: Any) -> List[str]:
        urls: List[str] = []
        for url_el in root:
            if self._get_type(url_el) != "url":
                continue

            for child in url_el:
                name = self._get_type(child)
                if name == "loc" and child.text:
                    urls.append(child.text.strip())
                elif self.sitemap_alternate_links and name == "link":
                    href = child.get("href")
                    if href:
                        urls.append(href.strip())
        return urls

    @staticmethod
    def _get_type(el: Any) -> str:
        return etree.QName(el.tag).localname

    def _sm_body(self, body: bytes, content_type: Optional[str] = None) -> SitemapResult:
        """Parse a sitemap body and return its URLs and any child sitemaps."""
        try:
            body = self._decompress(body, content_type)
        except OSError as e:
            self.logger.warning(f"Failed to decompress sitemap: {e}")
            return SitemapResult()

        try:
            root = etree.fromstring(body)
        except etree.XMLSyntaxError as e:
            self.logger.warning(f"Failed to parse sitemap XML: {e}")
            return SitemapResult()

        root_name = self._get_type(root)
        if root_name == "sitemapindex":
            locs = []
            for sm_el in root:
                if self._get_type(sm_el) == "sitemap":
                    for child in sm_el:
                        if self._get_type(child) == "loc" and child.text:
                            locs.append(child.text.strip())
                            break
            return SitemapResult(sitemaps=locs)
        if root_name == "urlset":
            return SitemapResult(urls=self._extract_urls(root))

        self.logger.warning(f"Unknown sitemap root element: {root_name!r}")
        return SitemapResult()

    async def _parse_sitemap(self, response: "Response") -> AsyncGenerator[Union[Dict[str, Any], Request, None], None]:
        if urlsplit(response.url).path.endswith("/robots.txt"):
            sitemaps = self._robots_body(response)
            if not sitemaps:
                self.logger.warning(f"No Sitemaps found in {response.url}")

            for sitemap_url in sitemaps:
                yield response.follow(sitemap_url, callback=self._parse_sitemap)
            return

        content_type = response.headers.get("content-type") if response.headers else None
        result = self._sm_body(response.body, content_type=content_type)

        # Descend into child sitemaps (apply sitemap_follow filter if present)
        for child_url in result.sitemaps:
            if self.sitemap_follow is not None and not self.sitemap_follow.matches(child_url):
                continue
            yield response.follow(child_url, callback=self._parse_sitemap)

        # Dispatch each URL through rules() (first match wins; unmatched drop unless rules empty)
        rules = self.rules()
        for url in result.urls:
            req = self._dispatch(response, url, rules)
            if req is not None:
                yield req

    @staticmethod
    def _dispatch(response: "Response", url: str, rules: List[CrawlRule]) -> Optional[Request]:
        if not rules:
            return response.follow(url)
        for rule in rules:
            if rule.link_extractor.matches(url):
                req = response.follow(url, callback=rule.callback)
                if rule.priority is not None:
                    req.priority = rule.priority
                if rule.process_request is not None:
                    req = rule.process_request(req, response)
                return req
        return None