import { createHash } from "node:crypto";
import { z } from "zod";
import { answerUiBlockSchema } from "../chat/answer-ui.js";
import type { SessionWorkspaceContext } from "../runtime/conversationWorkspace.js";
import type { SessionWorkspaceStore } from "./workspace-store.js";

const id = z.string().min(1).max(200);
const index = z.number().int().nonnegative();
const indices = (max: number) => z.array(index).max(max).refine((items) => new Set(items).size === items.length);
const tableStateSchema = (maxRows: number) => z.object({ query: z.string().max(1000), sort: z.object({ column: index, descending: z.boolean() }).strict().nullable(), selection: indices(maxRows) }).strict();
const nodeSchema = z.discriminatedUnion("type", [
  z.object({ type: z.literal("Chart"), view: z.enum(["line", "bar", "table"]), hidden: indices(8), range: z.tuple([index, index]), point: index, inspecting: z.boolean(), table: tableStateSchema(100).optional() }).strict(),
  z.object({ type: z.literal("Table"), query: z.string().max(1000), sort: z.object({ column: index, descending: z.boolean() }).strict().nullable(), selection: indices(50) }).strict(),
  z.object({ type: z.literal("FollowUpForm"), values: z.record(z.string().max(500)) }).strict(),
  z.object({ type: z.literal("Evidence"), open: z.boolean() }).strict(),
]);
export const answerViewStateIdentitySchema = z.object({ message_id: id, block_id: z.string().regex(/^[a-zA-Z0-9_-]{1,80}$/) }).strict();
export const answerViewStatePutSchema = answerViewStateIdentitySchema.extend({
  schema_digest: z.string().regex(/^[a-f0-9]{64}$/), expected_revision: z.number().int().nonnegative().max(Number.MAX_SAFE_INTEGER - 1),
  request_id: z.string().min(1).max(128), nodes: z.record(nodeSchema).refine((nodes) => Object.keys(nodes).length <= 48),
  block: z.object({ evidenceFilter: z.enum(["all", "observed", "human_input", "inference", "demo"]) }).strict(),
}).strict();
export type AnswerViewStatePut = z.infer<typeof answerViewStatePutSchema>;
export type AnswerViewState = {
  version: "answer-view-state@1"; message_id: string; block_id: string; revision: number;
  schema_digest: string; nodes: Record<string, z.infer<typeof nodeSchema>>;
  block: { evidenceFilter: string }; updated_at: string | null;
};
type Spec = z.infer<typeof answerUiBlockSchema>["spec"];
type Basis = { type: string; columns?: string[]; rowsDigest?: string; dataDigest?: string; labels?: string[]; series?: string[]; fields?: Record<string, string> };
type Entry = { state: AnswerViewState; basis: Record<string, Basis> };
type RecordFile = { version: "answer-view-state@1"; entries: Record<string, Entry>; requestOrder: string[]; requests: Record<string, { digest: string; response: AnswerViewState }> };
export class AnswerViewStateError extends Error {
  constructor(readonly status: number, readonly code: string, message: string, readonly current?: AnswerViewState) { super(message); }
}
const invalid = () => new AnswerViewStateError(400, "ANSWER_VIEW_STATE_INVALID", "invalid answer view state");
const isRecord = (value: unknown): value is Record<string, unknown> => typeof value === "object" && value !== null && !Array.isArray(value);
export function canonicalAnswerJson(value: unknown): string {
  if (Array.isArray(value)) return `[${value.map(canonicalAnswerJson).join(",")}]`;
  if (isRecord(value)) return `{${Object.keys(value).sort().map((key) => `${JSON.stringify(key)}:${canonicalAnswerJson(value[key])}`).join(",")}}`;
  return JSON.stringify(value);
}
const digest = (value: unknown) => createHash("sha256").update(canonicalAnswerJson(value)).digest("hex");
const bytes = (value: unknown) => Buffer.byteLength(JSON.stringify(value) ?? "null", "utf8");
const same = (a: unknown, b: unknown) => canonicalAnswerJson(a) === canonicalAnswerJson(b);
const own = <T>(record: Record<string, T>, key: string) => Object.hasOwn(record, key) ? record[key] : undefined;

function currentSpec(messages: unknown, messageId: string, blockId: string): Spec {
  const items = isRecord(messages) && Array.isArray(messages.messages) ? messages.messages : [];
  const message = items.find((item) => isRecord(item) && item.id === messageId && item.role === "assistant");
  const parts = isRecord(message) && Array.isArray(message.parts) ? message.parts : [];
  const finalPart = [...parts].reverse().find((part) => isRecord(part) && part.type === "data-final_answer" && isRecord(part.data));
  if (isRecord(finalPart) && isRecord(finalPart.data) && (!Array.isArray(finalPart.data.ui_block_ids) || !finalPart.data.ui_block_ids.includes(blockId))) {
    throw new AnswerViewStateError(404, "ANSWER_VIEW_BLOCK_NOT_FOUND", "answer UI block not found");
  }
  const candidates = parts.filter((part) => isRecord(part) && part.type === "data-answer_ui" && isRecord(part.data) && part.data.block_id === blockId);
  if (candidates.length === 1) {
    const data = (candidates[0] as { data: Record<string, unknown> }).data;
    const parsed = answerUiBlockSchema.safeParse({ version: data.version, block_id: data.block_id, phase: data.phase, fallback_text: data.fallback_text, spec: data.spec });
    if (parsed.success && parsed.data.phase !== "process") return parsed.data.spec;
  }
  throw new AnswerViewStateError(404, "ANSWER_VIEW_BLOCK_NOT_FOUND", "answer UI block not found");
}
function basisFor(spec: Spec): Record<string, Basis> {
  return Object.fromEntries(Object.entries(spec.elements).map(([key, node]) => {
    const basis: Basis = { type: node.type };
    if (node.type === "Chart") { basis.labels = node.props.x_data as string[]; basis.series = (node.props.series as { name: string }[]).map((item) => item.name); basis.dataDigest = digest(node.props.series); }
    if (node.type === "Table") { basis.columns = node.props.columns as string[]; basis.rowsDigest = digest(node.props.rows); }
    if (node.type === "FollowUpForm") basis.fields = Object.fromEntries((node.props.fields as { name: string; kind: string }[]).map((item) => [item.name, item.kind]));
    return [key, basis];
  }));
}
function validateNodes(nodes: AnswerViewState["nodes"], spec: Spec) {
  for (const [key, state] of Object.entries(nodes)) {
    const node = own(spec.elements, key);
    if (!node || node.type !== state.type) throw invalid();
    if (state.type === "Chart") {
      const count = (node.props.x_data as unknown[]).length;
      const seriesCount = (node.props.series as unknown[]).length;
      if (state.hidden.some((index) => index >= seriesCount) || state.hidden.length >= seriesCount || state.range[0] > state.range[1] || state.range[1] >= count || state.point >= count) throw invalid();
      if (state.table && ((state.table.sort && state.table.sort.column >= 1 + seriesCount - state.hidden.length) || state.table.selection.some((index) => index >= state.range[1] - state.range[0] + 1))) throw invalid();
    }
    if (state.type === "Table" && ((state.sort && state.sort.column >= (node.props.columns as unknown[]).length) || state.selection.some((index) => index >= (node.props.rows as unknown[]).length))) throw invalid();
    if (state.type === "FollowUpForm") {
      const names = new Set((node.props.fields as { name: string }[]).map((field) => field.name));
      if (Object.keys(state.values).some((name) => !names.has(name))) throw invalid();
    }
  }
}
function migrate(entry: Entry | undefined, spec: Spec, identity: { message_id: string; block_id: string }): AnswerViewState {
  const currentBasis = basisFor(spec);
  const nodes: AnswerViewState["nodes"] = {};
  for (const [key, value] of Object.entries(entry?.state.nodes ?? {})) {
    const before = entry && own(entry.basis, key);
    const after = own(currentBasis, key);
    if (!before || !after || before.type !== after.type || value.type !== after.type) continue;
    const state = structuredClone(value);
    if (state.type === "Table") {
      if (!same(before.columns, after.columns)) state.sort = null;
      if (before.rowsDigest !== after.rowsDigest || !same(before.columns, after.columns)) state.selection = [];
    }
    if (state.type === "Chart") {
      if (!same(before.series, after.series)) state.hidden = [];
      if (!same(before.labels, after.labels) || before.dataDigest !== after.dataDigest) {
        const labels = after.labels!;
        const old = before.labels!;
        const relocate = (index: number) => old.filter((label) => label === old[index]).length === 1 && labels.filter((label) => label === old[index]).length === 1 ? labels.indexOf(old[index]) : -1;
        const start = relocate(state.range[0]); const end = relocate(state.range[1]); const point = relocate(state.point);
        if (start < 0 || end < start || point < 0) {
          state.point = 0; state.range = [0, labels.length - 1]; state.inspecting = false;
        } else { state.point = point; state.range = [start, end]; }
      }
    }
    if (state.type === "Chart" && state.table && value.type === "Chart") {
      const columns = (basis: Basis, hidden: number[]) => basis.series!.filter((_, index) => !hidden.includes(index));
      if (!same(columns(before, value.hidden), columns(after, state.hidden))) state.table.sort = null;
      if (before.dataDigest !== after.dataDigest || !same(before.labels, after.labels) || !same(value.range, state.range) || !same(value.hidden, state.hidden)) state.table.selection = [];
    }
    if (state.type === "FollowUpForm") state.values = Object.fromEntries(Object.entries(state.values).filter(([name]) => own(before.fields!, name) && own(before.fields!, name) === own(after.fields!, name)));
    Object.defineProperty(nodes, key, { value: state, enumerable: true, configurable: true, writable: true });
  }
  const kinds = Object.values(spec.elements).filter((node) => node.type === "Evidence").map((node) => node.props.kind);
  const filter = entry?.state.block.evidenceFilter ?? "all";
  return { version: "answer-view-state@1", ...identity, revision: entry?.state.revision ?? 0, schema_digest: digest(spec), nodes,
    block: { evidenceFilter: filter === "all" || kinds.includes(filter) ? filter : "all" }, updated_at: entry?.state.updated_at ?? null };
}

export class AnswerViewStateStore {
  constructor(private readonly workspaces: SessionWorkspaceStore) {}
  async get(context: SessionWorkspaceContext, identity: { message_id: string; block_id: string }): Promise<AnswerViewState> {
    const parsed = answerViewStateIdentitySchema.safeParse(identity);
    if (!parsed.success) throw invalid();
    return this.workspaces.updateRecord(context, "answer-view-state.json", async ({ messages, current }) => {
      const spec = currentSpec(messages, identity.message_id, identity.block_id);
      const file = current as RecordFile | null;
      return { result: migrate(file ? own(file.entries, JSON.stringify([identity.message_id, identity.block_id])) : undefined, spec, identity) };
    });
  }
  async put(context: SessionWorkspaceContext, value: unknown): Promise<AnswerViewState> {
    if (bytes(value) > 32768) throw new AnswerViewStateError(413, "ANSWER_VIEW_STATE_TOO_LARGE", "answer view state exceeds size limit");
    const parsed = answerViewStatePutSchema.safeParse(value);
    if (!parsed.success) throw invalid();
    const input = parsed.data;
    return this.workspaces.updateRecord(context, "answer-view-state.json", async ({ messages, current }) => {
      const spec = currentSpec(messages, input.message_id, input.block_id);
      const file: RecordFile = current as RecordFile | null ?? { version: "answer-view-state@1", entries: {}, requestOrder: [], requests: {} };
      const key = JSON.stringify([input.message_id, input.block_id]);
      const requestDigest = digest(input);
      const request = own(file.requests, input.request_id);
      if (request) {
        if (request.digest !== requestDigest) throw new AnswerViewStateError(409, "ANSWER_VIEW_REQUEST_CONFLICT", "request ID was already used");
        return { result: request.response };
      }
      const state = migrate(own(file.entries, key), spec, { message_id: input.message_id, block_id: input.block_id });
      if (input.schema_digest !== state.schema_digest) throw new AnswerViewStateError(409, "ANSWER_VIEW_SCHEMA_CHANGED", "answer UI schema changed", state);
      if (input.expected_revision !== state.revision) throw new AnswerViewStateError(409, "ANSWER_VIEW_STATE_CONFLICT", "answer view revision changed", state);
      validateNodes(input.nodes, spec);
      const kinds = Object.values(spec.elements).filter((node) => node.type === "Evidence").map((node) => node.props.kind);
      if (input.block.evidenceFilter !== "all" && !kinds.includes(input.block.evidenceFilter)) throw invalid();
      const response: AnswerViewState = { ...state, nodes: input.nodes, block: input.block, revision: state.revision + 1, updated_at: new Date().toISOString() };
      if (bytes(response) > 32768) throw new AnswerViewStateError(413, "ANSWER_VIEW_STATE_TOO_LARGE", "answer view state exceeds size limit");
      Object.defineProperty(file.entries, key, { value: { state: response, basis: basisFor(spec) }, enumerable: true, configurable: true, writable: true });
      Object.defineProperty(file.requests, input.request_id, { value: { digest: requestDigest, response }, enumerable: true, configurable: true, writable: true });
      // Bounded replay window: older requests still carry stale revisions and cannot reapply.
      file.requestOrder ??= Object.keys(file.requests).filter((key) => key !== input.request_id);
      file.requestOrder.push(input.request_id);
      while (file.requestOrder.length > 32 || bytes(file.requests) > 131072) delete file.requests[file.requestOrder.shift()!];
      // Account for the exact pretty-printed bytes written by atomicWriteJson.
      if (Buffer.byteLength(JSON.stringify(file, null, 2), "utf8") > 1048576) throw new AnswerViewStateError(413, "ANSWER_VIEW_STATE_TOO_LARGE", "session view state exceeds size limit");
      return { result: response, value: file };
    });
  }
}