import { canonicalJson, sha256, type JsonObject } from "../core/json.js";
import type { JazzThoughtStore } from "../jazz/store.js";
import type { ThoughtAgentDeclaration } from "./types.js";
import {
  contextPacketFromSnapshot,
  snapshotManifestMatchesRunContext,
} from "./context.js";
import { declarationFingerprint } from "./declarations.js";
import {
  createOutputContractRegistry,
  outputContractForDeclaration,
} from "./output-contracts.js";
import { composePiModelInput } from "./pi.js";

export interface RunContextInspection {
  run: {
    id: string;
    agentId: string;
    agentVersion: number;
    triggerEventId: string;
    status: string;
  };
  snapshot: {
    id: string;
    source: string;
    path: string;
    contentType: string;
    sha256: string;
    sizeBytes: number;
    createdAt: string;
  };
  packet: {
    contentIncluded: boolean;
    systemChars: number;
    currentTextChars: number;
    messageCount: number;
    messageRoles: string[];
    messagesChars: number;
    imageArtifacts: number;
    manifest: JsonObject;
    systemTextSha256?: string;
    currentTextSha256: string;
    messagesSha256?: string;
    systemText?: string;
    currentText?: string;
    messages?: JsonObject[];
    imageArtifactReferences?: JsonObject[];
  };
  modelInput: {
    exact: boolean;
    incompleteReasons: string[];
    systemPromptChars?: number;
    currentPromptChars?: number;
    systemPromptSha256?: string;
    currentPromptSha256?: string;
    priorMessageCount: number;
    priorMessageRoles: string[];
    proposalCapabilities?: JsonObject;
    systemPrompt?: string;
    currentPrompt?: string;
    priorMessages?: JsonObject[];
  };
}

export async function inspectRunContext(
  store: JazzThoughtStore,
  runId: string,
  options: { includeContent?: boolean } = {},
): Promise<RunContextInspection> {
  const run = await store.getRun(runId);
  if (!run) throw new Error(`Run not found: ${runId}`);
  const snapshotIdentity = objectField(run.contextManifest.contextSnapshot);
  const snapshotId = stringField(snapshotIdentity?.id);
  if (!snapshotId) throw new Error(`Run has no durable context snapshot: ${runId}`);
  const snapshot = await store.getDocumentVersion(snapshotId);
  if (!snapshot
    || snapshot.id !== snapshotId
    || snapshot.documentId !== snapshotId
    || snapshot.source !== `context:${run.agentId}`
    || sha256(snapshot.content) !== snapshot.sha256
    || Buffer.byteLength(snapshot.content) !== snapshot.sizeBytes) {
    throw new Error(`Run context snapshot storage evidence is missing or inconsistent: ${runId}`);
  }
  const packet = contextPacketFromSnapshot(snapshot.content, snapshotId);
  if (!snapshotManifestMatchesRunContext(packet.manifest, run.contextManifest)) {
    throw new Error(`Run context snapshot does not match the run manifest: ${runId}`);
  }
  const messages = packet.messages as unknown as JsonObject[] | undefined;
  const imageArtifacts = packet.imageArtifacts as unknown as JsonObject[] | undefined;
  const messagesJson = messages ? canonicalJson(messages) : undefined;
  const includeContent = options.includeContent === true;
  const modelInput = await reconstructModelInput(store, run, packet, includeContent);
  return {
    run: {
      id: run.id,
      agentId: run.agentId,
      agentVersion: run.agentVersion,
      triggerEventId: run.triggerEventId,
      status: run.status,
    },
    snapshot: {
      id: snapshot.id,
      source: snapshot.source,
      path: snapshot.path,
      contentType: snapshot.contentType,
      sha256: snapshot.sha256,
      sizeBytes: snapshot.sizeBytes,
      createdAt: snapshot.createdAt,
    },
    packet: {
      contentIncluded: includeContent,
      systemChars: packet.systemText?.length ?? 0,
      currentTextChars: packet.text.length,
      messageCount: packet.messages?.length ?? 0,
      messageRoles: packet.messages?.map((message) => message.role) ?? [],
      messagesChars: messagesJson?.length ?? 0,
      imageArtifacts: packet.imageArtifacts?.length ?? 0,
      manifest: packet.manifest,
      ...(packet.systemText ? { systemTextSha256: sha256(packet.systemText) } : {}),
      currentTextSha256: sha256(packet.text),
      ...(messagesJson ? { messagesSha256: sha256(messagesJson) } : {}),
      ...(includeContent ? {
        ...(packet.systemText ? { systemText: packet.systemText } : {}),
        currentText: packet.text,
        ...(messages ? { messages } : {}),
        ...(imageArtifacts ? { imageArtifactReferences: imageArtifacts } : {}),
      } : {}),
    },
    modelInput,
  };
}

async function reconstructModelInput(
  store: JazzThoughtStore,
  run: Awaited<ReturnType<JazzThoughtStore["getRun"]>> & {},
  packet: ReturnType<typeof contextPacketFromSnapshot>,
  includeContent: boolean,
): Promise<RunContextInspection["modelInput"]> {
  const reasons: string[] = [];
  const stored = (await store.listAgents()).find((agent) => agent.id === run.agentId && agent.version === run.agentVersion);
  if (!stored) reasons.push("matching stored declaration is unavailable");
  const declaration = stored?.spec as unknown as ThoughtAgentDeclaration | undefined;
  if (stored && stored.specHash !== sha256(canonicalJson(stored.spec))) reasons.push("stored declaration hash is inconsistent");
  if (declaration) {
    if (declaration.id !== run.agentId || declaration.version !== run.agentVersion) reasons.push("stored declaration identity is inconsistent");
    if (run.promptHash !== sha256(declaration.systemPrompt)) reasons.push("stored declaration prompt does not match run prompt hash");
    if (run.contextManifest.declarationFingerprint !== declarationFingerprint(declaration)) {
      reasons.push("stored declaration does not match run fingerprint");
    }
    if (declaration.mode !== "pi") reasons.push("run is not a Pi model input");
    if ((declaration.tools?.length ?? 0) > 0) reasons.push("read-only prefetched evidence content is not durable");
    if ((packet.imageArtifacts?.length ?? 0) > 0) reasons.push("resolved current-image bytes are not durable in the context snapshot");
  }
  const traces = await store.listTrace(run.id);
  const systemTrace = promptTraceMetadata(traces, "system_prompt");
  const promptTrace = promptTraceMetadata(traces, "prompt");
  if (!systemTrace || !promptTrace) reasons.push("durable prompt hash traces are unavailable or ambiguous");
  let composed: ReturnType<typeof composePiModelInput> | undefined;
  if (declaration && reasons.length === 0) {
    try {
      const contract = createOutputContractRegistry().resolve(outputContractForDeclaration(declaration));
      composed = composePiModelInput(declaration, packet, {
        readOnlyToolCount: 0,
        readOnlyEvidenceText: "",
        resolvedContextImageCount: 0,
        outputContractPrompt: contract.prompt,
      });
      if (systemTrace!.chars !== composed.systemPrompt.length || systemTrace!.sha256 !== sha256(composed.systemPrompt)) {
        reasons.push("reconstructed system prompt does not match durable trace evidence");
      }
      if (promptTrace!.chars !== composed.prompt.length || promptTrace!.sha256 !== sha256(composed.prompt)) {
        reasons.push("reconstructed current prompt does not match durable trace evidence");
      }
    } catch {
      reasons.push("stored declaration cannot reconstruct the run model input");
    }
  }
  const exact = reasons.length === 0 && composed !== undefined;
  return {
    exact,
    incompleteReasons: reasons,
    ...(systemTrace ? { systemPromptChars: systemTrace.chars, systemPromptSha256: systemTrace.sha256 } : {}),
    ...(promptTrace ? { currentPromptChars: promptTrace.chars, currentPromptSha256: promptTrace.sha256 } : {}),
    priorMessageCount: packet.messages?.length ?? 0,
    priorMessageRoles: packet.messages?.map((message) => message.role) ?? [],
    ...(composed?.proposalCapabilities ? {
      proposalCapabilities: composed.proposalCapabilities as unknown as JsonObject,
    } : {}),
    ...(includeContent && exact ? {
      systemPrompt: composed!.systemPrompt,
      currentPrompt: composed!.prompt,
      ...(packet.messages ? { priorMessages: packet.messages as unknown as JsonObject[] } : {}),
    } : {}),
  };
}

function promptTraceMetadata(
  traces: Awaited<ReturnType<JazzThoughtStore["listTrace"]>>,
  type: "system_prompt" | "prompt",
): { chars: number; sha256: string } | undefined {
  const matches = traces.filter((trace) => trace.type === type);
  if (matches.length !== 1) return undefined;
  const data = objectField(matches[0]!.payload.data);
  return Number.isSafeInteger(data?.chars)
    && typeof data?.sha256 === "string"
    && /^[a-f0-9]{64}$/.test(data.sha256)
    ? { chars: Number(data.chars), sha256: data.sha256 }
    : undefined;
}

function objectField(value: unknown): JsonObject | undefined {
  return value && typeof value === "object" && !Array.isArray(value) ? value as JsonObject : undefined;
}

function stringField(value: unknown): string | undefined {
  return typeof value === "string" && value.length > 0 ? value : undefined;
}