import { RecordBatchReader } from "apache-arrow";
import type { GraphMeta } from "./selection_queries";

export type ArrowGraphLoadResult = {
  meta: GraphMeta | null;
  vertexIds: Uint32Array;
  xs: Float32Array;
  ys: Float32Array;
  edgeData: Uint32Array;
  routeLineVertices: Float32Array;
  labels: (string | undefined)[];
  iris: string[];
};

type ArrowBatchLog = (phase: string, extra?: Record<string, unknown>) => void;

type ArrowLikeVector = {
  data?: Array<{ valueOffsets?: ArrayLike<number | bigint> }>;
  getChildAt?: (index: number) => ArrowLikeVector | null;
  get?: (index: number) => unknown;
  toArray?: () => ArrayLike<number>;
};

const graphTransportVersion = "1";

export async function readGraphArrow(response: Response, logPhase?: ArrowBatchLog): Promise<ArrowGraphLoadResult> {
  validateArrowResponse(response);

  const reader = await openArrowReader(response);

  let meta: GraphMeta | null = null;
  let vertexIds: Uint32Array | null = null;
  let xs: Float32Array | null = null;
  let ys: Float32Array | null = null;
  let edgeData: Uint32Array | null = null;
  let routeLineVertices: Float32Array | null = null;
  let labels: (string | undefined)[] | null = null;
  let iris: string[] | null = null;

  let batchIndex = 0;
  for await (const batch of reader) {
    const batchStartedAt = performance.now();

    switch (batchIndex) {
      case 0: {
        meta = decodeMetaBatch(batch);
        const nodeCount = typeof meta.nodes === "number" ? meta.nodes : 0;
        const edgeCount = typeof meta.edges === "number" ? meta.edges : 0;
        const routeLineSegments = readScalarNumber(batch, "meta_route_line_segments") ?? 0;

        vertexIds = new Uint32Array(nodeCount);
        xs = new Float32Array(nodeCount);
        ys = new Float32Array(nodeCount);
        edgeData = new Uint32Array(edgeCount * 2);
        routeLineVertices = new Float32Array(routeLineSegments * 4);
        labels = new Array<string | undefined>(nodeCount);
        iris = new Array<string>(nodeCount);

        logPhase?.("arrow_meta_batch_decoded", {
          batch_index: batchIndex,
          rows: batch.numRows,
          node_count: nodeCount,
          edge_count: edgeCount,
          route_line_segments: routeLineSegments,
          decode_ms: Math.round(performance.now() - batchStartedAt),
        });
        break;
      }
      case 1: {
        ensureAllocated(meta, vertexIds, xs, ys, edgeData, routeLineVertices, labels, iris);
        const nodeIDs = readUint32List(batch, "node_id");
        const nodeXs = readFloat32List(batch, "node_x");
        const nodeYs = readFloat32List(batch, "node_y");
        const nodeIRIs = readStringList(batch, "node_iri");
        const nodeLabels = readNullableStringList(batch, "node_label");

        if (nodeIDs.length !== vertexIds.length || nodeXs.length !== xs.length || nodeYs.length !== ys.length) {
          throw new Error("Arrow node batch length mismatch");
        }
        if (nodeIRIs.length !== iris.length || nodeLabels.length !== labels.length) {
          throw new Error("Arrow node metadata batch length mismatch");
        }

        vertexIds.set(nodeIDs);
        xs.set(nodeXs);
        ys.set(nodeYs);
        for (let i = 0; i < iris.length; i++) {
          iris[i] = nodeIRIs[i] ?? "";
          labels[i] = nodeLabels[i];
        }

        logPhase?.("arrow_node_batch_decoded", {
          batch_index: batchIndex,
          rows: batch.numRows,
          nodes: vertexIds.length,
          decode_ms: Math.round(performance.now() - batchStartedAt),
        });
        break;
      }
      case 2: {
        ensureAllocated(meta, vertexIds, xs, ys, edgeData, routeLineVertices, labels, iris);
        const edgeSources = readUint32List(batch, "edge_source");
        const edgeTargets = readUint32List(batch, "edge_target");
        if (edgeSources.length !== edgeTargets.length) {
          throw new Error("Arrow edge batch source/target length mismatch");
        }
        if (edgeData.length !== edgeSources.length * 2) {
          throw new Error("Arrow edge batch size mismatch");
        }
        for (let i = 0; i < edgeSources.length; i++) {
          edgeData[i * 2] = edgeSources[i];
          edgeData[i * 2 + 1] = edgeTargets[i];
        }
        logPhase?.("arrow_edge_batch_decoded", {
          batch_index: batchIndex,
          rows: batch.numRows,
          edges: edgeSources.length,
          decode_ms: Math.round(performance.now() - batchStartedAt),
        });
        break;
      }
      case 3: {
        ensureAllocated(meta, vertexIds, xs, ys, edgeData, routeLineVertices, labels, iris);
        const routeX1 = readFloat32List(batch, "route_x1");
        const routeY1 = readFloat32List(batch, "route_y1");
        const routeX2 = readFloat32List(batch, "route_x2");
        const routeY2 = readFloat32List(batch, "route_y2");
        if (
          routeX1.length !== routeY1.length ||
          routeX1.length !== routeX2.length ||
          routeX1.length !== routeY2.length
        ) {
          throw new Error("Arrow route batch axis length mismatch");
        }
        if (routeLineVertices.length !== routeX1.length * 4) {
          throw new Error("Arrow route batch size mismatch");
        }
        for (let i = 0; i < routeX1.length; i++) {
          routeLineVertices[i * 4] = routeX1[i];
          routeLineVertices[i * 4 + 1] = routeY1[i];
          routeLineVertices[i * 4 + 2] = routeX2[i];
          routeLineVertices[i * 4 + 3] = routeY2[i];
        }
        logPhase?.("arrow_route_batch_decoded", {
          batch_index: batchIndex,
          rows: batch.numRows,
          route_line_segments: routeX1.length,
          decode_ms: Math.round(performance.now() - batchStartedAt),
        });
        break;
      }
      default:
        throw new Error(`Unexpected Arrow batch index ${batchIndex}`);
    }

    batchIndex++;
  }

  if (batchIndex !== 4) {
    throw new Error(`Expected 4 Arrow batches, received ${batchIndex}`);
  }
  ensureAllocated(meta, vertexIds, xs, ys, edgeData, routeLineVertices, labels, iris);

  return {
    meta,
    vertexIds,
    xs,
    ys,
    edgeData,
    routeLineVertices,
    labels,
    iris,
  };
}

async function openArrowReader(response: Response) {
  if (response.body) {
    return RecordBatchReader.from(response.body);
  }
  return RecordBatchReader.from(await response.arrayBuffer());
}

function validateArrowResponse(response: Response): void {
  const contentType = response.headers.get("content-type") || "";
  if (!contentType.includes("application/vnd.apache.arrow.stream")) {
    throw new Error(`Unexpected graph content type: ${contentType || "(missing)"}`);
  }
  const version = response.headers.get("X-Graph-Transport-Version");
  if (version !== graphTransportVersion) {
    throw new Error(`Unexpected graph transport version: ${version || "(missing)"}`);
  }
}

function decodeMetaBatch(batch: any): GraphMeta {
  return {
    backend: readScalarString(batch, "meta_backend"),
    graph_query_id: readScalarString(batch, "meta_graph_query_id"),
    node_limit: readScalarNumber(batch, "meta_node_limit"),
    edge_limit: readScalarNumber(batch, "meta_edge_limit"),
    nodes: readScalarNumber(batch, "meta_nodes"),
    edges: readScalarNumber(batch, "meta_edges"),
  };
}

function ensureAllocated(
  meta: GraphMeta | null,
  vertexIds: Uint32Array | null,
  xs: Float32Array | null,
  ys: Float32Array | null,
  edgeData: Uint32Array | null,
  routeLineVertices: Float32Array | null,
  labels: (string | undefined)[] | null,
  iris: string[] | null
): asserts meta is GraphMeta & {} & {
  nodes?: number;
  edges?: number;
} {
  if (!meta || !vertexIds || !xs || !ys || !edgeData || !routeLineVertices || !labels || !iris) {
    throw new Error("Arrow graph stream is missing the meta batch");
  }
}

function readScalarString(batch: any, name: string): string | undefined {
  const vector = batch.getChild(name);
  if (!vector || batch.numRows < 1) return undefined;
  const value = vector.get(0);
  return typeof value === "string" ? value : undefined;
}

function readScalarNumber(batch: any, name: string): number | undefined {
  const vector = batch.getChild(name);
  if (!vector || batch.numRows < 1) return undefined;
  const value = vector.get(0);
  return typeof value === "number" && Number.isFinite(value) ? value : undefined;
}

function readUint32List(batch: any, name: string): Uint32Array {
  return readPrimitiveList(batch, name, Uint32Array);
}

function readFloat32List(batch: any, name: string): Float32Array {
  return readPrimitiveList(batch, name, Float32Array);
}

function readPrimitiveList<T extends Uint32Array | Float32Array>(
  batch: any,
  name: string,
  ctor: { new(length: number): T }
): T {
  const vector = batch.getChild(name) as ArrowLikeVector | null;
  if (!vector) return new ctor(0);

  const [begin, end] = getListBounds(vector);
  const child = vector.getChildAt?.(0);
  if (!child || typeof child.toArray !== "function") {
    return new ctor(0);
  }
  const values = child.toArray();
  if (typeof (values as any).subarray === "function") {
    return (values as T).subarray(begin, end) as T;
  }

  const out = new ctor(end - begin);
  for (let i = begin; i < end; i++) {
    out[i - begin] = Number((values as ArrayLike<number>)[i]) as T[number];
  }
  return out;
}

function readStringList(batch: any, name: string): string[] {
  return readNullableStringList(batch, name).map((value) => value ?? "");
}

function readNullableStringList(batch: any, name: string): (string | undefined)[] {
  const vector = batch.getChild(name) as ArrowLikeVector | null;
  if (!vector) return [];

  const [begin, end] = getListBounds(vector);
  const child = vector.getChildAt?.(0);
  if (!child || typeof child.get !== "function") {
    return [];
  }

  const out = new Array<string | undefined>(Math.max(0, end - begin));
  for (let i = begin; i < end; i++) {
    const value = child.get(i);
    out[i - begin] = typeof value === "string" ? value : undefined;
  }
  return out;
}

function getListBounds(vector: ArrowLikeVector): [number, number] {
  const offsets = vector.data?.[0]?.valueOffsets;
  if (!offsets || offsets.length < 2) return [0, 0];
  return [offsetToNumber(offsets[0]), offsetToNumber(offsets[1])];
}

function offsetToNumber(value: number | bigint | undefined): number {
  if (typeof value === "bigint") return Number(value);
  return typeof value === "number" ? value : 0;
}
