import { z } from "zod";
import { executionContextSchema } from "./context.js";
import { artifactRefSchema } from "./finding.js";
import {
  defineRecordCodec,
  recordEnvelopeSchema,
  type RecordCodec,
} from "./envelope.js";

export const scopeRefSchema = z
  .object({
    id: z.string(),
    detectorId: z.string().optional(),
    surfaceId: z.string().optional(),
  })
  .strict();

export const witnessRefSchema = z
  .object({
    id: z.string(),
    artifactRef: artifactRefSchema.optional(),
  })
  .strict();

export const capabilityProfileRefSchema = z
  .object({
    id: z.string(),
  })
  .strict();

export const coverageReasonSchema = z.enum([
  "empty-selection",
  "unproven-precondition",
  "inert-perturbation",
  "failed-known-bad",
  "unsupported-capability",
  "unreachable-target",
  "missing-artifact",
]);

export const COVERAGE_OUTCOME_VERSION = "coverage-outcome/v1" as const;

export type JsonObservation =
  | null
  | boolean
  | number
  | string
  | JsonObservation[]
  | { [key: string]: JsonObservation };

export const jsonObservationSchema: z.ZodType<JsonObservation> = z.lazy(() =>
  z.union([
    z.null(),
    z.boolean(),
    z.number(),
    z.string(),
    z.array(jsonObservationSchema),
    z.record(z.string(), jsonObservationSchema),
  ]),
);

export const interactionArtifactRefSchema = z
  .object({
    relativePath: z.string(),
    contentHash: z.string(),
    mediaType: z.string(),
    dimensions: z.object({ width: z.number(), height: z.number() }).strict().optional(),
  })
  .strict();

export type InteractionArtifactRef = z.infer<typeof interactionArtifactRefSchema>;

export const interactionCoverageWitnessRefSchema = z.object({ id: z.string() }).strict();

export const interactionControlRefSchema = z
  .object({ accessibleName: z.string(), locator: z.string() })
  .strict();

const preconditionEvidenceBase = {
  expected: jsonObservationSchema,
  observed: jsonObservationSchema,
  setup: jsonObservationSchema,
  locator: z.string(),
  witness: interactionCoverageWitnessRefSchema,
  artifact: interactionArtifactRefSchema,
} as const;

const B = preconditionEvidenceBase;

export const preconditionObligationEvidenceSchema = z.discriminatedUnion("unproven", [
  z.object({
    unproven: z.literal("missing-expected"),
    observed: B.observed,
    locator: B.locator,
    setup: B.setup.optional(),
    witness: B.witness.optional(),
    artifact: B.artifact.optional(),
  }).strict(),
  z.object({
    unproven: z.literal("missing-observed"),
    expected: B.expected,
    locator: B.locator,
    observed: B.observed.optional(),
    setup: B.setup.optional(),
    witness: B.witness.optional(),
    artifact: B.artifact.optional(),
  }).strict(),
  z.object({
    unproven: z.literal("missing-setup"),
    expected: B.expected,
    locator: B.locator,
    observed: B.observed.optional(),
    setup: B.setup.optional(),
    witness: B.witness.optional(),
    artifact: B.artifact.optional(),
  }).strict(),
  z.object({
    unproven: z.literal("missing-locator"),
    expected: B.expected,
    observed: B.observed,
    setup: B.setup,
    witness: B.witness,
    artifact: B.artifact.optional(),
  }).strict(),
  z.object({
    unproven: z.literal("missing-witness"),
    expected: B.expected,
    locator: B.locator,
    observed: B.observed.optional(),
    setup: B.setup.optional(),
    artifact: B.artifact.optional(),
  }).strict(),
  z.object({
    unproven: z.literal("missing-artifact"),
    expected: B.expected,
    locator: B.locator,
    observed: B.observed,
    setup: B.setup,
    witness: B.witness,
  }).strict(),
  z.object({
    unproven: z.literal("missing-expected-and-observed"),
    locator: B.locator,
    setup: B.setup.optional(),
    witness: B.witness.optional(),
    artifact: B.artifact.optional(),
  }).strict(),
  z.object({
    unproven: z.literal("missing-setup-and-locator"),
    expected: B.expected,
    witness: B.witness,
    artifact: B.artifact.optional(),
  }).strict(),
  z.object({
    unproven: z.literal("missing-locator-and-witness"),
    expected: B.expected,
    observed: B.observed,
    setup: B.setup,
    artifact: B.artifact.optional(),
  }).strict(),
  z.object({
    unproven: z.literal("missing-witness-and-artifact"),
    expected: B.expected,
    locator: B.locator,
    observed: B.observed,
    setup: B.setup,
  }).strict(),
  z.object({
    unproven: z.literal("unrecognized-precondition-gap"),
    failedPrecondition: z.string(),
    observed: B.observed,
    locator: B.locator.optional(),
    setup: B.setup.optional(),
    witness: B.witness.optional(),
    artifact: B.artifact.optional(),
  }).strict(),
]);

export const interactionCoverageObligationSchema = z
  .object({
    id: z.string(),
    witnessRefs: z.array(interactionCoverageWitnessRefSchema).nonempty(),
  })
  .strict();

export const preconditionCoverageObligationSchema = interactionCoverageObligationSchema
  .extend({ unproven: preconditionObligationEvidenceSchema })
  .strict();

export const activationInteractionCoverageKindSchema = z.enum([
  "unknown-command",
  "unsupported-activation",
]);

export const terminalInteractionCoverageKindSchema = z.enum([
  "unknown-terminal-state",
  "unknown-idempotency",
]);

export const terminalInteractionCoverageEvidenceSchema = z
  .object({
    kind: terminalInteractionCoverageKindSchema,
    phase: z.literal("terminal"),
    obligations: z.array(interactionCoverageObligationSchema).nonempty(),
  })
  .strict();

export const activationInteractionCoverageEvidenceSchema = z
  .object({
    kind: activationInteractionCoverageKindSchema,
    phase: z.literal("activation"),
    obligations: z.array(interactionCoverageObligationSchema).nonempty(),
  })
  .strict();

export const preconditionInteractionCoverageEvidenceSchema = z
  .object({
    kind: z.literal("unproven-precondition"),
    phase: z.literal("precondition"),
    obligations: z.array(preconditionCoverageObligationSchema).nonempty(),
  })
  .strict();

export const interactionCoverageSchema = z.discriminatedUnion("phase", [
  z
    .object({
      gapKind: activationInteractionCoverageKindSchema,
      phase: z.literal("activation"),
      evidence: activationInteractionCoverageEvidenceSchema,
      control: interactionControlRefSchema,
    })
    .strict()
    .refine((row) => row.evidence.kind === row.gapKind, {
      message: "interaction coverage gapKind must match evidence kind",
    }),
  z
    .object({
      gapKind: terminalInteractionCoverageKindSchema,
      phase: z.literal("terminal"),
      evidence: terminalInteractionCoverageEvidenceSchema,
      control: interactionControlRefSchema,
    })
    .strict()
    .refine((row) => row.evidence.kind === row.gapKind, {
      message: "interaction coverage gapKind must match evidence kind",
    }),
  z
    .object({
      gapKind: z.literal("unproven-precondition"),
      phase: z.literal("precondition"),
      evidence: preconditionInteractionCoverageEvidenceSchema,
      control: interactionControlRefSchema.optional(),
    })
    .strict(),
]);

export type InteractionCoverageWitnessRef = z.infer<typeof interactionCoverageWitnessRefSchema>;

export type InteractionCoverageObligation = z.infer<typeof interactionCoverageObligationSchema>;

export type PreconditionObligationEvidence = z.infer<typeof preconditionObligationEvidenceSchema>;

export type PreconditionCoverageObligation = z.infer<typeof preconditionCoverageObligationSchema>;

export type ActivationInteractionCoverageKind = z.infer<
  typeof activationInteractionCoverageKindSchema
>;

export type TerminalInteractionCoverageKind = z.infer<
  typeof terminalInteractionCoverageKindSchema
>;

export type TerminalInteractionCoverageEvidence = z.infer<
  typeof terminalInteractionCoverageEvidenceSchema
>;

export type ActivationInteractionCoverageEvidence = z.infer<
  typeof activationInteractionCoverageEvidenceSchema
>;

export type PreconditionInteractionCoverageEvidence = z.infer<
  typeof preconditionInteractionCoverageEvidenceSchema
>;

export type InteractionCoverageSerialization = z.infer<typeof interactionCoverageSchema>;

export const coverageOutcomeSchema = recordEnvelopeSchema
  .extend({
    schemaVersion: z.literal(COVERAGE_OUTCOME_VERSION),
    status: z.literal("coverage-incomplete"),
    scope: scopeRefSchema,
    context: executionContextSchema,
    reason: coverageReasonSchema,
    witnessRefs: z.array(witnessRefSchema),
    capabilityProfileRef: capabilityProfileRefSchema.optional(),
    interaction: interactionCoverageSchema.optional(),
    runId: z.string(),
  })
  .strict();

export type CoverageOutcome = z.infer<typeof coverageOutcomeSchema>;

export const coverageOutcomeCodec: RecordCodec<CoverageOutcome> =
  defineRecordCodec({
    schema: coverageOutcomeSchema,
    currentVersion: COVERAGE_OUTCOME_VERSION,
  });
