/**
 * Durable A2A Task-record module for the v0.4 GetTask lifecycle.
 *
 * The sesh shim's `getTask` reads `sesh_tasks_<scope-kind>_<scope-id>` KV at
 * key = taskId and returns the stored JSON **verbatim** (byte-passthrough via
 * `seshtask.GetRaw`). So the record itself must carry the A2A wire fields:
 * `kind`, `id`, `contextId`, `status.state`, optional `history`, and
 * `artifacts`.
 *
 * The shim mints the initial record on `message/send` as
 *   {"id":..,"kind":"task","contextId":..,"status":{"state":"TASK_STATE_SUBMITTED"}}
 * (see sesh `internal/shim/methods/send.go`) and only ever transitions it to
 * `TASK_STATE_CANCELED` via a CAS `kv.Update` (`cancel.go`). There is NO
 * COMPLETED/WORKING writer on the shim side — that is what this module adds on
 * the adapter side: drive SUBMITTED → WORKING → COMPLETED and embed artifacts
 * in `.artifacts[]`, all via read-modify-write CAS that preserves every
 * shim-minted field unmolested.
 *
 * Schema source (extracted, not guessed):
 *   - a2a-go `a2a/core.go`: `type Task`, `type TaskStatus`, `type Artifact`.
 *     TaskStatus serializes `{ "state": <TaskState>, "message"?, "timestamp"? }`.
 *     TaskState literals are the `TASK_STATE_*` strings (see taskstate.ts).
 *   - sesh `internal/shim/methods/send.go` / `gettask.go` / `cancel.go`: the
 *     minted shape + the CAS update pattern (`kv.Update(id, bytes, rev)` with
 *     retry on revision conflict).
 *
 * The SDK never opens its own KV — the caller (adapter) passes an
 * already-opened handle (typically `scope.tasksBucket(scope)`).
 */

import type { KV } from "@nats-io/kv";

import type { Artifact } from "./artifacts.ts";
import { SDKError } from "./errors.ts";
import { tasksBucket } from "./scope.ts";
import type { ScopeRef } from "./stream.ts";
import type { Message } from "./types.ts";
import type { TaskStateWire } from "./taskstate.ts";

// Re-export the wire types callers will need so they don't import twice.
export type { Artifact } from "./artifacts.ts";
export type { Message } from "./types.ts";
export type { TaskStateWire } from "./taskstate.ts";

/**
 * The slice of `@nats-io/kv`'s `Kvm` this module needs to open (create-if-
 * missing) a task bucket. Structural so callers can pass a real `Kvm`
 * instance without the SDK taking a hard dependency on its concrete class —
 * `new Kvm(js).create(name)` already satisfies this shape.
 */
export interface KvOpener {
  create(name: string): Promise<KV>;
}

/**
 * A2A TaskStatus. Mirrors a2a-go `core.go` `type TaskStatus`:
 *   State     TaskState  `json:"state"`
 *   Message   *Message   `json:"message,omitempty"`
 *   Timestamp *time.Time `json:"timestamp,omitempty"`
 */
export interface TaskStatus {
  state: TaskStateWire;
  message?: Message;
  /** RFC3339/ISO-8601 timestamp string, matching Go's `time.Time` JSON. */
  timestamp?: string;
}

/**
 * The durable Task record as stored in KV and returned verbatim by the shim's
 * GetTask. Mirrors a2a-go `core.go` `type Task`:
 *   ID        string      `json:"id"`
 *   ContextID string      `json:"contextId"`
 *   Kind      string      `json:"kind"`   (always "task")
 *   Status    TaskStatus  `json:"status"`
 *   History   []*Message  `json:"history,omitempty"`
 *   Artifacts []*Artifact `json:"artifacts,omitempty"`
 *   Metadata  map         `json:"metadata,omitempty"`
 *
 * An index signature is included so unknown/forward-compat fields the shim may
 * mint are preserved through CAS read-modify-write rather than dropped.
 */
export interface TaskRecord {
  id: string;
  kind: "task";
  contextId: string;
  status: TaskStatus;
  history?: Message[];
  artifacts?: Artifact[];
  metadata?: Record<string, unknown>;
  [extra: string]: unknown;
}

/** Number of CAS attempts before giving up (1 initial + 1 retry on conflict). */
const CAS_ATTEMPTS = 2 as const;

export interface UpdateTaskStatusOpts {
  /** Optional human-readable status message embedded at `status.message`. */
  message?: Message;
  /**
   * Optional `status.timestamp`. Pass a string to set explicitly, or `true`
   * to stamp `new Date().toISOString()`. Omit to leave timestamp untouched.
   */
  timestamp?: string | true;
}

export interface CompleteTaskOpts {
  /** Artifacts to embed in the record's `.artifacts[]` on completion. */
  artifacts?: Artifact[];
  /** Optional `status.message`. */
  message?: Message;
  /** See `UpdateTaskStatusOpts.timestamp`. */
  timestamp?: string | true;
}

// ── Serialisation ────────────────────────────────────────────────────────
// Symmetric with artifacts.ts: encode Uint8Array (raw Part bytes) as a JSON
// array of bytes so the record is plain, valid JSON the shim can pass through.
function replacer(_key: string, value: unknown): unknown {
  if (value instanceof Uint8Array) return Array.from(value);
  return value;
}

function encode(record: TaskRecord): Uint8Array {
  return new TextEncoder().encode(JSON.stringify(record, replacer));
}

function decode(bytes: Uint8Array): TaskRecord {
  return JSON.parse(new TextDecoder().decode(bytes)) as TaskRecord;
}

function resolveTimestamp(ts: string | true | undefined): string | undefined {
  if (ts === undefined) return undefined;
  return ts === true ? new Date().toISOString() : ts;
}

/**
 * A KV update failed because the stored revision moved out from under us
 * (optimistic-concurrency conflict). The `@nats-io/kv` client surfaces this
 * as a `wrong last sequence` API error rather than a typed class, so we
 * pattern-match the message. Returns true if the error looks like a CAS miss.
 */
function isCASConflict(err: unknown): boolean {
  const msg =
    typeof err === "object" && err !== null && "message" in err
      ? String((err as { message: unknown }).message)
      : String(err);
  return /wrong last sequence|key exists|10071/i.test(msg);
}

// ── Read ───────────────────────────────────────────────────────────────────
/**
 * Read the existing (shim-minted) Task record + its KV revision. Returns
 * `null` if no record exists (or the key was deleted/purged). The revision is
 * what callers feed back into `kv.update` for CAS.
 */
export async function readTask(
  kv: KV,
  taskId: string,
): Promise<{ task: TaskRecord; revision: number } | null> {
  const entry = await kv.get(taskId);
  if (!entry) return null;
  if (entry.operation === "DEL" || entry.operation === "PURGE") return null;
  return { task: decode(entry.value), revision: entry.revision };
}

// ── Open / mint ─────────────────────────────────────────────────────────
export interface OpenTaskOpts {
  /** A2A `contextId` to mint into the initial record. Defaults to `taskId`. */
  contextId?: string;
  /**
   * Optional initial `status.message` (e.g. the inbound prompt). Omitted from
   * the minted record when absent, matching the shim's bare SUBMITTED mint.
   */
  message?: Message;
}

/**
 * Open the scope's task bucket (create-if-missing) and mint the INITIAL
 * SUBMITTED Task record for `taskId` if it does not already exist.
 *
 * This is the responder-side counterpart to the Go sesh-shim's
 * `message/send` mint: on a shim-less mesh (plain external NATS hub) no code
 * writes the initial record, so the read-modify-write mutators
 * (`updateTaskStatus`/`completeTask`) hit `TaskNotFound` at completion. This
 * helper makes the adapter self-sufficient — call it at task-open, before the
 * SUBMITTED → WORKING transition.
 *
 * Idempotent: if a record already exists at `taskId` (shim-minted, or a
 * re-delivery of the same prompt) it is left **unmolested** — we never clobber
 * a richer existing record with a bare SUBMITTED one. The bucket open is
 * `kvm.create`, which is create-if-missing, unifying the historical
 * `kvm.create` vs `kvm.open` drift across adapters onto one path.
 *
 * NON-best-effort by design: a failure to ensure the bucket or write the
 * initial record throws `TaskOpenError`. A dropped task record must be loud at
 * open, never silently discovered at completion.
 *
 * Returns the opened KV handle (so the caller can drive the subsequent
 * WORKING/COMPLETED CAS mutations on it) plus the resolved initial record.
 */
export async function openTask(
  kvm: KvOpener,
  scope: ScopeRef,
  taskId: string,
  opts: OpenTaskOpts = {},
): Promise<{ kv: KV; task: TaskRecord }> {
  const bucket = tasksBucket(scope);

  let kv: KV;
  try {
    kv = await kvm.create(bucket);
  } catch (err) {
    throw new TaskOpenError(taskId, `open task bucket "${bucket}"`, err);
  }

  // Already present (shim-minted or re-delivery)? Leave it untouched.
  const existing = await readTask(kv, taskId);
  if (existing) return { kv, task: existing.task };

  const contextId = opts.contextId ?? taskId;
  const record: TaskRecord = {
    id: taskId,
    kind: "task",
    contextId,
    status: { state: "TASK_STATE_SUBMITTED" },
  };
  if (opts.message !== undefined) record.status.message = opts.message;

  try {
    await kv.create(taskId, encode(record));
  } catch (err) {
    // A concurrent opener (another in-flight delivery of the same task) may
    // have minted the record between our read and create. Treat a now-present
    // record as success rather than a failure — the record exists, which is
    // the post-condition callers depend on.
    const raced = await readTask(kv, taskId);
    if (raced) return { kv, task: raced.task };
    throw new TaskOpenError(taskId, `write initial SUBMITTED record`, err);
  }

  return { kv, task: record };
}

// ── CAS read-modify-write core ───────────────────────────────────────────
/**
 * Read the current record, apply `mutate` to a deep-ish copy, and CAS-write it
 * back with one retry on revision conflict. `mutate` MUST preserve all fields
 * it does not intend to change — it receives the parsed record by reference.
 *
 * Throws `TaskNotFound` if the record does not exist (the shim mints it on
 * `message/send`, so a missing record means the caller is racing ahead of the
 * inbound message or addressing the wrong task).
 */
async function casUpdate(
  kv: KV,
  taskId: string,
  mutate: (record: TaskRecord) => void,
): Promise<TaskRecord> {
  let lastErr: unknown;
  for (let attempt = 0; attempt < CAS_ATTEMPTS; attempt++) {
    const current = await readTask(kv, taskId);
    if (!current) throw new TaskNotFound(taskId);
    mutate(current.task);
    try {
      await kv.update(taskId, encode(current.task), current.revision);
      return current.task;
    } catch (err) {
      lastErr = err;
      if (isCASConflict(err) && attempt + 1 < CAS_ATTEMPTS) continue;
      throw err;
    }
  }
  throw new TaskCASConflict(taskId, lastErr);
}

// ── Public mutators ────────────────────────────────────────────────────────
/**
 * Transition the Task to `state` (e.g. WORKING, COMPLETED, FAILED) via CAS.
 * Preserves `id`, `kind`, `contextId`, `history`, `artifacts`, and any other
 * fields the shim minted — only `status` is touched. Returns the post-write
 * record.
 */
export async function updateTaskStatus(
  kv: KV,
  taskId: string,
  state: TaskStateWire,
  opts: UpdateTaskStatusOpts = {},
): Promise<TaskRecord> {
  const ts = resolveTimestamp(opts.timestamp);
  return casUpdate(kv, taskId, (record) => {
    const status: TaskStatus = { ...record.status, state };
    if (opts.message !== undefined) status.message = opts.message;
    if (ts !== undefined) status.timestamp = ts;
    record.status = status;
  });
}

/**
 * CAS-append `artifact` to the record's `.artifacts[]` (creating the array if
 * absent). Leaves `status` and every other field intact. Returns the
 * post-write record. The embedded artifact is what GetTask surfaces in
 * `.artifacts[]`; the separate `sesh_artifacts_*` bucket feeds the streaming
 * watch path, not GetTask.
 */
export async function appendTaskArtifact(
  kv: KV,
  taskId: string,
  artifact: Artifact,
): Promise<TaskRecord> {
  return casUpdate(kv, taskId, (record) => {
    const artifacts = Array.isArray(record.artifacts) ? [...record.artifacts] : [];
    artifacts.push(artifact);
    record.artifacts = artifacts;
  });
}

/**
 * Convenience: set state COMPLETED and embed `opts.artifacts` (appended to any
 * existing) in a single CAS, optionally stamping a status message/timestamp.
 * Returns the post-write record.
 */
export async function completeTask(
  kv: KV,
  taskId: string,
  opts: CompleteTaskOpts = {},
): Promise<TaskRecord> {
  const ts = resolveTimestamp(opts.timestamp);
  return casUpdate(kv, taskId, (record) => {
    const status: TaskStatus = { ...record.status, state: "TASK_STATE_COMPLETED" };
    if (opts.message !== undefined) status.message = opts.message;
    if (ts !== undefined) status.timestamp = ts;
    record.status = status;

    if (opts.artifacts && opts.artifacts.length > 0) {
      const artifacts = Array.isArray(record.artifacts) ? [...record.artifacts] : [];
      artifacts.push(...opts.artifacts);
      record.artifacts = artifacts;
    }
  });
}

// ── Errors ───────────────────────────────────────────────────────────────
/** No durable Task record exists for the given taskId. */
export class TaskNotFound extends SDKError {
  constructor(public readonly taskId: string) {
    super(`task ${taskId} not found in KV`);
    this.name = "TaskNotFound";
  }
}

/**
 * `openTask` could not ensure the task bucket / write the initial record.
 * Non-best-effort: surfaced loudly at open so a dropped task record is never
 * silently discovered at completion.
 */
export class TaskOpenError extends SDKError {
  constructor(
    public readonly taskId: string,
    public readonly phase: string,
    cause?: unknown,
  ) {
    super(`task ${taskId}: open failed during ${phase}`, cause);
    this.name = "TaskOpenError";
  }
}

/** CAS retries exhausted — the record kept changing out from under us. */
export class TaskCASConflict extends SDKError {
  constructor(
    public readonly taskId: string,
    cause?: unknown,
  ) {
    super(`task ${taskId}: CAS conflict after ${CAS_ATTEMPTS} attempts`, cause);
    this.name = "TaskCASConflict";
  }
}