import { z } from "zod";
import type { JsonObject } from "../core/json.js";
import { focusDeclarationSchema } from "../focuses/types.js";

export const PROPOSAL_DECLARATION_NAMES = ["memory-change", "self-correction", "focus-declaration"] as const;
export type ProposalDeclarationName = typeof PROPOSAL_DECLARATION_NAMES[number];

export const PROPOSAL_TOOL_NAMES = ["request_memory_change", "submit_correction", "propose_focus"] as const;
export type ProposalToolName = typeof PROPOSAL_TOOL_NAMES[number];

const idSchema = z.string().min(1).max(500);
const sha256Schema = z.string().regex(/^[a-f0-9]{64}$/);
const evidenceIdsSchema = z.array(idSchema).max(16).superRefine((ids, context) => {
  const seen = new Set<string>();
  for (let index = 0; index < ids.length; index += 1) {
    if (seen.has(ids[index]!)) context.addIssue({ code: "custom", path: [index], message: "Evidence ids must be unique" });
    seen.add(ids[index]!);
  }
});

export const memoryChangeArgumentsSchema = z.object({
  operation: z.enum(["append", "replace-document"]),
  proposed_text: z.string().min(1).max(32_768),
  reason: z.string().min(1).max(1_000),
  evidence_event_ids: evidenceIdsSchema,
}).strict();

export const CORRECTION_TARGET_LATEST = "latest";

export const correctionArgumentsSchema = z.object({
  target_output: z.union([idSchema, z.literal(CORRECTION_TARGET_LATEST)]),
  replacement: z.string().min(1).max(4_096),
  reason: z.string().min(1).max(1_000),
  evidence_event_ids: evidenceIdsSchema,
}).strict();

export const focusDeclarationArgumentsSchema = z.object({
  declaration: focusDeclarationSchema,
}).strict();

export const proposalMemoryTargetSchema = z.object({
  source: z.literal("filesystem:telegram-agent-context"),
  documentId: idSchema,
  path: z.literal("memory.md"),
  versionId: idSchema,
  sha256: sha256Schema,
  contentType: z.literal("text/markdown"),
}).strict();

export const outputContractIdentitySchema = z.object({
  id: idSchema,
  version: z.number().int().positive(),
  sha256: sha256Schema,
}).strict();

export const proposalCorrectionTargetSchema = z.object({
  runId: idSchema,
  outputEventId: idSchema,
  deliveryReceiptEventId: idSchema,
  sourceRootEventId: idSchema,
  outputContract: outputContractIdentitySchema,
}).strict();

export const proposalCapabilitiesSchema = z.object({
  enabled: z.array(z.enum(PROPOSAL_DECLARATION_NAMES)).max(3).superRefine((values, context) => {
    if (new Set(values).size !== values.length) context.addIssue({ code: "custom", message: "Proposal capability names must be unique" });
  }),
  evidenceEventIds: z.array(idSchema).max(300).superRefine((values, context) => {
    if (new Set(values).size !== values.length) context.addIssue({ code: "custom", message: "Proposal evidence ids must be unique" });
  }),
  memoryTarget: proposalMemoryTargetSchema.optional(),
  correctionTargets: z.array(proposalCorrectionTargetSchema).max(100).superRefine((targets, context) => {
    const ids = targets.map((target) => target.outputEventId);
    if (new Set(ids).size !== ids.length) context.addIssue({ code: "custom", message: "Correction target output ids must be unique" });
  }),
}).strict().superRefine((value, context) => {
  if (value.enabled.includes("memory-change") !== Boolean(value.memoryTarget)) {
    context.addIssue({ code: "custom", path: ["memoryTarget"], message: "Memory proposal capability requires exactly one memory target" });
  }
  if (!value.enabled.includes("self-correction") && value.correctionTargets.length > 0) {
    context.addIssue({ code: "custom", path: ["correctionTargets"], message: "Correction targets require self-correction capability" });
  }
});

export const capturedMemoryProposalSchema = z.object({
  toolCallId: idSchema,
  kind: z.literal("memory-change"),
  arguments: memoryChangeArgumentsSchema,
}).strict();

export const capturedCorrectionProposalSchema = z.object({
  toolCallId: idSchema,
  kind: z.literal("self-correction"),
  arguments: correctionArgumentsSchema,
}).strict();

export const capturedFocusDeclarationProposalSchema = z.object({
  toolCallId: idSchema,
  kind: z.literal("focus-declaration"),
  arguments: focusDeclarationArgumentsSchema,
}).strict();

export const capturedProposalSchema = z.discriminatedUnion("kind", [
  capturedMemoryProposalSchema,
  capturedCorrectionProposalSchema,
  capturedFocusDeclarationProposalSchema,
]);

export const capturedProposalsSchema = z.array(capturedProposalSchema).max(3).superRefine((proposals, context) => {
  const calls = proposals.map((proposal) => proposal.toolCallId);
  if (new Set(calls).size !== calls.length) context.addIssue({ code: "custom", message: "Proposal tool-call ids must be unique" });
  const kinds = proposals.map((proposal) => proposal.kind);
  if (new Set(kinds).size !== kinds.length) context.addIssue({ code: "custom", message: "Only one proposal of each kind is allowed" });
});

export type MemoryChangeArguments = z.infer<typeof memoryChangeArgumentsSchema>;
export type CorrectionArguments = z.infer<typeof correctionArgumentsSchema>;
export type FocusDeclarationArguments = z.infer<typeof focusDeclarationArgumentsSchema>;
export type ProposalMemoryTarget = z.infer<typeof proposalMemoryTargetSchema>;
export type ProposalCorrectionTarget = z.infer<typeof proposalCorrectionTargetSchema>;
export type ProposalCapabilities = z.infer<typeof proposalCapabilitiesSchema>;
export type CapturedProposal = z.infer<typeof capturedProposalSchema>;

export function validateCapturedProposalAgainstCapabilities(
  proposal: CapturedProposal,
  capabilities: ProposalCapabilities,
): CapturedProposal {
  const admittedEvidence = new Set(capabilities.evidenceEventIds);
  if (proposal.kind !== "focus-declaration") {
    for (const id of proposal.arguments.evidence_event_ids) {
      if (!admittedEvidence.has(id)) throw new Error("Proposal evidence id is outside the context snapshot");
    }
  }
  if (proposal.kind === "memory-change") {
    if (!capabilities.enabled.includes("memory-change") || !capabilities.memoryTarget) {
      throw new Error("Memory proposal capability is unavailable");
    }
    return proposal;
  }
  if (proposal.kind === "focus-declaration") {
    if (!capabilities.enabled.includes("focus-declaration")) {
      throw new Error("Focus proposal capability is unavailable");
    }
    return proposal;
  }
  if (!capabilities.enabled.includes("self-correction")) {
    throw new Error("Correction proposal capability is unavailable");
  }
  const resolvedTargetOutput = resolveCorrectionTarget(proposal.arguments.target_output, capabilities);
  if (!resolvedTargetOutput) {
    throw new Error("Correction target is outside the context snapshot");
  }
  return proposal;
}

export function resolveCorrectionTarget(
  targetOutput: string,
  capabilities: ProposalCapabilities,
): string | undefined {
  if (targetOutput === CORRECTION_TARGET_LATEST) {
    return capabilities.correctionTargets.at(-1)?.outputEventId;
  }
  return capabilities.correctionTargets.some((target) => target.outputEventId === targetOutput)
    ? targetOutput
    : undefined;
}

export function proposalCapabilitiesJson(value: ProposalCapabilities): JsonObject {
  return proposalCapabilitiesSchema.parse(value) as unknown as JsonObject;
}