/**
 * Messages KV helpers. Wraps `@nats-io/kv` with CAS-on-create idempotency and
 * Part materialisation (auto-offload to Object Store above 1 MB).
 *
 * The SDK never opens its own KV — the caller (adapter) passes an
 * already-opened handle. Bucket naming is sesh-ops's concern.
 */

import type { KV, KvWatchEntry } from "@nats-io/kv";
import { KvWatchInclude } from "@nats-io/kv";
import type { ObjectStore } from "@nats-io/obj";
import type { QueuedIterator } from "@nats-io/transport-node";

import { canonicalBytes } from "./canonicalJson.ts";
import { MessageIDConflict, MaterializePartError } from "./errors.ts";
import { objectKey, objectURL, putObject, type ObjectURLParts } from "./objects.ts";
import type { Message, Part } from "./types.ts";

// Re-export the canonical wire types so callers don't have to import twice.
export type { Message, Part, Role } from "./types.ts";

// ── Per-variant constructor inputs (sugar over the flat Part shape) ────
export interface TextPartInput {
  text: string;
  mediaType?: string;
  filename?: string;
  metadata?: Record<string, unknown>;
}

export interface DataPartInput {
  data: unknown;
  mediaType?: string;
  metadata?: Record<string, unknown>;
}

export interface FilePartURLInput {
  url: string;
  mediaType: string;
  filename?: string;
  metadata?: Record<string, unknown>;
}

export interface FilePartRawInput {
  raw: Uint8Array;
  mediaType: string;
  filename?: string;
  metadata?: Record<string, unknown>;
}

/**
 * Constructors. These produce the flat `Part` shape from `types.ts`. They
 * exist for ergonomics and to give callers a single, named entry-point per
 * variant rather than asking them to remember which discriminator field
 * makes a Part a text part vs a data part.
 */
export function TextPart(p: TextPartInput): Part {
  const out: Part = { text: p.text };
  if (p.mediaType !== undefined) out.mediaType = p.mediaType;
  if (p.filename !== undefined) out.filename = p.filename;
  if (p.metadata !== undefined) out.metadata = p.metadata;
  return out;
}

export function DataPart(p: DataPartInput): Part {
  const out: Part = { data: p.data };
  if (p.mediaType !== undefined) out.mediaType = p.mediaType;
  if (p.metadata !== undefined) out.metadata = p.metadata;
  return out;
}

export function FilePartURL(p: FilePartURLInput): Part {
  const out: Part = { url: p.url, mediaType: p.mediaType };
  if (p.filename !== undefined) out.filename = p.filename;
  if (p.metadata !== undefined) out.metadata = p.metadata;
  return out;
}

export function FilePartRaw(p: FilePartRawInput): Part {
  const out: Part = { raw: p.raw, mediaType: p.mediaType };
  if (p.filename !== undefined) out.filename = p.filename;
  if (p.metadata !== undefined) out.metadata = p.metadata;
  return out;
}

// ── Materialise (auto-offload large raw Parts) ────────────────────────
/** Max raw bytes carried inline before auto-offload to Object Store. */
export const PART_INLINE_MAX_BYTES = 1_048_576 as const; // 1 MB

export interface MaterializeOpts {
  inlineMaxBytes?: number;
  /** Required when raw exceeds threshold — used to compute obj:// URI + os key. */
  scope: Pick<ObjectURLParts, "scopeKind" | "scopeId">;
  taskId: string;
  /** Typically a fresh ULID. Called once per offloaded part. */
  artifactIdFor: (part: Part) => string;
}

/**
 * If `part` carries `raw` bytes larger than `inlineMaxBytes`, write them to
 * Object Store and return an equivalent URL-bearing Part. Otherwise return
 * the part unchanged. Pure for Text/Data/URL parts (no raw field).
 */
export async function materializePart(
  os: ObjectStore,
  part: Part,
  opts: MaterializeOpts,
): Promise<Part> {
  if (!part.raw) return part;
  const threshold = opts.inlineMaxBytes ?? PART_INLINE_MAX_BYTES;
  if (part.raw.byteLength <= threshold) return part;
  if (!part.mediaType) {
    throw new MaterializePartError("raw Part requires mediaType to materialize");
  }
  const artifactId = opts.artifactIdFor(part);
  const key = objectKey({ taskId: opts.taskId, artifactId });
  try {
    await putObject(os, key, part.raw);
  } catch (e) {
    throw new MaterializePartError(`failed to offload Part to Object Store (key=${key})`, e);
  }
  const url = objectURL({ ...opts.scope, taskId: opts.taskId, artifactId });
  const out: Part = { url, mediaType: part.mediaType };
  if (part.filename !== undefined) out.filename = part.filename;
  if (part.metadata !== undefined) out.metadata = part.metadata;
  return out;
}

// ── KV layout ─────────────────────────────────────────────────────────
/** Compose the KV key for a Message: `<task-id>.<message-id>`. */
export function messageKey(taskId: string, messageId: string): string {
  return `${taskId}.${messageId}`;
}

// ── Append / read / watch ─────────────────────────────────────────────
export type AppendResult =
  | { status: "created"; revision: number }
  | { status: "duplicate-identical"; revision: number };

/**
 * Strict create. Throws `MessageIDConflict` on duplicate, even if identical.
 * Use `idempotentAppend` for the more forgiving variant.
 */
export async function appendMessage(
  kv: KV,
  taskId: string,
  message: Message,
): Promise<{ revision: number }> {
  const key = messageKey(taskId, message.messageId);
  const bytes = canonicalBytes(message);
  try {
    const revision = await kv.create(key, bytes);
    return { revision };
  } catch {
    const existing = await kv.get(key);
    throw new MessageIDConflict(
      taskId,
      message.messageId,
      existing?.value ?? new Uint8Array(),
      bytes,
    );
  }
}

/**
 * Idempotent variant: on duplicate, byte-compare canonical-JSON of the
 * incoming message against the stored entry. If bytes match → silent
 * success. If they differ → `MessageIDConflict`.
 */
export async function idempotentAppend(
  kv: KV,
  taskId: string,
  message: Message,
): Promise<AppendResult> {
  const key = messageKey(taskId, message.messageId);
  const bytes = canonicalBytes(message);
  try {
    const revision = await kv.create(key, bytes);
    return { status: "created", revision };
  } catch {
    const existing = await kv.get(key);
    if (!existing) {
      // Race: deleted between create-fail and get. Retry once.
      const revision = await kv.create(key, bytes);
      return { status: "created", revision };
    }
    if (bytesEqual(existing.value, bytes)) {
      return { status: "duplicate-identical", revision: existing.revision };
    }
    throw new MessageIDConflict(taskId, message.messageId, existing.value, bytes);
  }
}

export async function readMessage(
  kv: KV,
  taskId: string,
  messageId: string,
): Promise<Message | null> {
  const entry = await kv.get(messageKey(taskId, messageId));
  if (!entry) return null;
  if (entry.operation === "DEL" || entry.operation === "PURGE") return null;
  return JSON.parse(new TextDecoder().decode(entry.value)) as Message;
}

export interface WatchMessagesOpts {
  /** If true, replay all stored history before yielding live updates. */
  includeHistory?: boolean;
}

/**
 * Subscribe to all messages for a task. Caller is responsible for breaking
 * out of the iterator (it will run forever otherwise). Yields parsed Message
 * values plus the raw KV watch entry (for revision / operation inspection).
 */
export async function* watchMessages(
  kv: KV,
  taskId: string,
  opts: WatchMessagesOpts = {},
): AsyncGenerator<{ message: Message; entry: KvWatchEntry }> {
  const include = opts.includeHistory
    ? KvWatchInclude.AllHistory
    : KvWatchInclude.LastValue;
  const iter: QueuedIterator<KvWatchEntry> = await kv.watch({
    key: `${taskId}.*`,
    include,
  });
  try {
    for await (const entry of iter) {
      if (entry.operation !== "PUT") continue;
      const message = JSON.parse(new TextDecoder().decode(entry.value)) as Message;
      yield { message, entry };
    }
  } finally {
    iter.stop();
  }
}

// ── internals ─────────────────────────────────────────────────────────
function bytesEqual(a: Uint8Array, b: Uint8Array): boolean {
  if (a.byteLength !== b.byteLength) return false;
  for (let i = 0; i < a.byteLength; i++) if (a[i] !== b[i]) return false;
  return true;
}

// Re-export so callers can `instanceof`-check without a second import.
export { MessageIDConflict, MaterializePartError } from "./errors.ts";