src/lib/bucket.ts

import path from "node:path";
import { AccessToken, RepoDesignation } from "@huggingface/hub";
import * as hub from "@huggingface/hub";
import { NS_routes_c } from "../../../../../packages/lib/helpers/client/route";
import { NS_misc } from "./misc";
import { NS_promise } from "./promise";
import { NS_ffmpeg_v2 } from "./ffmpeg-v2";
import { NS_runtime } from "../../../../../packages/lib/env.runtime";

const HF_BUCKET_ACCESS_TOKEN = () =>
  NS_runtime.ENV_VAR_DB?.HF_BUCKET_ACCESS_TOKEN ||
  process.env.HF_BUCKET_ACCESS_TOKEN;
const HF_BUCKET_REPO = () =>
  NS_runtime.ENV_VAR_DB?.HF_BUCKET_REPO || process.env.HF_BUCKET_REPO;

export const STORAGE_ROOT = () =>
  NS_runtime.ENV_VAR_DB?.STORAGE_ROOT ||
  process.env.STORAGE_ROOT ||
  path.join(process.cwd(), "S3");

export const hf_repo = () => {
  const data: RepoDesignation = {
    type: "bucket",
    name: HF_BUCKET_REPO() || ""
  };
  return data;
};

export const bucket_credentials = () => {
  const data: { accessToken: AccessToken } = {
    accessToken: HF_BUCKET_ACCESS_TOKEN()!
  };

  return data;
};

// await hub.createRepo({
//   repo: hf_repo(),
//   accessToken: bucket_credentials().accessToken,
//   private: false
// });

type T_UploadJob = {
  status: "in-progress" | "done" | "failed";
  tmp_path?: string | null;
  created_at: number;
  error?: string;
  percent?: number;
  transcoded_size?: string;
  progress_endpoint?: string;
};

const upload_jobs = new Map<string, T_UploadJob>();

export const NS_bucket = {
  track_upload_process: {
    v1: trackProgress,
    v2: trackProgress_v2
  },
  get_hf_bucket_file_info,
  get_hf_bucket_list_of_files,
  get_latest_hf_bucket_file_by_name,
  upload_to_bucket
};

function trackProgress(
  blob: Blob,
  onProgress: (pct: number) => void
): ReadableStream<Uint8Array> {
  const total = blob.size;
  let loaded = 0;
  let lastReported = -1;

  return new ReadableStream({
    async start(controller) {
      const reader = blob.stream().getReader();
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        loaded += value.byteLength;

        // only fire when % actually changes
        const pct = Math.round((loaded / total) * 100);
        if (pct !== lastReported) {
          lastReported = pct;
          onProgress(pct);
        }

        controller.enqueue(value);
      }
      controller.close();
    }
  });
}

function trackProgress_v2(
  blob: Blob,
  onProgress: (pct: number) => void
): ReadableStream {
  const total = blob.size;
  let loaded = 0;
  let lastReported = -1;

  return new ReadableStream({
    async start(controller) {
      const reader = blob.stream().getReader();
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
        loaded += value.byteLength;

        const pct = Math.round((loaded / total) * 100);
        if (pct % 5 === 0 && pct !== lastReported) {
          lastReported = pct;
          onProgress(pct);
        }

        controller.enqueue(value);
      }
      controller.close();
    }
  });
}

async function get_hf_bucket_file_info(arg0: { file_path: string }) {
  const { file_path } = arg0;

  const credentials = bucket_credentials();
  const repo = hf_repo();

  if (
    !credentials.accessToken ||
    typeof file_path !== "string" ||
    NS_routes_c.F_is_url(file_path)
  ) {
    return;
  }

  try {
    const info = await hub.fileDownloadInfo({
      repo,
      path: file_path,
      accessToken: credentials.accessToken
    });

    return info;
  } catch {
    // file doesn't exist
  }
}

async function get_hf_bucket_list_of_files() {
  const credentials = bucket_credentials();
  const repo = hf_repo();

  if (!credentials.accessToken) return;

  const available_files = await Array.fromAsync(
    hub.listFiles({
      repo,
      recursive: true,
      accessToken: credentials.accessToken
    })
  );

  return available_files;
}

async function get_latest_hf_bucket_file_by_name(arg0: {
  folder_path: string;
  human_filename: string;
  sbe_id?: string;
}) {
  const { folder_path, human_filename, sbe_id } = arg0;
  const credentials = bucket_credentials();
  const repo = hf_repo();
  if (
    !credentials.accessToken ||
    typeof folder_path !== "string" ||
    typeof human_filename !== "string"
  ) {
    return;
  }

  try {
    const files = await Array.fromAsync(
      hub.listFiles({
        repo,
        accessToken: credentials.accessToken,
        path: folder_path,
        recursive: false
      })
    );

    // match files whose name ends with -{human_filename}
    const matching = files.filter((f) => {
      const name = f.path.split("/").pop() ?? "";
      const withoutPrefix = name.slice(name.indexOf("-") + 1);
      return withoutPrefix.toLowerCase() === human_filename.toLowerCase();
    });

    if (matching.length === 0) return;

    // pool 1 — sbe_id prefixed
    const sbe_id_pool = sbe_id
      ? matching.filter((f) => {
          const name = f.path.split("/").pop() ?? "";
          const prefix = name.slice(0, name.indexOf("-"));
          return prefix === sbe_id;
        })
      : [];

    // pool 2 — timestamp prefixed, from what's left after removing sbe_id matches
    const remaining = matching.filter((f) => !sbe_id_pool.includes(f));

    const timestamp_pool = remaining
      .map((f) => {
        const name = f.path.split("/").pop() ?? "";
        const prefix = Number(name.slice(0, name.indexOf("-")));
        return { f, ts: NS_misc.isValidUnixTimestamp(prefix) ? prefix : null };
      })
      .filter((x) => x.ts !== null)
      .sort((a, b) => b.ts! - a.ts!)
      .map((x) => x.f);

    // prefer timestamp pool, fallback to sbe_id pool
    let pool =
      timestamp_pool.length > 0
        ? timestamp_pool
        : sbe_id_pool.length > 0
          ? sbe_id_pool
          : [];

    if (true) {
      const _ = matching.sort((a, b) => {
        const da = new Date(a.uploadedAt ?? 0).getTime();
        const db = new Date(b.uploadedAt ?? 0).getTime();
        return db - da;
      });

      if (_.length) {
        pool = [..._];
      }
    }

    const best_file = await get_hf_bucket_file_info({
      file_path: pool[0].path
    });

    return best_file;
  } catch {}
}

async function upload_to_bucket(arg0: {
  job_id: string;
  bucket_path: string;
  file_with_path: string;
  skip?: {
    probe?: boolean;
  };
  REQ0?: URL;
}) {
  const { job_id, bucket_path, skip, file_with_path, REQ0 } = arg0;

  const credentials = bucket_credentials();
  const repo = hf_repo();

  upload_jobs.set(job_id, {
    status: "in-progress",
    created_at: Date.now()
  });

  console.log("\n", `🧲 UPLOADING TO BUCKET: ${job_id}`, "\n");

  if (!skip?.probe) {
    const [E_, R_] = await NS_promise.F_catch_promise({
      promise: NS_ffmpeg_v2.F_probe_t1({
        playback: {
          url: file_with_path,
          type: "video",
          protocol: "local"
        }
      })
    });

    if (E_ || !R_) {
      upload_jobs.delete(job_id);
      console.log("\n", `❌ UPLOAD BUCKET: FFPROBE FAILED: ${job_id}`, "\n");
      return;
    }
  }

  const blob: any = Bun.file(file_with_path);

  const [E__, R__] = await NS_promise.F_catch_promise({
    promise: hub.uploadFiles({
      repo,
      accessToken: credentials.accessToken,
      files: [
        {
          path: bucket_path,
          content: blob
        }
      ]
    })
  });

  upload_jobs.delete(job_id);

  if (E__) {
    console.log("\n", `❌ UPLOAD BUCKET FAILED: ${job_id}`, E__, "\n");
    return { success: false };
  }

  console.log("\n", `🔥 UPLOAD BUCKET DONE: ${job_id}`, "\n");

  return { success: true };
}