import { createHash } from "node:crypto";
import path from "node:path";
import { canonicalJson, type JsonObject } from "../core/json.js";
import type { JazzThoughtStore } from "../jazz/store.js";
import { assertPrivateDestination, openPrivateDirectory } from "../security/private-files.js";
import {
  projectPrivateTrainingExamples,
  type JudgmentKind,
  type TrainingExample,
  type TrainingProjectionOptions,
} from "./judgments.js";

export interface PrivateTrainingExportOptions extends Omit<TrainingProjectionOptions, "includeSensitivePrivate"> {
  acknowledgeSensitivePrivate: true;
  publicContentRoots?: string[] | undefined;
  beforeFinalize?: (() => void | Promise<void>) | undefined;
}

export interface PrivateTrainingDatasetManifest {
  format: "thoughtstream.private-training-dataset-manifest.v1";
  datasetId: string;
  generatedAt: string;
  dataFile: string;
  sha256: string;
  examples: number;
  kinds: Partial<Record<JudgmentKind, number>>;
  models: string[];
  exactPrivateProvenance: true;
  externalExportAuthorityRequired: false;
}

export async function exportPrivateTrainingDataset(
  store: JazzThoughtStore,
  destination: string,
  options: PrivateTrainingExportOptions,
): Promise<PrivateTrainingDatasetManifest> {
  if (!destination || !path.isAbsolute(path.resolve(destination))) throw new Error("Private training export requires an explicit file destination");
  if (options.acknowledgeSensitivePrivate !== true) throw new Error("Private training export requires explicit sensitive-data acknowledgment");
  const examples = await projectPrivateTrainingExamples(store, {
    ...(options.includeRestrictedModelAdapters ? { includeRestrictedModelAdapters: true } : {}),
  });
  return writePrivateTrainingJsonl(destination, examples, options);
}

export async function writePrivateTrainingJsonl(
  destination: string,
  examples: TrainingExample[],
  options: PrivateTrainingExportOptions,
): Promise<PrivateTrainingDatasetManifest> {
  if (options.acknowledgeSensitivePrivate !== true) throw new Error("Private training export requires explicit sensitive-data acknowledgment");
  if (examples.some((example) => example.format !== "thoughtstream.private-training-example.v1" || !hasExactPrivateProvenance(example))) {
    throw new Error("Private training export accepts only complete exact-provenance private training examples");
  }
  const absolute = await assertPrivateDestination(path.resolve(destination), { publicContentRoots: options.publicContentRoots });
  const manifestPath = await assertPrivateDestination(`${absolute}.manifest.json`, { publicContentRoots: options.publicContentRoots });
  const serialized = examples.length > 0
    ? `${examples.map((example) => canonicalJson(example as unknown as JsonObject)).join("\n")}\n`
    : "";
  const digest = createHash("sha256").update(serialized).digest("hex");
  const manifest: PrivateTrainingDatasetManifest = {
    format: "thoughtstream.private-training-dataset-manifest.v1",
    datasetId: `sha256:${digest}`,
    generatedAt: new Date().toISOString(),
    dataFile: path.basename(absolute),
    sha256: digest,
    examples: examples.length,
    kinds: examples.reduce<Partial<Record<JudgmentKind, number>>>((counts, example) => {
      counts[example.kind] = (counts[example.kind] ?? 0) + 1;
      return counts;
    }, {}),
    models: [...new Set(examples.map((example) => `${example.provenance.provider}:${example.provenance.model}`))].sort(),
    exactPrivateProvenance: true,
    externalExportAuthorityRequired: false,
  };
  const directory = await openPrivateDirectory(path.dirname(absolute), {
    publicContentRoots: options.publicContentRoots,
    beforeFinalize: options.beforeFinalize,
  });
  try {
    await directory.write(path.basename(absolute), serialized);
    await directory.write(path.basename(manifestPath), `${canonicalJson(manifest as unknown as JsonObject)}\n`);
  } finally {
    await directory.close();
  }
  return manifest;
}

function hasExactPrivateProvenance(example: TrainingExample): boolean {
  const provenance = example.privateProvenance;
  return Boolean(provenance && [
    provenance.judgmentEventId,
    provenance.runId,
    provenance.outputEventId,
    provenance.triggerEventId,
    provenance.feedbackSourceEventId,
    provenance.deliveryReceiptEventId,
    provenance.contextSnapshotId,
  ].every((value) => typeof value === "string" && value.length > 0));
}