import { createHash } from "node:crypto";
import { createReadStream, createWriteStream } from "node:fs";
import { access, mkdir, readdir, readFile, rename, rm, stat, writeFile } from "node:fs/promises";
import { dirname, join, normalize, resolve, sep } from "node:path";
import type { Readable } from "node:stream";
import { pipeline } from "node:stream/promises";

export interface StoredObject {
  key: string;
  sha256: string;
  bytes: number;
}

/** Write-once, checksum-verified object storage. Filesystem backend by default; S3 can implement the same interface. */
export interface StorageBackend {
  put(key: string, data: Buffer | Readable, expectedSha256?: string): Promise<StoredObject>;
  get(key: string): Promise<Buffer>;
  stream(key: string): Promise<Readable>;
  head(key: string): Promise<StoredObject | null>;
  exists(key: string): Promise<boolean>;
  copy(fromKey: string, toKey: string): Promise<StoredObject>;
  list(prefix: string): Promise<string[]>;
  remove(key: string): Promise<void>;
  absolutePath(key: string): string;
}

export function sha256Hex(data: Buffer | string): string {
  return createHash("sha256").update(data).digest("hex");
}

export async function sha256File(path: string): Promise<string> {
  const h = createHash("sha256");
  for await (const chunk of createReadStream(path)) h.update(chunk as Buffer);
  return h.digest("hex");
}

export function assertSafeKey(key: string): string {
  const n = normalize(key).replace(/\\/g, "/");
  if (!n || n.startsWith("/") || n.startsWith("..") || n.split("/").includes("..") || n.includes("\0")) {
    throw new Error(`unsafe storage key: ${key}`);
  }
  return n;
}

export class FilesystemStorage implements StorageBackend {
  readonly root: string;
  constructor(root: string) {
    this.root = resolve(root);
  }

  absolutePath(key: string): string {
    const safe = assertSafeKey(key);
    const abs = resolve(this.root, safe);
    if (!abs.startsWith(this.root + sep) && abs !== this.root) throw new Error(`key escapes root: ${key}`);
    return abs;
  }

  /** Digest sidecar. The suffix cannot collide with user artifact names (research code often writes `*.sha256`). */
  private sidecar(abs: string): string {
    return `${abs}.~digest`;
  }

  async put(key: string, data: Buffer | Readable, expectedSha256?: string): Promise<StoredObject> {
    const abs = this.absolutePath(key);
    if (await this.exists(key)) {
      const existing = await this.head(key);
      if (existing && expectedSha256 && existing.sha256 !== expectedSha256) {
        throw new Error(`object exists with different digest: ${key}`);
      }
      if (existing) return existing; // write-once: identical content is idempotent
    }
    await mkdir(dirname(abs), { recursive: true });
    const tmp = `${abs}.tmp-${process.pid}-${Date.now()}`;
    const h = createHash("sha256");
    let bytes = 0;
    if (Buffer.isBuffer(data)) {
      h.update(data);
      bytes = data.length;
      await writeFile(tmp, data, { flag: "wx" });
    } else {
      await pipeline(
        data,
        async function* (src) {
          for await (const chunk of src) {
            const b = chunk as Buffer;
            h.update(b);
            bytes += b.length;
            yield b;
          }
        },
        createWriteStream(tmp, { flags: "wx" }),
      );
    }
    const sha256 = h.digest("hex");
    if (expectedSha256 && expectedSha256 !== sha256) {
      await rm(tmp, { force: true });
      throw new Error(`digest mismatch for ${key}: expected ${expectedSha256}, got ${sha256}`);
    }
    await rename(tmp, abs);
    await writeFile(this.sidecar(abs), `${sha256}  ${key}\n`);
    return { key: assertSafeKey(key), sha256, bytes };
  }

  async get(key: string): Promise<Buffer> {
    const abs = this.absolutePath(key);
    const buf = await readFile(abs);
    const meta = await this.head(key);
    if (meta && sha256Hex(buf) !== meta.sha256) throw new Error(`integrity failure reading ${key}`);
    return buf;
  }

  async stream(key: string): Promise<Readable> {
    return createReadStream(this.absolutePath(key));
  }

  async head(key: string): Promise<StoredObject | null> {
    const abs = this.absolutePath(key);
    try {
      const s = await stat(abs);
      let sha256: string;
      try {
        const side = await readFile(this.sidecar(abs), "utf8");
        sha256 = side.slice(0, 64);
      } catch {
        sha256 = await sha256File(abs);
      }
      return { key: assertSafeKey(key), sha256, bytes: s.size };
    } catch {
      return null;
    }
  }

  async exists(key: string): Promise<boolean> {
    try {
      await access(this.absolutePath(key));
      return true;
    } catch {
      return false;
    }
  }

  async copy(fromKey: string, toKey: string): Promise<StoredObject> {
    const buf = await this.get(fromKey);
    return this.put(toKey, buf);
  }

  async list(prefix: string): Promise<string[]> {
    const abs = this.absolutePath(prefix);
    const out: string[] = [];
    const walk = async (dir: string, rel: string) => {
      let entries: import("node:fs").Dirent[];
      try {
        entries = (await readdir(dir, { withFileTypes: true })) as import("node:fs").Dirent[];
      } catch {
        return;
      }
      for (const e of entries) {
        const name = String(e.name);
        if (name.endsWith(".~digest") || name.includes(".tmp-")) continue;
        const r = rel ? `${rel}/${name}` : name;
        if (e.isDirectory()) await walk(join(dir, name), r);
        else out.push(`${assertSafeKey(prefix)}/${r}`.replace(/^\.\//, ""));
      }
    };
    await walk(abs, "");
    return out.sort();
  }

  async remove(key: string): Promise<void> {
    const abs = this.absolutePath(key);
    await rm(abs, { force: true });
    await rm(this.sidecar(abs), { force: true });
  }
}

export function runArtifactKey(runId: string, attemptNo: number, path: string): string {
  return assertSafeKey(`runs/${runId}/a${attemptNo}/${path}`);
}

export function publicArtifactKey(publicId: string, versionNo: number, path: string): string {
  return assertSafeKey(`public/${publicId}/v${versionNo}/${path}`);
}

export function getStorage(root = process.env.ARTIFACT_ROOT ?? "./data/artifacts"): FilesystemStorage {
  return new FilesystemStorage(root);
}