import { z } from "zod";

// ──────────────────────────────────────────────
// Shared enums (v0.1 & v0.2)
// ──────────────────────────────────────────────

export const confidenceEnum = z.enum(["low", "medium", "high"]);
const importanceEnum = z.enum([
  "incidental",
  "supporting",
  "important",
  "critical",
]);
const expectedInfoValueEnum = z.enum(["low", "medium", "high"]);

// ──────────────────────────────────────────────
// v0.1 — extraction-only schema (preserved)
// ──────────────────────────────────────────────

const confidenceEnumV1 = z.enum(["low", "medium", "high"]);

const itemSchemaV1 = z.object({
  id: z.string().min(1),
  description: z.string().min(1),
  confidence: confidenceEnumV1,
});

export const reconstructionSchema = z.object({
  observations: z.array(itemSchemaV1),
  reportedClaims: z.array(
    itemSchemaV1.extend({
      attributedTo: z
        .union([z.string().min(1), z.null()])
        .optional()
        .nullable(),
    }),
  ),
  assumptions: z.array(itemSchemaV1),
  entities: z.array(itemSchemaV1),
  transitions: z.array(
    itemSchemaV1.extend({
      entity: z.string().min(1),
      previousState: z.string().min(1),
      currentState: z.string().min(1),
      explanationStatus: z.string().min(1),
    }),
  ),
  expectedButMissing: z.array(itemSchemaV1),
  presentButUnexpected: z.array(itemSchemaV1),
  contradictions: z.array(itemSchemaV1),
  openUncertainties: z.array(itemSchemaV1),
});

// v0.1 analyse response (used internally)
export const analyseResponseSchema = z.object({
  reconstruction: z.union([reconstructionSchema, z.null()]),
  modelName: z.string(),
  responseDurationMs: z.number(),
  validationStatus: z.enum(["valid", "partial", "invalid"]),
  rawResponse: z.string().optional(),
  errors: z.array(z.string()).optional(),
});

export const healthResponseSchema = z.object({
  configPresent: z.boolean(),
  baseUrl: z.string().nullable(),
  model: z.string().nullable(),
  reachable: z.boolean(),
  error: z.string().nullable(),
});

// ──────────────────────────────────────────────
// v0.2 — reasoning classification + reconstruction
// ──────────────────────────────────────────────

export const inputTypes =
  /** @type {z.ZodType<typeof import("@/lib/reconstruction/schema").INPUT_TYPE_VALUE>} */ (
    z.enum([
      "observed_problem",
      "unexplained_change",
      "contradiction",
      "decision_request",
      "causal_claim",
      "reported_claim",
      "fault_report",
      "ambiguous_statement",
      "question",
      "desired_outcome",
      "insufficient_context",
      "other",
    ])
  );

export const reasoningModes =
  /** @type {z.ZodType<typeof import("@/lib/reconstruction/schema").REASONING_MODE_VALUE>} */ (
    z.enum([
      "establish_baseline",
      "identify_difference",
      "reconstruct_transition",
      "decompose_aggregate",
      "validate_measurement",
      "validate_claim",
      "investigate_contradiction",
      "clarify_meaning",
      "decision_support",
      "fault_investigation",
      "identify_missing_information",
      "test_possible_explanations",
      "other",
    ])
  );

const evidenceRecordSchema = z.object({
  id: z.string().min(1),
  description: z.string().min(1),
  evidenceType: z.enum([
    "direct_observation",
    "reported_statement",
    "interpretation",
    "assumption",
    "inferred_relationship",
  ]),
  source: z.string().optional(),
  attribution: z.string().nullable().optional(),
  confidence: confidenceEnum,
  importance: importanceEnum,
});

const reconstructionSchemaV2 = z.object({
  summary: z.string().min(1),
  actors: z.array(itemSchemaV1),
  systemsOrObjects: z.array(itemSchemaV1),
  expectedStates: z.array(itemSchemaV1),
  observedStates: z.array(itemSchemaV1),
  differences: z.array(itemSchemaV1),
  knownTransitions: z.array(
    itemSchemaV1.extend({
      entity: z.string().min(1),
      previousState: z.string().min(1),
      currentState: z.string().min(1),
      explanationStatus: z.string().min(1),
    }),
  ),
  unexplainedTransitions: z.array(
    itemSchemaV1.extend({
      entity: z.string().min(1).optional(),
      previousState: z.string().min(1).optional(),
      currentState: z.string().min(1).optional(),
    }),
  ),
  contradictions: z.array(itemSchemaV1),
  importantUnknowns: z.array(itemSchemaV1),
  plausibleInterpretations: z.array(
    z.object({
      id: z.string().min(1),
      description: z.string().min(1),
      supportingEvidenceIds: z.array(z.string()),
      assumptionsRequired: z.array(z.string()).optional().default([]),
      confidence: confidenceEnum,
    }),
  ),
});

const inputClassificationSchema = z.object({
  primaryType: inputTypes,
  secondaryTypes: z.array(inputTypes).optional().default([]),
  reasoningModes: z.array(reasoningModes).optional().default([]),
  classificationReason: z.string().min(1),
  confidence: confidenceEnum,
});

const nextQuestionSchema = z.object({
  id: z.string().min(1),
  question: z.string().min(1),
  targets: z.array(z.string()),
  reason: z.string().min(1),
  expectedInformationValue: expectedInfoValueEnum,
  reasoningMode: reasoningModes.optional().default("other"),
});

// v0.2 complete analysis response (what the model produces)
export const reconstructionV2Schema = z.object({
  inputClassification: inputClassificationSchema,
  reconstruction: reconstructionSchemaV2,
  evidence: z.array(evidenceRecordSchema),
  nextQuestion: nextQuestionSchema,
});

// Outer wrapper for API return (includes diagnostics + v0.2 data)
export const analyseResponseV2Schema = z.object({
  inputClassification: inputClassificationSchema.optional(),
  reconstruction: reconstructionSchemaV2.optional().nullable(),
  evidence: z.array(evidenceRecordSchema).optional(),
  nextQuestion: nextQuestionSchema.optional(),
  modelName: z.string(),
  responseDurationMs: z.number(),
  validationStatus: z.enum(["valid", "partial", "invalid"]),
  rawResponse: z.string().optional(),
  errors: z.array(z.string()).optional(),
  promptVersion: z.string().optional(),
});

// ──────────────────────────────────────────────
// Parsing helpers
// ──────────────────────────────────────────────

export function parseReconstruction(raw) {
  if (typeof raw === "string") {
    try {
      raw = JSON.parse(raw);
    } catch {
      throw new SyntaxError("Model response is not valid JSON");
    }
  }
  return reconstructionSchema.parse(raw);
}

export function parseReconstructionV2(raw) {
  if (typeof raw === "string") {
    try {
      raw = JSON.parse(raw);
    } catch {
      throw new SyntaxError("Model response is not valid JSON");
    }
  }
  return reconstructionV2Schema.parse(raw);
}