import {
  canonicalJson,
  fnv1a64,
  prepareHistoricalDataset,
  type CustomizationCapabilityFinding,
  type CustomizationResourceLimits,
  type HistoricalDatasetInput,
  type HistoricalObservationRecord,
  type HistoricalRecordDisposition,
  type JsonValue,
  type PreparedHistoricalObservation,
} from "@actuarial-ts/core";
import { parseHistoricalDatasetInput, parseHistoricalObservationRecord } from "./customizationContracts.js";
import {
  CustomizationExecutionError,
  externalSortCustomizationRecords,
  type CustomizationExternalRecord,
  type CustomizationExternalStateAdapter,
} from "./customizationExternalState.js";

export type PreparedHistoryStreamEvent =
  | { readonly kind: "observation"; readonly observation: PreparedHistoricalObservation }
  | { readonly kind: "disposition"; readonly disposition: HistoricalRecordDisposition }
  | { readonly kind: "finding"; readonly finding: CustomizationCapabilityFinding };

function partitionKey(record: HistoricalObservationRecord): string {
  return canonicalJson([record.sourceNamespace, record.claimId]);
}

function orderKey(record: HistoricalObservationRecord, precedence: HistoricalDatasetInput["revisionPrecedence"]): string {
  const order = precedence === "sequence"
    ? String(record.revision.sequence ?? -1).padStart(20, "0")
    : record.revision.correctedAt ?? "";
  return canonicalJson([record.recordId, record.evaluationDate ?? record.effectiveDate ?? "", order, record.revision.id]);
}

/**
 * Sorts raw revisions in host-owned state and prepares one claim partition at
 * a time. Retained SDK state is bounded by maximumPartitionRecords; oversized
 * claims fail rather than silently spilling or dropping observations.
 */
export async function* streamPreparedHistoricalDataset(input: {
  readonly runId: string;
  readonly header: Omit<HistoricalDatasetInput, "records">;
  readonly records: AsyncIterable<unknown> | Iterable<unknown>;
  readonly adapter: CustomizationExternalStateAdapter;
  readonly limits: CustomizationResourceLimits;
  readonly signal?: AbortSignal;
}): AsyncGenerator<PreparedHistoryStreamEvent> {
  const header = parseHistoricalDatasetInput({ ...input.header, records: [] });
  async function* externalRecords(): AsyncGenerator<CustomizationExternalRecord<JsonValue>> {
    for await (const raw of input.records) {
      const record = parseHistoricalObservationRecord(raw);
      const stable = canonicalJson([record.sourceNamespace, record.claimId, record.recordId, record.revision.id]);
      yield {
        id: `record-${fnv1a64(stable)}`,
        partitionKey: partitionKey(record),
        orderKey: orderKey(record, header.revisionPrecedence),
        value: record as unknown as JsonValue,
      };
    }
  }
  let currentPartition: string | undefined;
  let current: HistoricalObservationRecord[] = [];
  let partitionBytes = 0;
  const encoder = new TextEncoder();
  const flush = function* (): Generator<PreparedHistoryStreamEvent> {
    if (current.length === 0) return;
    const sourceNamespace = current[0]!.sourceNamespace;
    const claimId = current[0]!.claimId;
    const prepared = prepareHistoricalDataset({
      ...header,
      ...(header.openingBalances === undefined ? {} : { openingBalances: header.openingBalances.filter((balance) => balance.sourceNamespace === sourceNamespace && balance.claimId === claimId) }),
      records: current,
    });
    for (const disposition of prepared.dispositions) yield { kind: "disposition", disposition };
    for (const finding of prepared.findings) yield { kind: "finding", finding };
    for (const observation of prepared.observations) yield { kind: "observation", observation };
    current = [];
    partitionBytes = 0;
  };
  for await (const external of externalSortCustomizationRecords({
    runId: input.runId,
    records: externalRecords(),
    adapter: input.adapter,
    limits: input.limits,
    ...(input.signal === undefined ? {} : { signal: input.signal }),
  })) {
    if (currentPartition !== undefined && external.partitionKey !== currentPartition) yield* flush();
    currentPartition = external.partitionKey;
    partitionBytes += encoder.encode(canonicalJson(external.value)).byteLength;
    if (partitionBytes > input.limits.maximumInMemoryBytes) throw new CustomizationExecutionError("RESOURCE_LIMIT", "Claim partition exceeds maximumInMemoryBytes");
    current.push(parseHistoricalObservationRecord(external.value));
  }
  yield* flush();
}