import { existsSync, readFileSync, realpathSync, statSync } from "node:fs";
import {
  dirname,
  isAbsolute,
  normalize,
  relative,
  resolve,
  sep,
} from "node:path";

export interface ArtifactProvider {
  isFile(path: string, options?: ArtifactReadOptions): boolean;
  readText(path: string, options?: ArtifactReadOptions): string;
}

export type ArtifactReadOptions = {
  expandHome?: boolean;
};

export type ArtifactBundle = {
  version: 1;
  files: Record<string, string>;
};

export type ArtifactLimits = {
  maxArtifactCount?: number;
  maxArtifactBytes?: number;
  maxTotalBytes?: number;
};

export type LocalArtifactProviderOptions = ArtifactLimits & {
  workspace: string;
  strict?: boolean;
};

export type BundleArtifactProviderOptions = ArtifactLimits;

export const DEFAULT_ARTIFACT_LIMITS = {
  maxArtifactCount: 100,
  maxArtifactBytes: 1024 * 1024,
  maxTotalBytes: 4 * 1024 * 1024,
} as const;

function artifactError(
  message: string,
  code?: string,
): Error & { code?: string } {
  return Object.assign(new Error(message), code ? { code } : {});
}

function toWorkspacePath(path: string): string {
  if (
    !path ||
    path.includes("\0") ||
    path.includes("\\") ||
    path.startsWith("~") ||
    isAbsolute(path)
  ) {
    throw artifactError(`Artifact path must be workspace-relative: ${path}`);
  }

  const normalized = normalize(path);
  if (normalized === ".." || normalized.startsWith(`..${sep}`)) {
    throw artifactError(`Artifact path escapes the workspace: ${path}`);
  }

  return normalized.replace(/^\.\//, "");
}

function assertWithinWorkspace(workspace: string, path: string): void {
  const relation = relative(workspace, path);
  if (
    relation === ".." ||
    relation.startsWith(`..${sep}`) ||
    isAbsolute(relation)
  ) {
    throw artifactError(`Artifact path escapes the workspace: ${path}`);
  }
}

function assertRealPathWithinWorkspace(workspace: string, path: string): void {
  let existingPath = path;
  while (!existsSync(existingPath)) {
    const parent = dirname(existingPath);
    if (parent === existingPath) break;
    existingPath = parent;
  }
  assertWithinWorkspace(workspace, realpathSync(existingPath));
}

function assertPositiveLimit(name: string, value: number | undefined): void {
  if (value !== undefined && (!Number.isSafeInteger(value) || value < 0)) {
    throw new Error(`${name} must be a non-negative safe integer`);
  }
}

class ArtifactBudget {
  readonly #limits: ArtifactLimits;
  #count = 0;
  #bytes = 0;

  constructor(limits: ArtifactLimits) {
    assertPositiveLimit("maxArtifactCount", limits.maxArtifactCount);
    assertPositiveLimit("maxArtifactBytes", limits.maxArtifactBytes);
    assertPositiveLimit("maxTotalBytes", limits.maxTotalBytes);
    this.#limits = limits;
  }

  add(path: string, content: string): void {
    const bytes = Buffer.byteLength(content);
    if (
      this.#limits.maxArtifactBytes !== undefined &&
      bytes > this.#limits.maxArtifactBytes
    ) {
      throw artifactError(
        `Artifact ${path} exceeds the ${this.#limits.maxArtifactBytes} byte limit`,
      );
    }
    if (
      this.#limits.maxArtifactCount !== undefined &&
      this.#count + 1 > this.#limits.maxArtifactCount
    ) {
      throw artifactError(
        `Artifact count exceeds the ${this.#limits.maxArtifactCount} limit`,
      );
    }
    if (
      this.#limits.maxTotalBytes !== undefined &&
      this.#bytes + bytes > this.#limits.maxTotalBytes
    ) {
      throw artifactError(
        `Artifact bytes exceed the ${this.#limits.maxTotalBytes} byte limit`,
      );
    }

    this.#count += 1;
    this.#bytes += bytes;
  }
}

export class LocalArtifactProvider implements ArtifactProvider {
  readonly #workspace: string;
  readonly #strict: boolean;
  readonly #budget: ArtifactBudget;

  constructor(options: LocalArtifactProviderOptions) {
    this.#workspace = options.strict
      ? realpathSync(options.workspace)
      : resolve(options.workspace);
    this.#strict = options.strict ?? false;
    this.#budget = new ArtifactBudget(options);
  }

  #resolve(path: string, options: ArtifactReadOptions): string {
    if (!this.#strict) {
      return options.expandHome && path.startsWith("~")
        ? resolve(process.env.HOME ?? "", path.slice(1))
        : resolve(this.#workspace, path);
    }

    const candidate = resolve(this.#workspace, toWorkspacePath(path));
    assertWithinWorkspace(this.#workspace, candidate);
    assertRealPathWithinWorkspace(this.#workspace, candidate);
    return candidate;
  }

  isFile(path: string, options: ArtifactReadOptions = {}): boolean {
    const resolved = this.#resolve(path, options);
    return existsSync(resolved) && statSync(resolved).isFile();
  }

  readText(path: string, options: ArtifactReadOptions = {}): string {
    const resolved = this.#resolve(path, options);
    if (this.#strict) {
      // Resolve again at read time so a symlink swap cannot bypass confinement.
      assertWithinWorkspace(this.#workspace, realpathSync(resolved));
    }
    const content = readFileSync(resolved, "utf8");
    this.#budget.add(path, content);
    return content;
  }
}

export class BundleArtifactProvider implements ArtifactProvider {
  readonly #files = new Map<string, string>();

  constructor(
    bundle: ArtifactBundle,
    options: BundleArtifactProviderOptions = {},
  ) {
    if (bundle.version !== 1)
      throw new Error("Unsupported artifact bundle version");

    const budget = new ArtifactBudget({
      ...DEFAULT_ARTIFACT_LIMITS,
      ...options,
    });
    for (const [path, content] of Object.entries(bundle.files)) {
      const normalized = toWorkspacePath(path);
      if (this.#files.has(normalized)) {
        throw artifactError(`Duplicate artifact path: ${path}`);
      }
      budget.add(normalized, content);
      this.#files.set(normalized, content);
    }
  }

  isFile(path: string): boolean {
    return this.#files.has(toWorkspacePath(path));
  }

  readText(path: string): string {
    const normalized = toWorkspacePath(path);
    const content = this.#files.get(normalized);
    if (content === undefined) {
      throw artifactError(`Artifact not found: ${path}`, "ENOENT");
    }
    return content;
  }
}

export function createArtifactBundle(
  files: Record<string, string>,
): ArtifactBundle {
  return { version: 1, files: { ...files } };
}