import { useEffect, useState } from "react";
import { tableFromIPC } from "apache-arrow";
import { API, USE_ARROW } from "@/lib/constants";
import type { Block, Meta } from "@/types/trajectory";
import { framesFromArrowTable } from "@/lib/trajectoryData";

export function useBlocks(
  meta: Meta | null,
  simId: string = "default"
): {
  blocks: (Block | null)[];
  loadProgress: number;
  ready: boolean;
} {
  const [blocks, setBlocks] = useState<(Block | null)[]>([null, null, null]);
  const [loadProgress, setLoadProgress] = useState(0);
  const [ready, setReady] = useState(false);

  useEffect(() => {
    if (!meta) return;
    const nBlocks = meta.num_blocks;
    const nParticles = meta.num_particles;
    const simParam = `sim=${encodeURIComponent(simId)}`;
    let loaded = 0;
    const result: (Block | null)[] = new Array(nBlocks).fill(null);
    setBlocks(result);
    setReady(false);

    function onBlockLoaded(i: number, block: Block | null) {
      result[i] = block;
      loaded++;
      setLoadProgress(Math.round((loaded / nBlocks) * 100));
      if (loaded === nBlocks) {
        setBlocks([...result]);
        setReady(true);
      } else {
        setBlocks((prev) => {
          const next = Array.from({ length: nBlocks }, (_, j) => (j === i ? block : prev[j] ?? result[j] ?? null));
          return next;
        });
      }
    }

    async function loadBlock(i: number) {
      try {
        if (USE_ARROW) {
          const [metaRes, arrowRes] = await Promise.all([
            fetch(`${API}/api/block/${i}?frames=false&${simParam}`),
            fetch(`${API}/api/block/${i}/arrow?${simParam}`),
          ]);
          if (metaRes.ok && arrowRes.ok) {
            const metaBlock = (await metaRes.json()) as Pick<
              Block,
              "block" | "step_start" | "step_end" | "accumulation"
            >;
            const buf = await arrowRes.arrayBuffer();
            const table = tableFromIPC(new Uint8Array(buf));
            const frames = framesFromArrowTable(table, nParticles);
            onBlockLoaded(i, { ...metaBlock, frames });
            return;
          }
        }
      } catch {
        /* fallback to JSON */
      }
      try {
        const full = (await fetch(`${API}/api/block/${i}?${simParam}`).then((r) =>
          r.json()
        )) as Block;
        onBlockLoaded(i, full);
      } catch {
        onBlockLoaded(i, null);
      }
    }

    for (let i = 0; i < nBlocks; i++) loadBlock(i);
  }, [meta, simId]);

  return { blocks, loadProgress, ready };
}
