import { useEffect, useState } from "react";
import type { RemoteH5FileX, RemoteH5Group } from "../remote-h5-file";

// Number of samples we lazily pull from the multi-million-sample timeseries.
export const WINDOW_SAMPLES = 30000;

export type NwbData = {
  rootGroup?: RemoteH5Group;
  meta?: Record<string, string>;
  timeseries?: {
    t: Float64Array;
    pupil: Float64Array;
    running: Float64Array;
    totalSamples: number;
  };
  trials?: { startTime: Float64Array; stimFrequency: Float64Array };
  step: string;
  error?: string;
  done: boolean;
};

const PUPIL = "/processing/behavior/PupilTracking/pupil_diameter";
const RUNNING = "/processing/behavior/running_speed";
const TRIALS = "/intervals/trials";

const toF64 = (x: unknown): Float64Array => {
  if (x instanceof Float64Array) return x;
  if (ArrayBuffer.isView(x)) return Float64Array.from(x as unknown as number[]);
  if (Array.isArray(x)) return Float64Array.from(x as number[]);
  return new Float64Array(0);
};

const readScalar = async (
  file: RemoteH5FileX,
  path: string,
): Promise<string> => {
  try {
    const v = await file.getDatasetData(path, {});
    if (v === undefined || v === null) return "";
    if (typeof v === "string") return v;
    if (ArrayBuffer.isView(v) || Array.isArray(v)) {
      return Array.from(v as unknown as number[]).join(", ");
    }
    return String(v);
  } catch {
    return "";
  }
};

export const useNwbData = (file: RemoteH5FileX | undefined): NwbData => {
  const [data, setData] = useState<NwbData>({ step: "idle", done: false });

  useEffect(() => {
    if (!file) {
      setData({ step: "idle", done: false });
      return;
    }
    let canceled = false;
    const update = (patch: Partial<NwbData>) =>
      !canceled && setData((d) => ({ ...d, ...patch }));

    const run = async () => {
      try {
        update({ step: "reading file structure (root group)", done: false });
        const rootGroup = await file.getGroup("/");
        update({ rootGroup });

        update({ step: "reading metadata fields" });
        const subjectPrefix = "/general/subject";
        const meta: Record<string, string> = {
          session_description: await readScalar(file, "/session_description"),
          identifier: await readScalar(file, "/identifier"),
          session_start_time: await readScalar(file, "/session_start_time"),
          institution: await readScalar(file, "/general/institution"),
          lab: await readScalar(file, "/general/lab"),
          subject_id: await readScalar(file, `${subjectPrefix}/subject_id`),
          species: await readScalar(file, `${subjectPrefix}/species`),
          sex: await readScalar(file, `${subjectPrefix}/sex`),
          age: await readScalar(file, `${subjectPrefix}/age`),
        };
        update({ meta });

        update({
          step: `lazily reading first ${WINDOW_SAMPLES.toLocaleString()} samples of pupil & running traces`,
        });
        const pupilDs = await file.getDataset(`${PUPIL}/data`);
        const totalSamples = pupilDs?.shape?.[0] ?? 0;
        const n = Math.min(WINDOW_SAMPLES, totalSamples || WINDOW_SAMPLES);
        const slice: [number, number][] = [[0, n]];
        const t = toF64(
          await file.getDatasetData(`${PUPIL}/timestamps`, { slice }),
        );
        const pupil = toF64(
          await file.getDatasetData(`${PUPIL}/data`, { slice }),
        );
        const running = toF64(
          await file.getDatasetData(`${RUNNING}/data`, { slice }),
        );
        update({ timeseries: { t, pupil, running, totalSamples } });

        update({ step: "reading trials table (stimulus frequency)" });
        const startTime = toF64(
          await file.getDatasetData(`${TRIALS}/start_time`, {}),
        );
        const stimFrequency = toF64(
          await file.getDatasetData(`${TRIALS}/stim_frequency`, {}),
        );
        update({ trials: { startTime, stimFrequency } });

        update({ step: "done", done: true });
      } catch (err: unknown) {
        update({
          step: "error",
          error: err instanceof Error ? err.message : String(err),
          done: true,
        });
      }
    };
    run();
    return () => {
      canceled = true;
    };
  }, [file]);

  return data;
};