import {
  canonicalJson, fnv1a64,
  isDiagnosticToken,
  isWellFormedDiagnosticString,
  snapshotDiagnosticJson,
  type CustomizationResourceLimits,
  type JsonValue,
} from "@actuarial-ts/core";
import { parseCustomizationResourceLimits } from "./customizationContracts.js";

export type CustomizationExecutionErrorCode =
  | "CANCELLED"
  | "RESOURCE_LIMIT"
  | "EXTERNAL_STATE_FAILURE"
  | "EXTERNAL_STATE_ORDER"
  | "EXTERNAL_STATE_TRUNCATED";

export class CustomizationExecutionError extends Error {
  readonly code: CustomizationExecutionErrorCode;

  constructor(code: CustomizationExecutionErrorCode, message: string) {
    super(message);
    this.name = "CustomizationExecutionError";
    this.code = code;
    Object.freeze(this);
  }
}

export interface CustomizationExternalRecord<T extends JsonValue = JsonValue> {
  readonly id: string;
  readonly partitionKey: string;
  readonly orderKey: string;
  readonly value: T;
}

export interface CustomizationExternalStateSession<T extends JsonValue = JsonValue> {
  /** Trusted host accounting of retained working bytes and all spill files. Required for grouped history. */
  readonly resourceUsage?: () => { readonly workingBytes: number; readonly spillBytes: number };
  /** Store one SDK-owned record. The adapter must not mutate it. */
  readonly append: (record: CustomizationExternalRecord<T>) => void | Promise<void>;
  /** Finish writes before reads begin. */
  readonly seal: () => void | Promise<void>;
  /** Reopen all records exactly once in partitionKey/orderKey/id order. */
  readonly read: () => AsyncIterable<CustomizationExternalRecord<T>>;
  /** Release the run's temporary namespace for every terminal outcome. */
  readonly dispose: (outcome: "completed" | "failed" | "cancelled") => void | Promise<void>;
}

export interface CustomizationExternalStateAdapter {
  readonly open: <T extends JsonValue>(request: {
    readonly runId: string;
    readonly signal?: AbortSignal;
  }) => CustomizationExternalStateSession<T> | Promise<CustomizationExternalStateSession<T>>;
}

export interface ExternalSortCustomizationRecordsInput<T extends JsonValue = JsonValue> {
  readonly runId: string;
  readonly records: Iterable<CustomizationExternalRecord<T>> | AsyncIterable<CustomizationExternalRecord<T>>;
  readonly adapter: CustomizationExternalStateAdapter;
  readonly limits: CustomizationResourceLimits;
  readonly signal?: AbortSignal;
}

const encoder = new TextEncoder();

function cancelled(signal: AbortSignal | undefined): void {
  if (signal?.aborted)
    throw new CustomizationExecutionError("CANCELLED", "Customization execution was cancelled");
}

function validateRecord<T extends JsonValue>(
  value: CustomizationExternalRecord<T>,
  maximumRecordBytes: number,
): CustomizationExternalRecord<T> {
  if (!isDiagnosticToken(value.id))
    throw new CustomizationExecutionError("EXTERNAL_STATE_FAILURE", "External record id is invalid");
  if (
    !isWellFormedDiagnosticString(value.partitionKey) ||
    !isWellFormedDiagnosticString(value.orderKey)
  )
    throw new CustomizationExecutionError(
      "EXTERNAL_STATE_FAILURE",
      "External record sort keys must be valid strings without U+0000",
    );
  const owned = snapshotDiagnosticJson(value);
  const bytes = encoder.encode(JSON.stringify(owned)).byteLength;
  if (bytes > maximumRecordBytes)
    throw new CustomizationExecutionError(
      "RESOURCE_LIMIT",
      `External record ${value.id} requires ${bytes} encoded bytes, exceeding maximumRecordBytes ${maximumRecordBytes}`,
    );
  return owned;
}

function compareRecords(
  left: CustomizationExternalRecord,
  right: CustomizationExternalRecord,
): number {
  for (const key of ["partitionKey", "orderKey", "id"] as const) {
    const result = left[key] === right[key] ? 0 : left[key] < right[key] ? -1 : 1;
    if (result !== 0) return result;
  }
  return 0;
}

/**
 * Validates, owns, externally stages, and deterministically rereads records.
 * The adapter supplies storage mechanics only; it cannot choose actuarial
 * ordering or publish a completed run. Early consumer return is cancellation.
 */
export async function* externalSortCustomizationRecords<T extends JsonValue>(
  input: ExternalSortCustomizationRecordsInput<T>,
): AsyncGenerator<CustomizationExternalRecord<T>> {
  const runId = input.runId;
  const records = input.records;
  const adapter = input.adapter;
  const signal = input.signal;
  let limits: CustomizationResourceLimits;
  try {
    limits = parseCustomizationResourceLimits(input.limits);
  } catch (error) {
    throw new CustomizationExecutionError(
      "RESOURCE_LIMIT",
      error instanceof Error ? error.message : "Invalid customization resource limits",
    );
  }
  if (!isDiagnosticToken(runId))
    throw new CustomizationExecutionError("EXTERNAL_STATE_FAILURE", "runId is invalid");
  cancelled(signal);

  let session: CustomizationExternalStateSession<T> | undefined;
  let outcome: "completed" | "failed" | "cancelled" = "failed";
  let primary: unknown;
  try {
    session = await adapter.open<T>({
      runId,
      ...(signal === undefined ? {} : { signal }),
    });
    let written = 0;
    let writeIdentity = 0n, readIdentity = 0n;
    const mask = (1n << 64n) - 1n;
    const identity = (record: CustomizationExternalRecord) => BigInt(`0x${fnv1a64(canonicalJson(record))}`);
    for await (const raw of records) {
      cancelled(signal);
      written += 1;
      if (written > limits.maximumInputRecords)
        throw new CustomizationExecutionError(
          "RESOURCE_LIMIT",
          `Input exceeds maximumInputRecords ${limits.maximumInputRecords}`,
        );
      const record = validateRecord(raw, limits.maximumRecordBytes);
      writeIdentity = (writeIdentity + identity(record)) & mask;
      await session.append(record);
    }
    cancelled(signal);
    await session.seal();

    let read = 0;
    let partitionRecords = 0;
    let prior: CustomizationExternalRecord<T> | undefined;
    for await (const raw of session.read()) {
      cancelled(signal);
      const record = validateRecord(raw, limits.maximumRecordBytes);
      if (prior !== undefined && compareRecords(prior, record) > 0)
        throw new CustomizationExecutionError(
          "EXTERNAL_STATE_ORDER",
          "External records are not in nondecreasing partitionKey/orderKey/id order",
        );
      partitionRecords = prior?.partitionKey === record.partitionKey ? partitionRecords + 1 : 1;
      if (partitionRecords > limits.maximumPartitionRecords)
        throw new CustomizationExecutionError(
          "RESOURCE_LIMIT",
          `Partition ${record.partitionKey} exceeds maximumPartitionRecords ${limits.maximumPartitionRecords}`,
        );
      read += 1;
      if (read > written)
        throw new CustomizationExecutionError(
          "EXTERNAL_STATE_FAILURE",
          "External state returned more records than were written",
        );
      readIdentity = (readIdentity + identity(record)) & mask;
      prior = record;
      yield record;
    }
    if (read !== written)
      throw new CustomizationExecutionError(
        "EXTERNAL_STATE_TRUNCATED",
        `External state returned ${read} of ${written} records`,
      );
    if (readIdentity !== writeIdentity) throw new CustomizationExecutionError("EXTERNAL_STATE_FAILURE", "External state changed the staged record multiset");
    outcome = "completed";
  } catch (error) {
    primary = error;
    outcome =
      error instanceof CustomizationExecutionError && error.code === "CANCELLED"
        ? "cancelled"
        : "failed";
    throw error;
  } finally {
    if (session !== undefined) {
      try {
        const terminal =
          outcome === "completed"
            ? "completed"
            : primary === undefined
              ? "cancelled"
              : outcome;
        await session.dispose(terminal);
      } catch (error) {
        if (primary === undefined)
          throw new CustomizationExecutionError(
            "EXTERNAL_STATE_FAILURE",
            `External state cleanup failed: ${error instanceof Error ? error.message : String(error)}`,
          );
      }
    }
  }
}