"""
Asynchronous Graph API Client.

Provides asynchronous interface for Graph API operations including
SSE-based monitoring for long-running ingestion tasks.
"""

import asyncio
import json
import time
from collections.abc import AsyncGenerator
from typing import Any, cast

import httpx
from httpx_sse import aconnect_sse

from robosystems.logger import logger

from .base import BaseGraphClient
from .config import GraphClientConfig
from .exceptions import (
  GraphTimeoutError,
  GraphTransientError,
)


class GraphClient(BaseGraphClient):
  """Asynchronous client for Graph API operations."""

  def __init__(
    self,
    base_url: str | None = None,
    config: GraphClientConfig | None = None,
    **kwargs,
  ):
    """
    Initialize asynchronous Graph client.

    Args:
        base_url: Base URL for the API
        config: Client configuration
        **kwargs: Additional config overrides
    """
    super().__init__(base_url, config, **kwargs)

    # Configure httpx client limits
    limits = httpx.Limits(
      max_connections=self.config.max_connections,
      max_keepalive_connections=self.config.max_keepalive_connections,
      keepalive_expiry=self.config.keepalive_expiry,
    )

    # Create httpx client
    self.client = httpx.AsyncClient(
      base_url=self.config.base_url,
      timeout=httpx.Timeout(self.config.timeout),
      limits=limits,
      headers=self.config.headers,
      verify=self.config.verify_ssl,
    )

    # Routing metadata (set by factory for debugging)
    self._route_target: str | None = None
    self._graph_id: str | None = None
    self._database_name: str | None = None
    self._instance_id: str | None = None
    self._purpose: str | None = None

  async def __aenter__(self):
    """Async context manager entry."""
    return self

  async def __aexit__(self, exc_type, exc_val, exc_tb):
    """Async context manager exit."""
    await self.close()

  async def close(self):
    """Close the client and cleanup resources."""
    await self.client.aclose()

  async def _execute_with_retry(self, func, *args, **kwargs):
    """
    Execute an async function with retry logic.

    Args:
        func: Async function to execute
        *args: Positional arguments for func
        **kwargs: Keyword arguments for func

    Returns:
        Function result

    Raises:
        GraphAPIError: If all retries fail
    """
    last_error = None

    for attempt in range(self.config.max_retries + 1):
      try:
        # Check circuit breaker
        self._check_circuit_breaker()

        # Execute function
        result = await func(*args, **kwargs)

        # Record success
        self._record_success()

        return result

      except Exception as e:
        last_error = e

        # Convert to appropriate exception type
        if isinstance(e, httpx.TimeoutException):
          last_error = GraphTimeoutError(f"Request timeout: {e}")
        elif isinstance(e, httpx.ConnectError):
          last_error = GraphTransientError(f"Connection error: {e}")
        elif isinstance(e, httpx.RequestError):
          last_error = GraphTransientError(f"Request error: {e}")

        # Check if we should retry
        if not self._should_retry(last_error, attempt):
          self._record_failure()
          raise last_error

        # Calculate retry delay
        if attempt < self.config.max_retries:
          delay = self._calculate_retry_delay(attempt)
          logger.warning(
            f"Request failed (attempt {attempt + 1}/{self.config.max_retries + 1}), "
            f"retrying in {delay:.2f}s: {last_error}"
          )
          await asyncio.sleep(delay)

    # All retries failed
    self._record_failure()
    if last_error is None:
      raise RuntimeError("Retry logic failed without capturing an exception")
    raise last_error

  @staticmethod
  def _classify_operation(path: str) -> str:
    """Classify a Graph API path into an operation name for metrics."""
    # Normalize: /databases/{id}/query -> query, /databases -> create_database, etc.
    parts = path.strip("/").split("/")
    if len(parts) == 0:
      return "unknown"
    if parts[-1] == "query":
      return "query"
    if parts[-1] == "tables" and len(parts) >= 2:
      return "table_ops"
    if parts[-1] == "materialize":
      return "materialize"
    if parts[-1] == "schema":
      return "schema"
    if parts[-1] == "backup":
      return "backup"
    if parts[-1] == "restore":
      return "restore"
    if parts[-1] == "metrics":
      return "metrics"
    if parts[-1] == "health":
      return "health"
    if parts[-1] == "info":
      return "info"
    if parts[-1] == "databases" or (len(parts) == 2 and parts[0] == "databases"):
      return "database_ops"
    if "memory" in parts[-1]:
      return "memory"
    if "fork" in parts[-1] or "copy" in parts[-1]:
      return "fork"
    if "ingest" in path or "s3" in parts[-1]:
      return "ingest"
    if parts[-1] in ("tasks", "task"):
      return "task"
    return "other"

  def _emit_graph_api_metrics(
    self,
    method: str,
    path: str,
    duration: float,
    status_code: int,
    error: bool,
    error_type: str | None,
  ):
    """Emit OTel metrics for a Graph API call."""
    try:
      from robosystems.middleware.otel.metrics import get_endpoint_metrics

      get_endpoint_metrics().record_graph_api_call(
        method=method,
        operation=self._classify_operation(path),
        duration=duration,
        status_code=status_code,
        route_target=self._route_target or "unknown",
        error=error,
        error_type=error_type,
      )
    except Exception:
      pass  # Metrics are best-effort, never break graph API calls

  async def _request(
    self,
    method: str,
    path: str,
    json_data: dict[str, Any] | None = None,
    params: dict[str, Any] | None = None,
    timeout: float | None = None,
    retries: int | None = None,
    headers: dict[str, str] | None = None,
  ) -> httpx.Response:
    """
    Make HTTP request with optional retry logic.

    Args:
        method: HTTP method
        path: API path
        json_data: JSON body
        params: Query parameters
        timeout: Request timeout
        retries: Number of retries. None uses config default, 0 disables retries.
                 Use retries=0 for non-idempotent operations like materialization.
        headers: Additional headers to include in the request.

    Returns:
        Response object
    """
    # Build request kwargs with proper typing
    request_kwargs: dict[str, Any] = {
      "method": method,
      "url": path,
    }

    if json_data is not None:
      request_kwargs["json"] = json_data
    if params is not None:
      request_kwargs["params"] = params
    if timeout is not None:
      request_kwargs["timeout"] = timeout
    if headers is not None:
      request_kwargs["headers"] = headers

    async def make_request():
      # Debug log the request
      logger.debug(f"Making request: {method} {path}")
      if self.client.headers:
        debug_headers = dict(self.client.headers)
        if "X-Graph-API-Key" in debug_headers:
          debug_headers["X-Graph-API-Key"] = (
            debug_headers["X-Graph-API-Key"][:8] + "..."
          )
        logger.debug(f"Client headers: {debug_headers}")

      response = await self.client.request(**request_kwargs)

      # Raise for error status codes
      if response.status_code >= 400:
        try:
          error_data = response.json()
        except Exception:
          error_data = {"detail": response.text}

        error = self._handle_response_error(response.status_code, error_data)
        raise error

      return response

    # Time the entire request including retries
    start_time = time.time()
    status_code = 0
    error_occurred = False
    error_type = None

    try:
      # Skip retry logic if explicitly disabled (retries=0)
      if retries == 0:
        response = await make_request()
      else:
        response = await self._execute_with_retry(make_request)

      status_code = response.status_code
      return response

    except Exception as e:
      error_occurred = True
      status_code = getattr(e, "status_code", 0)
      # Classify error type for metrics
      from .exceptions import GraphClientError, GraphTimeoutError, GraphTransientError

      if isinstance(e, GraphTimeoutError):
        error_type = "timeout"
      elif isinstance(e, GraphTransientError):
        error_type = "transient"
      elif isinstance(e, GraphClientError):
        error_type = "client"
      else:
        error_type = "server"
      raise

    finally:
      duration = time.time() - start_time
      self._emit_graph_api_metrics(
        method=method,
        path=path,
        duration=duration,
        status_code=status_code,
        error=error_occurred,
        error_type=error_type,
      )

  # API Methods

  async def health_check(self) -> dict[str, Any]:
    """Check API health status."""
    response = await self._request("GET", "/health")
    return response.json()

  async def query(
    self,
    cypher: str,
    graph_id: str = "sec",
    parameters: dict[str, Any] | None = None,
    streaming: bool = False,
  ) -> dict[str, Any] | AsyncGenerator[Any]:
    """
    Execute a Cypher query.

    Args:
        cypher: Cypher query to execute
        graph_id: Target graph database ID
        parameters: Query parameters
        streaming: If True, use streaming for large result sets

    Returns:
        Query results (dict for regular, async generator for streaming)
    """
    payload: dict[str, Any] = {"cypher": cypher, "database": graph_id}
    if parameters:
      payload["parameters"] = parameters

    # Add streaming parameter to URL
    params = {"streaming": "true"} if streaming else {}

    if not streaming:
      response = await self._request(
        "POST", f"/databases/{graph_id}/query", json_data=payload, params=params
      )
      # Debug logging for response content
      logger.debug(f"Response status: {response.status_code}")
      logger.debug(f"Response content type: {response.headers.get('content-type')}")
      logger.debug(
        f"Response content length: {len(response.content) if response.content else 'None'}"
      )
      logger.debug(
        f"Response content (first 200 chars): {repr(response.content[:200]) if response.content else 'None'}"
      )

      # Handle truly empty response body (but be more careful about false positives)
      if response.content is None or len(response.content) == 0:
        logger.warning("Received empty response body from Graph API")
        return {"data": [], "columns": [], "row_count": 0}

      try:
        json_result = response.json()
        logger.debug(f"Successfully parsed JSON response: {json_result}")
        return json_result
      except Exception as e:
        logger.error(
          f"Failed to parse response as JSON: {e}, content: {response.content[:100]!r}"
        )
        return {
          "error": f"Invalid JSON response: {e!s}",
          "data": [],
          "columns": [],
          "row_count": 0,
        }

    # For true streaming, use the streaming endpoint
    # This allows the graph database instance to do the chunking
    async def stream_chunks() -> AsyncGenerator[dict[str, Any]]:
      """Stream NDJSON chunks from Graph API server."""
      stream_start = time.time()
      metrics_recorded = False

      try:
        async with self.client.stream(
          "POST",
          f"/databases/{graph_id}/query",
          json=payload,
          params=params,
          timeout=httpx.Timeout(300.0, connect=10.0),  # 5 min stream timeout
        ) as response:
          if response.status_code >= 400:
            error_text = await response.aread()
            try:
              error_data = json.loads(error_text)
            except Exception:
              error_data = {
                "detail": error_text.decode()
                if isinstance(error_text, bytes)
                else error_text
              }
            raise self._handle_response_error(response.status_code, error_data)

          # Record time-to-first-byte as the query metric
          ttfb = time.time() - stream_start
          self._emit_graph_api_metrics(
            method="POST",
            path=f"/databases/{graph_id}/query",
            duration=ttfb,
            status_code=response.status_code,
            error=False,
            error_type=None,
          )
          metrics_recorded = True

          # Stream NDJSON lines
          async for line in response.aiter_lines():
            if line:
              try:
                chunk = json.loads(line)
                yield chunk
              except json.JSONDecodeError as e:
                logger.error(f"Failed to parse NDJSON line: {e}, line: {line[:100]}")
                continue

      except Exception as e:
        if not metrics_recorded:
          from .exceptions import GraphTimeoutError, GraphTransientError

          if isinstance(e, GraphTimeoutError):
            error_type = "timeout"
          elif isinstance(e, GraphTransientError):
            error_type = "transient"
          else:
            error_type = "server"

          self._emit_graph_api_metrics(
            method="POST",
            path=f"/databases/{graph_id}/query",
            duration=time.time() - stream_start,
            status_code=getattr(e, "status_code", 0),
            error=True,
            error_type=error_type,
          )
        raise

    return stream_chunks()

  async def get_info(self) -> dict[str, Any]:
    """
    Get comprehensive cluster information.

    Returns cluster configuration, status, and capabilities.
    """
    response = await self._request("GET", "/info")
    return response.json()

  async def _monitor_task_sse(
    self, sse_path: str, task_id: str, task_type: str, timeout: int
  ) -> dict[str, Any]:
    """
    Monitor any task progress via SSE stream (generic).

    Args:
        sse_path: SSE endpoint path
        task_id: Task ID for logging
        task_type: Type of task (ingestion, backup, restore, etc.)
        timeout: Maximum time to wait (seconds)

    Returns:
        Dict with results or error
    """
    # Use the existing monitoring logic with task_type instead of table_name
    return await self._monitor_ingestion_sse(
      sse_path=sse_path,
      task_id=task_id,
      table_name=task_type,  # Reuse existing param for task type
      timeout=timeout,
    )

  async def _monitor_ingestion_sse(
    self, sse_path: str, task_id: str, table_name: str, timeout: int
  ) -> dict[str, Any]:
    """
    Monitor ingestion progress via SSE stream.

    Args:
        sse_path: SSE endpoint path (e.g., "/databases/sec/ingest/monitor/task123")
        task_id: Task ID for logging
        table_name: Table name for logging
        timeout: Maximum time to wait (seconds)

    Returns:
        Dict with results or error
    """
    start_time = time.time()

    # Build full SSE URL
    sse_url = f"{self.config.base_url}{sse_path}"

    try:
      # Wrap entire SSE monitoring with asyncio.wait_for for hard timeout
      # This ensures we timeout even if the server stops sending events entirely
      # (the previous timeout checks only ran after receiving an event)
      return await asyncio.wait_for(
        self._sse_event_loop(
          sse_url=sse_url,
          task_id=task_id,
          table_name=table_name,
          timeout=timeout,
          start_time=start_time,
        ),
        timeout=timeout + 30,  # Add 30s buffer for connection setup
      )

    except TimeoutError:
      elapsed = time.time() - start_time
      logger.warning(
        f"SSE monitoring hard timeout after {elapsed:.0f}s for {table_name} "
        f"(task {task_id}), falling back to status poll"
      )
      return await self._poll_task_status_fallback(
        task_id, table_name, timeout, start_time
      )
    except Exception as e:
      logger.warning(
        f"SSE monitoring error for task {task_id}: {e}, falling back to status poll"
      )
      return await self._poll_task_status_fallback(
        task_id, table_name, timeout, start_time
      )

  async def _sse_event_loop(
    self,
    sse_url: str,
    task_id: str,
    table_name: str,
    timeout: int,
    start_time: float,
  ) -> dict[str, Any]:
    """
    Internal SSE event processing loop.

    Separated from _monitor_ingestion_sse to allow wrapping with asyncio.wait_for.
    """
    last_heartbeat = start_time
    last_progress_log = start_time

    try:
      # Use a separate client for SSE to avoid interfering with main client
      # Include authentication headers from the main client config
      async with httpx.AsyncClient(
        timeout=httpx.Timeout(timeout),
        headers=self.config.headers,  # Include API key and other headers
      ) as sse_client:
        # Use aconnect_sse for async operations
        async with aconnect_sse(sse_client, "GET", sse_url) as event_source:
          async for sse_event in event_source.aiter_sse():
            current_time = time.time()

            # Parse event data
            try:
              data = json.loads(sse_event.data) if sse_event.data else {}
            except json.JSONDecodeError:
              logger.warning(f"Invalid JSON in SSE event: {sse_event.data}")
              continue

            # Handle different event types
            if sse_event.event == "heartbeat":
              last_heartbeat = current_time
              elapsed = current_time - start_time
              logger.debug(
                f"[Heartbeat] Task {task_id} still running after {elapsed:.0f}s"
              )

            elif sse_event.event == "progress":
              # Log progress every 30 seconds
              if current_time - last_progress_log > 30:
                progress = data.get("progress_percent", 0)
                records = data.get("records_processed", 0)
                estimated = data.get("estimated_records", 0)

                logger.info(
                  f"[Progress] {table_name}: {progress}% complete, "
                  f"{records:,}/{estimated:,} records processed"
                )
                last_progress_log = current_time

            elif sse_event.event == "completed":
              result = data.get("result", {})
              duration = data.get("duration_seconds", 0)
              records_loaded = result.get("records_loaded", 0)

              if records_loaded > 0:
                logger.info(
                  f"[OK] Ingestion completed for {table_name}: "
                  f"{records_loaded:,} records in {duration:.1f}s"
                )
              else:
                # With IGNORE_ERRORS, record count is not available
                logger.info(
                  f"[OK] Ingestion completed for {table_name} in {duration:.1f}s"
                )

              return {
                "status": "completed",
                "task_id": task_id,
                "records_loaded": records_loaded,
                "duration_seconds": duration,
                "result": result,
              }

            elif sse_event.event == "failed":
              error = data.get("error", "Unknown error")
              # "No files found" is expected for optional tables - don't log as error
              if "No files found" in error:
                logger.info(f"No files found for {table_name} (optional table)")
              else:
                logger.error(f"[FAILED] Ingestion failed for {table_name}: {error}")

              return {"status": "failed", "task_id": task_id, "error": error}

            elif sse_event.event == "error":
              error = data.get("error", "Stream error")
              logger.error(f"SSE stream error: {error}")

              return {
                "status": "failed",
                "task_id": task_id,
                "error": f"SSE stream error: {error}",
              }

            # Check for timeout (soft timeout - checked after each event)
            if current_time - start_time > timeout:
              logger.error(f"Ingestion timeout after {timeout}s for {table_name}")
              return {
                "status": "failed",
                "task_id": task_id,
                "error": f"Timeout after {timeout} seconds",
              }

            # Check for stale connection (no heartbeat for 2 minutes)
            if current_time - last_heartbeat > 120:
              logger.warning(
                f"No heartbeat received for 2 minutes for task {task_id}, "
                f"falling back to status poll"
              )
              return await self._poll_task_status_fallback(
                task_id, table_name, timeout, start_time
              )

      # If we exit the loop without a completion event
      logger.warning(
        f"SSE stream ended unexpectedly for task {task_id}, falling back to status poll"
      )
      return await self._poll_task_status_fallback(
        task_id, table_name, timeout, start_time
      )

    except TimeoutError:
      logger.error(f"SSE connection timeout for task {task_id}")
      return {"status": "failed", "task_id": task_id, "error": "SSE connection timeout"}

  async def _poll_task_status_fallback(
    self,
    task_id: str,
    table_name: str,
    timeout: int,
    start_time: float,
  ) -> dict[str, Any]:
    """
    Fallback when SSE stream goes stale: poll task status directly via HTTP GET.

    This handles cases where the SSE stream is starved (e.g., outbound network
    saturated by a large S3 upload) but the task itself is still running fine.
    """
    poll_interval = 5
    max_polls = max(timeout // poll_interval, 60)  # Poll until timeout, minimum 5 min

    for i in range(max_polls):
      # Check overall timeout
      if time.time() - start_time > timeout:
        logger.error(f"Task {task_id} timed out after {timeout}s during status polling")
        return {
          "status": "failed",
          "task_id": task_id,
          "error": f"Timeout after {timeout} seconds",
        }

      try:
        task = await self.get_task_status(task_id)
        status = task.get("status")

        if status == "completed":
          result = task.get("result", {})
          logger.info(
            f"[OK] Task {task_id} completed (detected via status poll after SSE stale)"
          )
          return {
            "status": "completed",
            "task_id": task_id,
            "result": result,
            "duration_seconds": time.time() - start_time,
          }

        if status == "failed":
          error = task.get("error", "Unknown error")
          logger.error(f"[FAILED] Task {task_id} failed: {error}")
          return {"status": "failed", "task_id": task_id, "error": error}

        # Still running — keep polling
        if i == 0:
          logger.info(
            f"Task {task_id} still running (status={status}), "
            f"polling every {poll_interval}s..."
          )

      except Exception as e:
        logger.warning(f"Status poll failed for task {task_id}: {e}")

      await asyncio.sleep(poll_interval)

    # Exhausted polls
    return {
      "status": "failed",
      "task_id": task_id,
      "error": f"Task still running after {max_polls * poll_interval}s of status polling",
    }

  async def list_databases(self) -> dict[str, Any]:
    """List all databases."""
    response = await self._request("GET", "/databases")
    return response.json()

  async def get_database(self, graph_id: str) -> dict[str, Any]:
    """Get specific database information."""
    response = await self._request("GET", f"/databases/{graph_id}")
    return response.json()

  async def create_database(
    self,
    graph_id: str,
    schema_type: str = "entity",
    repository_name: str | None = None,
    custom_schema_ddl: str | None = None,
    is_subgraph: bool = False,
  ) -> dict[str, Any]:
    """
    Create a new database.

    Args:
        graph_id: Database identifier
        schema_type: Type of schema (entity/shared/custom)
        repository_name: Repository name for shared databases
        custom_schema_ddl: Custom DDL for database creation
        is_subgraph: Whether this is a subgraph (bypasses max_databases check)

    Returns:
        Creation result
    """
    payload = {
      "graph_id": graph_id,
      "schema_type": schema_type,
      "is_subgraph": is_subgraph,
    }
    if repository_name:
      payload["repository_name"] = repository_name
    if custom_schema_ddl:
      payload["custom_schema_ddl"] = custom_schema_ddl

    response = await self._request("POST", "/databases", json_data=payload)
    return response.json()

  async def delete_database(
    self, graph_id: str, preserve_duckdb: bool = False, staging_only: bool = False
  ) -> dict[str, Any]:
    """Delete a database.

    Args:
        graph_id: Graph database identifier to delete
        preserve_duckdb: If True, preserve DuckDB staging database for retry scenarios.
            Useful when you want to rebuild LadybugDB from existing staging.
        staging_only: If True, delete only DuckDB staging, preserve LadybugDB graph.
            Useful for re-staging with different settings.

    Note: preserve_duckdb and staging_only are mutually exclusive.
    """
    params = {}
    if preserve_duckdb:
      params["preserve_duckdb"] = "true"
    if staging_only:
      params["staging_only"] = "true"
    response = await self._request("DELETE", f"/databases/{graph_id}", params=params)
    return response.json()

  # =========================================================================
  # Memory Management
  # =========================================================================

  async def boost_memory(self, graph_id: str, target: str = "both") -> dict[str, Any]:
    """
    Boost memory for staging (DuckDB) or materialization (LadybugDB) operations.

    Call this before starting a batch of staging or materialization operations.
    The boost will remain active until restore_memory or release_memory is called.

    Args:
        graph_id: Graph database identifier
        target: Which system to boost - "duckdb", "ladybug", or "both"

    Returns:
        Dictionary with boost status (duckdb_boosted, ladybug_boosted, message)
    """
    response = await self._request(
      "POST",
      f"/databases/{graph_id}/memory/boost",
      json_data={"target": target},
      timeout=30.0,
    )
    return response.json()

  async def restore_memory(self, graph_id: str) -> dict[str, Any]:
    """
    Restore memory limits to defaults after staging/materialization operations.

    This only reconfigures memory limits - connections stay open and buffers
    remain allocated. Use release_memory() to actually free memory.

    Args:
        graph_id: Graph database identifier

    Returns:
        Dictionary with restore status (duckdb_restored, ladybug_restored, message)
    """
    response = await self._request(
      "POST",
      f"/databases/{graph_id}/memory/restore",
      timeout=30.0,
    )
    return response.json()

  async def release_memory(
    self, graph_id: str, target: str = "both", aggressive: bool = True
  ) -> dict[str, Any]:
    """
    Release memory by closing connections and freeing buffers to the OS.

    Unlike restore_memory (which only reconfigures limits), this actually
    closes connections to force the database engines to release their
    buffer memory back to the operating system.

    Call this after staging or materialization operations complete.

    Args:
        graph_id: Graph database identifier
        target: Which system to release - "duckdb", "ladybug", or "both"
        aggressive: For LadybugDB - run GC and malloc_trim for maximum release

    Returns:
        Dictionary with release status and statistics
    """
    response = await self._request(
      "POST",
      f"/databases/{graph_id}/memory/release",
      json_data={"target": target, "aggressive": aggressive},
      timeout=60.0,  # Longer timeout for aggressive cleanup
    )
    return response.json()

  async def memory_status(self, graph_id: str) -> dict[str, Any]:
    """
    Check if memory is currently boosted for a graph.

    Args:
        graph_id: Graph database identifier

    Returns:
        Dictionary with boost status (duckdb_boosted, ladybug_boosted)
    """
    response = await self._request(
      "GET",
      f"/databases/{graph_id}/memory/status",
    )
    return response.json()

  async def ingest(
    self,
    graph_id: str,
    file_path: str | None = None,
    table_name: str | None = None,
    pipeline_run_id: str | None = None,
    bucket: str | None = None,
    files: list[str] | None = None,
    mode: str = "sync",
    priority: int = 5,
    ignore_errors: bool = True,
  ) -> dict[str, Any]:
    """
    Unified data ingestion with flexible execution modes.

    Args:
        graph_id: Target database ID
        file_path: Local file path (for sync mode)
        table_name: Target table name (for sync mode)
        pipeline_run_id: Pipeline ID (for async mode)
        bucket: S3 bucket (for async mode)
        files: S3 file keys (for async mode)
        mode: Execution mode ("sync" or "async")
        priority: Task priority (1-10)
        ignore_errors: Use IGNORE_ERRORS for duplicates

    Returns:
        Ingestion response
    """
    payload = {
      "mode": mode,
      "priority": priority,
      "ignore_errors": ignore_errors,
    }

    if mode == "sync":
      if not file_path or not table_name:
        raise ValueError("Sync mode requires file_path and table_name")
      payload["file_path"] = file_path
      payload["table_name"] = table_name
    else:  # async
      if not pipeline_run_id or not bucket or not files:
        raise ValueError("Async mode requires pipeline_run_id, bucket, and files")
      payload["pipeline_run_id"] = pipeline_run_id
      payload["bucket"] = bucket
      payload["files"] = files

    # Longer timeout for ingestion operations
    timeout = self.config.timeout * 30 if mode == "sync" else self.config.timeout

    response = await self._request(
      "POST",
      f"/databases/{graph_id}/ingest",
      json_data=payload,
      timeout=timeout,
    )
    return response.json()

  async def get_task_status(self, task_id: str) -> dict[str, Any]:
    """Get background task status."""
    response = await self._request("GET", f"/tasks/{task_id}/status")
    return response.json()

  async def list_tasks(
    self, status: str | None = None, limit: int = 100
  ) -> dict[str, Any]:
    """List tasks with optional status filter."""
    params: dict[str, Any] = {"limit": limit}
    if status:
      params["status"] = status

    response = await self._request("GET", "/tasks", params=params)
    return response.json()

  async def cancel_task(self, task_id: str) -> dict[str, Any]:
    """Cancel a pending task."""
    response = await self._request("DELETE", f"/tasks/{task_id}")
    return response.json()

  async def get_queue_info(self) -> dict[str, Any]:
    """Get ingestion queue information."""
    response = await self._request("GET", "/tasks/queue/info")
    return response.json()

  # Additional methods for compatibility with APIRepository

  async def execute_query(
    self, cypher: str, params: dict[str, Any] | None = None
  ) -> list[dict[str, Any]]:
    """
    Execute a query and return data rows (APIRepository compatibility).

    Args:
        cypher: Cypher query
        params: Query parameters

    Returns:
        List of result rows
    """
    # Use _database_name if set (for subgraphs), otherwise fall back to graph_id
    database = getattr(self, "_database_name", None) or self.graph_id or "sec"

    result = cast(dict[str, Any], await self.query(cypher, database, params))
    return result.get("data", [])

  async def execute_single(
    self, cypher: str, params: dict[str, Any] | None = None
  ) -> dict[str, Any] | None:
    """
    Execute a query expecting a single result.

    Args:
        cypher: Cypher query
        params: Query parameters

    Returns:
        Single result or None
    """
    results = await self.execute_query(cypher, params)
    return results[0] if results else None

  async def get_schema(self) -> list[dict[str, Any]]:
    """
    Get database schema information.

    Returns:
        Schema information including tables and properties
    """
    # Get graph ID from instance variable
    graph_id = self.graph_id or "sec"

    response = await self._request("GET", f"/databases/{graph_id}/schema")
    schema_data = response.json()

    # Return just the tables array for compatibility
    return schema_data.get("tables", [])

  async def vector_search(
    self,
    graph_id: str,
    table_name: str,
    embedding: list[float],
    limit: int = 20,
    select: list[str] | None = None,
  ) -> list[dict[str, Any]]:
    """
    Search for similar rows by embedding via the Graph API vector index.

    Calls the LanceDB vector search endpoint on graph instances (master
    or replicas). The index must have been built via vector_build first.

    Args:
        graph_id: Graph database identifier (e.g., "sec")
        table_name: Table whose vector index to search (e.g., "Element")
        embedding: Query embedding vector (e.g., 384-dim float array)
        limit: Maximum number of results to return
        select: Columns to include in results. If None, returns all.

    Returns:
        List of result dicts with matched columns + distance.
    """
    json_data: dict[str, Any] = {"embedding": embedding, "limit": limit}
    if select is not None:
      json_data["select"] = select
    response = await self._request(
      "POST",
      f"/databases/{graph_id}/tables/{table_name}/vector/search",
      json_data=json_data,
      timeout=10.0,
    )
    data = response.json()
    return data.get("results", [])

  async def vector_build(
    self,
    graph_id: str,
    table_name: str,
    query: str,
    memory_limit: str = "8GB",
    timeout: int = 900,
  ) -> dict[str, Any]:
    """
    Build a vector index from a DuckDB query on the graph instance.

    Runs the provided SQL query against the DuckDB staging database and
    builds an IVF-PQ vector index from the results. The query must return
    a "vector" column — all other columns become searchable metadata.

    Args:
        graph_id: Graph database identifier
        table_name: Name for the lance index (e.g., "Element")
        query: DuckDB SQL that selects rows to index. Must include a
            "vector" column (e.g., ``embedding::FLOAT[384] AS vector``).
        memory_limit: DuckDB memory limit for the query
        timeout: Request timeout in seconds (build can take 10+ minutes)

    Returns:
        Dict with row_count, num_partitions, index_size_mb, duration_ms.
    """
    response = await self._request(
      "POST",
      f"/databases/{graph_id}/tables/{table_name}/vector/build",
      json_data={"query": query, "memory_limit": memory_limit},
      timeout=float(timeout),
    )
    return response.json()

  async def vector_export(
    self,
    graph_id: str,
    table_name: str,
    s3_bucket: str | None = None,
    s3_key: str | None = None,
    timeout: int = 300,
  ) -> dict[str, Any]:
    """
    Export a vector index as tar.gz, optionally uploading to S3.

    When s3_bucket and s3_key are provided, the Graph API instance uploads
    the tar.gz directly to S3 (required because the caller may not have
    access to the instance's filesystem).

    Args:
        graph_id: Graph database identifier
        table_name: Table whose index to export
        s3_bucket: S3 bucket to upload to (optional)
        s3_key: S3 object key (optional)
        timeout: Request timeout in seconds

    Returns:
        Dict with size_mb, duration_ms, and s3_uri if uploaded.
    """
    json_data: dict[str, Any] = {}
    if s3_bucket and s3_key:
      json_data["s3_bucket"] = s3_bucket
      json_data["s3_key"] = s3_key
    response = await self._request(
      "POST",
      f"/databases/{graph_id}/tables/{table_name}/vector/export",
      json_data=json_data if json_data else None,
      timeout=float(timeout),
    )
    return response.json()

  async def install_schema(
    self,
    graph_id: str,
    base_schema: str = "base",
    extensions: list[str] | None = None,
    custom_ddl: str | None = None,
  ) -> dict[str, Any]:
    """
    Install or update database schema.

    Args:
        graph_id: Target database ID
        base_schema: Base schema type (default: "base")
        extensions: List of schema extensions to install
        custom_ddl: Custom DDL statements to execute

    Returns:
        Schema installation result
    """
    # The API expects a 'type' field in the request
    # When using base_schema and extensions, the type should be 'custom'
    # When using custom_ddl, the type should be 'ddl'
    if custom_ddl:
      payload = {"type": "ddl", "ddl": custom_ddl}
    else:
      payload = {
        "type": "custom",
        "metadata": {"base_schema": base_schema, "extensions": extensions or []},
      }

    response = await self._request(
      "POST", f"/databases/{graph_id}/schema", json_data=payload
    )
    return response.json()

  # Additional endpoints not in original API

  async def export_database(self, graph_id: str) -> bytes:
    """
    Export a database file.

    Args:
        graph_id: Database to export

    Returns:
        Database file contents as bytes
    """
    response = await self._request("GET", f"/databases/{graph_id}/backup")
    return response.content

  async def get_database_info(self, graph_id: str) -> dict[str, Any]:
    """
    Get comprehensive database information and statistics.

    Args:
        graph_id: Database to get information for

    Returns:
        Database information including size, schema, and metadata
    """
    response = await self._request("GET", f"/databases/{graph_id}")
    return response.json()

  async def get_database_metrics(self, graph_id: str) -> dict[str, Any]:
    """
    Get metrics for a specific database (optimized for billing).

    Args:
        graph_id: Database to get metrics for

    Returns:
        Database metrics including size, counts, and timestamps
    """
    response = await self._request("GET", f"/databases/{graph_id}/metrics")
    return response.json()

  async def get_metrics(self) -> dict[str, Any]:
    """
    Get comprehensive metrics for the entire cluster node.

    Returns:
        System, database, query, and ingestion metrics for monitoring
    """
    response = await self._request("GET", "/metrics")
    return response.json()

  # Helper methods for common operations

  async def database_exists(self, graph_id: str) -> bool:
    """Check if a database exists."""
    try:
      await self.get_database(graph_id)
      return True
    except Exception as e:
      # Check if it's an HTTP exception with 404 status
      if hasattr(e, "status_code") and getattr(e, "status_code", None) == 404:
        return False
      raise

  async def ensure_database_exists(
    self, graph_id: str, schema_type: str = "entity"
  ) -> None:
    """Ensure a database exists, creating it if necessary."""
    if not await self.database_exists(graph_id):
      await self.create_database(graph_id, schema_type)

  async def swap_database(
    self, graph_id: str, lock_token: str | None = None
  ) -> dict[str, Any]:
    """Promote a WIP database to active (blue-green swap).

    Requires that {graph_id}-wip.lbug exists on the Graph API node.

    Args:
        graph_id: Base graph ID (not the -wip variant).
        lock_token: If the caller already holds the materialization lock,
            pass the token so the swap endpoint doesn't re-acquire it.

    Returns:
        Swap result with status and message.
    """
    headers: dict[str, str] = {}
    if lock_token:
      headers["X-Materialization-Lock-Token"] = lock_token
    response = await self._request(
      "POST",
      f"/databases/{graph_id}/swap",
      headers=headers or None,
    )
    return response.json()

  async def execute_ddl(self, ddl: str, graph_id: str | None = None) -> dict[str, Any]:
    """
    Execute DDL (Data Definition Language) statements.

    This is useful for creating tables, relationships, etc.

    Args:
        ddl: DDL statement (CREATE NODE TABLE, CREATE REL TABLE, etc.)
        graph_id: Target graph ID (uses self.graph_id if not provided)

    Returns:
        Query result
    """
    target_graph = graph_id or self.graph_id or "sec"
    return cast(dict[str, Any], await self.query(ddl, target_graph))

  async def node_exists(
    self, label: str, filters: dict[str, Any] | None = None
  ) -> bool:
    """
    Check if a node exists with the given label and filters.

    Args:
        label: Node label to check
        filters: Optional filters to match against

    Returns:
        True if node exists, False otherwise
    """
    database = self.graph_id or "sec"

    # Build WHERE clause from filters
    where_clause = ""
    params = {}

    if filters:
      conditions = []
      for key, value in filters.items():
        param_name = f"param_{key}"
        conditions.append(f"n.{key} = ${param_name}")
        params[param_name] = value

      if conditions:
        where_clause = f"WHERE {' AND '.join(conditions)}"

    # Build query to check existence
    cypher = f"""
      MATCH (n:{label})
      {where_clause}
      RETURN COUNT(n) > 0 AS exists
      LIMIT 1
    """

    result = cast(dict[str, Any], await self.query(cypher, database, params))
    data = result.get("data", [])

    if data and len(data) > 0:
      return data[0].get("exists", False)

    return False

  async def create_backup(
    self,
    graph_id: str,
    backup_format: str = "full_dump",
    compression: bool = True,
    encryption: bool = False,
    backup_type: str = "standard",
    s3_destination: dict[str, str] | None = None,
    checkpoint: bool = True,
    vacuum: bool = False,
  ) -> dict[str, Any]:
    """
    Create a backup of a database.

    Args:
        graph_id: Graph database identifier
        backup_format: Backup format (only 'full_dump' supported)
        compression: Enable compression
        encryption: Enable encryption
        backup_type: "standard", "replica", or "shared_repository"
        s3_destination: Target S3 location (required for replica/shared_repository)
        checkpoint: Run CHECKPOINT before backup
        vacuum: Run VACUUM to compact database before backup (DuckDB only)

    Returns:
        Task information including task_id and monitor_url
    """
    payload = {
      "backup_format": backup_format,
      "compression": compression,
      "encryption": encryption,
      "backup_type": backup_type,
      "checkpoint": checkpoint,
      "vacuum": vacuum,
    }
    if s3_destination:
      payload["s3_destination"] = s3_destination

    response = await self.client.post(
      f"/databases/{graph_id}/backup",
      json=payload,
      headers=self.config.headers,
    )
    response.raise_for_status()
    return response.json()

  async def backup_with_sse(
    self,
    graph_id: str,
    backup_format: str = "full_dump",
    compression: bool = True,
    encryption: bool = False,
    timeout: int = 3600,  # 1 hour default
    backup_type: str = "standard",
    s3_destination: dict[str, str] | None = None,
    checkpoint: bool = True,
    vacuum: bool = False,
  ) -> dict[str, Any]:
    """
    Create a backup and monitor via SSE.

    Args:
        graph_id: Graph database identifier
        backup_format: Backup format (only 'full_dump' supported)
        compression: Enable compression
        encryption: Enable encryption
        timeout: Maximum time to wait for completion (seconds)
        backup_type: "standard", "replica", or "shared_repository"
        s3_destination: Target S3 location (required for replica/shared_repository)
        checkpoint: Run CHECKPOINT before backup
        vacuum: Run VACUUM to compact database before backup (DuckDB only)

    Returns:
        Dict with backup results:
        - status: "completed" or "failed"
        - backup_size_mb: Size of backup (if successful)
        - duration_seconds: Total time taken
        - error: Error message (if failed)
    """
    try:
      # Start the backup task
      logger.info(f"Starting backup for database {graph_id}")

      start_response = await self.create_backup(
        graph_id=graph_id,
        backup_format=backup_format,
        compression=compression,
        encryption=encryption,
        backup_type=backup_type,
        s3_destination=s3_destination,
        checkpoint=checkpoint,
        vacuum=vacuum,
      )

      task_id = start_response["task_id"]
      monitor_url = start_response.get("monitor_url")

      if not monitor_url:
        # Fallback for compatibility
        monitor_url = f"/tasks/{task_id}/monitor"

      logger.info(f"Started backup task {task_id}, monitoring via SSE...")

      # Monitor via SSE using the generic monitor
      return await self._monitor_task_sse(
        sse_path=monitor_url, task_id=task_id, task_type="backup", timeout=timeout
      )

    except Exception as e:
      logger.error(f"Failed to start/monitor backup: {e}")
      return {"status": "failed", "error": str(e)}

  async def download_backup(
    self,
    graph_id: str,
  ) -> dict[str, Any]:
    """
    Download the current database as a backup.

    Args:
        graph_id: Graph database identifier

    Returns:
        Dict containing backup_data and metadata
    """
    response = await self.client.post(
      f"/databases/{graph_id}/backup-download",
      headers=self.config.headers,
    )
    response.raise_for_status()

    return {
      "backup_data": response.content,
      "size_bytes": int(response.headers.get("X-Backup-Size", len(response.content))),
      "database": response.headers.get("X-Database", graph_id),
      "format": response.headers.get("X-Backup-Format", "full_dump"),
    }

  async def restore_backup(
    self,
    graph_id: str,
    s3_bucket: str,
    s3_key: str,
    create_system_backup: bool = True,
    force_overwrite: bool = False,
    encrypted: bool = True,
    compressed: bool = True,
  ) -> dict[str, Any]:
    """
    Restore a database from S3 backup.

    Args:
        graph_id: Graph database identifier
        s3_bucket: S3 bucket containing the backup
        s3_key: S3 key path to the backup
        create_system_backup: Create system backup before restore
        force_overwrite: Force overwrite existing database
        encrypted: Whether the backup is encrypted
        compressed: Whether the backup is compressed

    Returns:
        Task information including task_id and monitor_url
    """
    data = {
      "s3_bucket": s3_bucket,
      "s3_key": s3_key,
      "create_system_backup": str(create_system_backup).lower(),
      "force_overwrite": str(force_overwrite).lower(),
      "encrypted": str(encrypted).lower(),
      "compressed": str(compressed).lower(),
    }

    response = await self.client.post(
      f"/databases/{graph_id}/restore",
      data=data,
    )
    response.raise_for_status()
    return response.json()

  async def restore_with_sse(
    self,
    graph_id: str,
    s3_bucket: str,
    s3_key: str,
    create_system_backup: bool = True,
    force_overwrite: bool = False,
    encrypted: bool = True,
    compressed: bool = True,
    timeout: int = 3600,  # 1 hour default
  ) -> dict[str, Any]:
    """
    Restore a database from S3 backup and monitor via SSE.

    This method is designed for long-running restore tasks. It uses
    Server-Sent Events to receive real-time progress updates.

    Args:
        graph_id: Graph database identifier
        s3_bucket: S3 bucket containing the backup
        s3_key: S3 key path to the backup
        create_system_backup: Create system backup before restore
        force_overwrite: Force overwrite existing database
        encrypted: Whether the backup is encrypted
        compressed: Whether the backup is compressed
        timeout: Maximum time to wait for completion (seconds)

    Returns:
        Dict with restore results:
        - status: "completed" or "failed"
        - duration_seconds: Total time taken
        - error: Error message (if failed)
    """
    try:
      # Step 1: Start the restore task
      logger.info(f"Starting restore for database {graph_id} from {s3_bucket}/{s3_key}")

      restore_response = await self.restore_backup(
        graph_id=graph_id,
        s3_bucket=s3_bucket,
        s3_key=s3_key,
        create_system_backup=create_system_backup,
        force_overwrite=force_overwrite,
        encrypted=encrypted,
        compressed=compressed,
      )

      task_id = restore_response["task_id"]
      monitor_url = restore_response.get("monitor_url", f"/tasks/{task_id}/monitor")

      logger.info(f"Started restore task {task_id}, monitoring via SSE...")

      # Step 2: Monitor via SSE
      return await self._monitor_task_sse(
        sse_path=monitor_url, task_id=task_id, task_type="restore", timeout=timeout
      )

    except Exception as e:
      logger.error(f"Failed to start/monitor restore: {e}")
      return {"status": "failed", "error": str(e)}

  # DuckDB Table Management Methods

  async def create_table(
    self,
    graph_id: str,
    table_name: str,
    s3_pattern: str | list[str],
    file_id_map: dict[str, str] | None = None,
    null_columns: list[str] | None = None,
    timeout: int = 1800,  # 30 minutes default for large file sets
  ) -> dict[str, Any]:
    """
    Create a DuckDB staging table from S3 files with SSE monitoring.

    This method starts a background table creation task and monitors it via SSE
    until completion. This allows creating tables from thousands of S3 files
    without HTTP timeout issues.

    Args:
        graph_id: Graph database identifier
        table_name: Name for the table
        s3_pattern: S3 glob pattern (string) or list of S3 file paths
        file_id_map: Optional map of s3_key -> file_id for provenance tracking
        timeout: Maximum time to wait for completion (seconds), default 30 min

    Returns:
        Dict with creation results:
        - status: "completed" or "failed"
        - result: Table creation details (if successful)
        - duration_seconds: Total time taken
        - error: Error message (if failed)
    """
    try:
      # Step 1: Start the background table creation task
      file_count = len(s3_pattern) if isinstance(s3_pattern, list) else 1
      logger.info(
        f"Starting table creation for {table_name} ({file_count} {'files' if file_count > 1 else 'pattern'})"
      )

      json_data: dict[str, Any] = {
        "graph_id": graph_id,
        "table_name": table_name,
        "s3_pattern": s3_pattern,
        "timeout_seconds": timeout,  # Pass timeout to server for background task
      }
      if file_id_map is not None:
        json_data["file_id_map"] = file_id_map
      if null_columns is not None:
        json_data["null_columns"] = null_columns

      start_response = await self._request(
        "POST",
        f"/databases/{graph_id}/tables",
        json_data=json_data,
        timeout=30.0,  # Short timeout for starting the task
        retries=0,  # Non-idempotent: retries could start duplicate staging tasks
      )

      start_data = start_response.json()
      task_id = start_data["task_id"]
      sse_path = start_data["sse_url"]

      logger.info(f"Started staging task {task_id}, monitoring via SSE...")

      # Step 2: Monitor via SSE
      return await self._monitor_task_sse(
        sse_path=sse_path, task_id=task_id, task_type="staging", timeout=timeout
      )

    except Exception as e:
      logger.error(f"Failed to start/monitor table creation: {e}")
      return {"status": "failed", "error": str(e)}

  async def insert_into_table(
    self,
    graph_id: str,
    table_name: str,
    s3_pattern: str | list[str],
    timeout: int = 1800,  # 30 minutes default for large file sets
    deduplicate: bool = True,
    null_columns: list[str] | None = None,
    file_id_map: dict[str, str] | None = None,
  ) -> dict[str, Any]:
    """
    Insert data into an existing DuckDB staging table with SSE monitoring.

    This method starts a background table insert task and monitors it via SSE
    until completion. Used for incremental data append to existing tables.

    Deduplication uses NOT EXISTS on the dedup key (identifier for nodes,
    src/dst for relationships). Safe for tables with FLOAT[384] embeddings.

    Prerequisites:
    - Table must already exist (created via create_table)
    - Schema must be compatible with the new files

    Args:
        graph_id: Graph database identifier
        table_name: Name of the existing table
        s3_pattern: S3 glob pattern (string) or list of S3 file paths
        timeout: Maximum time to wait for completion (seconds), default 30 min
        deduplicate: Dedup rows using NOT EXISTS on dedup key. Default True.

    Returns:
        Dict with insert results:
        - status: "completed" or "failed"
        - result: Table insert details (if successful)
        - duration_seconds: Total time taken
        - error: Error message (if failed)
    """
    try:
      # Step 1: Start the background table insert task
      file_count = len(s3_pattern) if isinstance(s3_pattern, list) else 1
      logger.info(
        f"Starting table insert for {table_name} ({file_count} {'files' if file_count > 1 else 'pattern'})"
      )

      json_data: dict[str, Any] = {
        "graph_id": graph_id,
        "table_name": table_name,
        "s3_pattern": s3_pattern,
        "deduplicate": deduplicate,
        "timeout_seconds": timeout,  # Pass timeout to server for background task
      }
      if null_columns is not None:
        json_data["null_columns"] = null_columns
      if file_id_map is not None:
        json_data["file_id_map"] = file_id_map

      start_response = await self._request(
        "POST",
        f"/databases/{graph_id}/tables/{table_name}/insert",
        json_data=json_data,
        timeout=30.0,  # Short timeout for starting the task
        retries=0,  # Non-idempotent: retries could insert duplicate data
      )

      start_data = start_response.json()
      task_id = start_data["task_id"]
      sse_path = start_data["sse_url"]

      logger.info(f"Started insert task {task_id}, monitoring via SSE...")

      # Step 2: Monitor via SSE
      return await self._monitor_task_sse(
        sse_path=sse_path, task_id=task_id, task_type="insert", timeout=timeout
      )

    except Exception as e:
      logger.error(f"Failed to start/monitor table insert: {e}")
      return {"status": "failed", "error": str(e)}

  async def list_tables(self, graph_id: str) -> list[dict[str, Any]]:
    """
    List all DuckDB staging tables for a graph.

    Args:
        graph_id: Graph database identifier

    Returns:
        List of table info dictionaries
    """
    response = await self._request("GET", f"/databases/{graph_id}/tables")
    return response.json()

  async def query_table(
    self,
    graph_id: str,
    sql: str,
    parameters: list[Any] | None = None,
    timeout: float | None = None,
  ) -> dict[str, Any]:
    """
    Execute SQL query on DuckDB staging tables.

    Args:
        graph_id: Graph database identifier
        sql: SQL query to execute
        parameters: Optional query parameters for safe value substitution
        timeout: Optional request timeout in seconds. Defaults to client config timeout.
                 Use longer timeout for DDL operations like CREATE TABLE AS SELECT.

    Returns:
        Query results with columns and rows
    """
    json_data = {"graph_id": graph_id, "sql": sql}
    if parameters is not None:
      json_data["parameters"] = parameters

    response = await self._request(
      "POST",
      f"/databases/{graph_id}/tables/query",
      json_data=json_data,
      timeout=timeout,
    )
    return response.json()

  async def delete_table(self, graph_id: str, table_name: str) -> dict[str, Any]:
    """
    Delete a DuckDB staging table.

    Args:
        graph_id: Graph database identifier
        table_name: Table name to delete

    Returns:
        Deletion response
    """
    response = await self._request(
      "DELETE", f"/databases/{graph_id}/tables/{table_name}"
    )
    return response.json()

  async def delete_file_data(
    self, graph_id: str, table_name: str, file_id: str
  ) -> dict[str, Any]:
    """
    Delete rows from a DuckDB table by file_id.

    Args:
        graph_id: Graph database identifier
        table_name: Table name to delete from
        file_id: File ID to delete rows for

    Returns:
        Deletion response with rows_deleted count
    """
    response = await self._request(
      "DELETE", f"/databases/{graph_id}/tables/{table_name}/files/{file_id}"
    )
    return response.json()

  async def materialize_table(
    self,
    graph_id: str,
    table_name: str,
    ignore_errors: bool = True,
    file_ids: list[str] | None = None,
    batch_num: int | None = None,
    num_batches: int | None = None,
    materialize_embeddings: bool = False,
    timeout: float = 600.0,
    source_graph_id: str | None = None,
  ) -> dict[str, Any]:
    """
    Materialize a DuckDB staging table into the graph database.

    Supports both selective materialization (filtering by file_ids) and full
    materialization (copying entire table). For large tables, use hash-based
    batching (batch_num/num_batches) for deterministic partitioning.

    Args:
        graph_id: Graph database identifier
        table_name: Table name to materialize
        ignore_errors: Continue on row errors
        file_ids: Optional list of file IDs to materialize. If None, materializes all rows.
        batch_num: Current batch number (0-indexed) for hash-based batching.
        num_batches: Total number of batches. Each row goes to one batch based on hash(key) % num_batches.
        materialize_embeddings: Include embedding columns in materialization. Default false.
        timeout: Request timeout in seconds. Default 600s (10 min) for large bulk operations.
        source_graph_id: Read DuckDB staging from this graph instead of graph_id.
            Used by blue-green materialization.

    Returns:
        Materialization response with rows materialized and timing
    """
    json_data: dict[str, Any] = {"ignore_errors": ignore_errors}

    if file_ids is not None:
      json_data["file_ids"] = file_ids

    if batch_num is not None and num_batches is not None:
      json_data["batch_num"] = batch_num
      json_data["num_batches"] = num_batches

    if materialize_embeddings:
      json_data["materialize_embeddings"] = True

    if source_graph_id is not None:
      json_data["source_graph_id"] = source_graph_id

    response = await self._request(
      "POST",
      f"/databases/{graph_id}/tables/{table_name}/materialize",
      json_data=json_data,
      timeout=timeout,
      retries=0,  # Non-idempotent: retries cause duplicate data
    )
    return response.json()

  async def build_vector_index(
    self,
    graph_id: str,
    table_name: str,
    backend: str = "hnsw",
    column: str = "embedding",
    query: str | None = None,
    timeout: float = 600.0,
  ) -> dict[str, Any]:
    """Build or rebuild a vector index on a table.

    Args:
        graph_id: Graph database identifier
        table_name: Table to index
        backend: "hnsw" for LadybugDB HNSW, "lance" for LanceDB IVF-PQ
        column: Embedding column name (hnsw only)
        query: DuckDB SQL query (required for lance, ignored for hnsw)
        timeout: Request timeout in seconds

    Returns:
        Build response with row_count, duration_ms, backend
    """
    json_data: dict[str, Any] = {"backend": backend}
    if backend == "hnsw":
      json_data["column"] = column
    elif query:
      json_data["query"] = query

    response = await self._request(
      "POST",
      f"/databases/{graph_id}/tables/{table_name}/vector/build",
      json_data=json_data,
      timeout=timeout,
    )
    return response.json()

  async def fork_from_parent(
    self,
    parent_graph_id: str,
    subgraph_id: str,
    tables: list[str] | None = None,
    ignore_errors: bool = True,
  ) -> dict[str, Any]:
    """
    Fork data from parent graph's DuckDB directly into subgraph's LadybugDB.

    This operation:
    1. Attaches parent graph's DuckDB staging database
    2. Copies specified tables (or all tables) from parent DuckDB to subgraph LadybugDB
    3. Runs on the same EC2 instance where both DuckDB and LadybugDB databases live

    Args:
        parent_graph_id: Parent graph to copy data from
        subgraph_id: Subgraph to copy data to
        tables: List of table names to copy (empty list = all tables)
        ignore_errors: Continue ingestion on row errors

    Returns:
        Fork response with tables copied, row counts, and timing
    """
    response = await self._request(
      "POST",
      f"/databases/{subgraph_id}/tables/{subgraph_id}/fork-from/{parent_graph_id}",
      json_data={
        "tables": tables or [],
        "ignore_errors": ignore_errors,
      },
    )
    return response.json()

  # Migration endpoints

  async def migration_export(
    self, source_version: str, target_version: str
  ) -> dict[str, Any]:
    """
    Start migration export on this instance.

    Args:
        source_version: Current LadybugDB version
        target_version: Target LadybugDB version

    Returns:
        Task ID and monitor URL
    """
    response = await self._request(
      "POST",
      "/migration/export",
      params={"source_version": source_version, "target_version": target_version},
      timeout=30.0,
      retries=0,
    )
    return response.json()

  async def migration_import(self) -> dict[str, Any]:
    """
    Start migration import on this instance.

    Returns:
        Task ID and monitor URL
    """
    response = await self._request(
      "POST",
      "/migration/import",
      timeout=30.0,
      retries=0,
    )
    return response.json()

  async def migration_status(self) -> dict[str, Any]:
    """
    Check migration status on this instance.

    Returns:
        Migration status including pending flag, manifest, and pre-migration files
    """
    response = await self._request("GET", "/migration/status")
    return response.json()

  async def migration_cleanup(self) -> dict[str, Any]:
    """
    Delete .pre-migration files on this instance after verifying migration success.

    System backups in S3 remain as the safety net.

    Returns:
        Cleanup result with files deleted and bytes freed
    """
    response = await self._request("POST", "/migration/cleanup")
    return response.json()
