import { z } from "zod";
import { consultationDomainSchema, type ConsultationDomain } from "./consultation-domain-registry.ts";
import { publicThinkingSectionSchema } from "./consultation-thinking-plan.ts";

export const publicActivityPhaseSchema = z.enum([
  // The first byte of a consultation stream, before the turn is classified or
  // the chart is read (BUG-1074). The client already shows the same queued
  // row from the send frame, so this frame changes nothing on screen.
  "received",
  "loading-method",
  "chart-calculation",
  "evidence-validation",
  "answer-composition",
]);
export type PublicActivityPhase = z.infer<typeof publicActivityPhaseSchema>;

export type WorkflowReceipt = Readonly<{
  route: string;
  status: string;
  preciseTiming: string;
  missingLayers: readonly string[];
  domains?: readonly ConsultationDomain[];
  // Requested but not calculated, because the run's wall clock could not pay
  // for them. Present so a partial plan cannot be read as a complete one.
  omittedDomains?: readonly ConsultationDomain[];
  minuteSensitiveThemes?: readonly string[];
  planOverrideIgnored?: boolean;
}>;

export const workflowReceiptSchema: z.ZodType<WorkflowReceipt> = z.object({
  route: z.string().max(120),
  status: z.string().max(120),
  preciseTiming: z.string().max(120),
  missingLayers: z.array(z.string().max(120)).max(30),
  domains: z.array(consultationDomainSchema).min(1).max(6).optional(),
  omittedDomains: z.array(consultationDomainSchema).min(1).max(6).optional(),
  minuteSensitiveThemes: z.array(z.string().max(40)).max(24).optional(),
  planOverrideIgnored: z.boolean().optional(),
}).strict();

export const techniqueAuditStatusSchema = z.enum(["executed", "blocked", "not_applicable"]);
export type TechniqueAuditStatus = z.infer<typeof techniqueAuditStatusSchema>;

export const techniqueAuditRowSchema = z.object({
  technique: z.string().trim().min(1).max(160),
  status: techniqueAuditStatusSchema,
  note: z.string().max(320).optional(),
}).strict();
export type TechniqueAuditRow = z.infer<typeof techniqueAuditRowSchema>;


const executionStepSchema = z.object({
  sequence: z.number().int().min(1).max(32),
  kind: z.enum(["skill", "tool", "validation", "abort"]),
  name: z.string().max(120),
  status: z.enum(["completed", "failed"]),
  durationMs: z.number().int().min(0).optional(),
}).strict();

const stepBudgetSchema = z.object({
  planned: z.number().int().min(1).max(32),
  used: z.number().int().min(0).max(32),
  remaining: z.number().int().min(0).max(32),
  truncated: z.boolean(),
}).strict();

export const agentExecutionReceiptSchema = z.object({
  runId: z.string().min(1).max(120),
  runtime: z.literal("mastra-agentic"),
  skill: z.object({
    name: z.literal("jyotish-vedic-astrology"),
    loaded: z.boolean(),
    version: z.string().max(120).optional(),
    // How many reference documents the model opened after loading the skill. Required rather than
    // optional: the count was tracked in runtime state and surfaced nowhere, so "did the model
    // consult the method at all" was unanswerable from a finished run. Zero is a real answer.
    referenceReads: z.number().int().min(0).max(64),
    // How many strict-method sections the server delivered with the evidence. Reported beside
    // referenceReads rather than folded into it, because "the model went looking" and "the method
    // was in front of it" are different facts and only one of them is under the model's control.
    methodologySections: z.number().int().min(0).max(12),
  }).strict(),
  steps: z.array(executionStepSchema).max(32),
  stepBudget: stepBudgetSchema.optional(),
  workflow: workflowReceiptSchema,
  techniqueTruth: z.string().max(120).optional(),
  techniqueAuditTable: z.array(techniqueAuditRowSchema).max(80).optional(),
}).strict();
export type AgentExecutionReceipt = z.infer<typeof agentExecutionReceiptSchema>;

const runStartedSchema = z.object({ type: z.literal("run.started"), runId: z.string(), requestId: z.string() }).strict();
const skillStartedSchema = z.object({ type: z.literal("skill.started"), name: z.literal("jyotish-vedic-astrology") }).strict();
const skillCompletedSchema = z.object({ type: z.literal("skill.completed"), name: z.literal("jyotish-vedic-astrology") }).strict();
const toolStartedSchema = z.object({ type: z.literal("tool.started"), callId: z.string(), tool: z.literal("run-jyotish-consultation"), label: z.string() }).strict();
const activitySchema = z.object({ type: z.literal("activity"), phase: publicActivityPhaseSchema, label: z.string().max(120) }).strict();
const toolCompletedSchema = z.object({
  type: z.literal("tool.completed"), callId: z.string(), tool: z.literal("run-jyotish-consultation"),
  status: z.enum(["ready", "degraded", "blocked"]), durationMs: z.number().int().min(0),
}).strict();
const toolFailedSchema = z.object({
  type: z.literal("tool.failed"), callId: z.string(), tool: z.literal("run-jyotish-consultation"),
  code: z.enum(["calculation_failed", "timeout", "cancelled"]),
}).strict();
/** `replace: true` swaps the whole answer for `text` (a moderated answer, 2026-09-30). */
const answerDeltaSchema = z.object({ type: z.literal("answer.delta"), text: z.string(), replace: z.literal(true).optional() }).strict();
const thinkingDeltaSchema = z.object({ type: z.literal("thinking.delta"), text: z.string() }).strict();
const thinkingSectionEventSchema = publicThinkingSectionSchema.extend({
  type: z.literal("thinking.section"),
}).strict();
const phaseStartedSchema = z.object({
  type: z.literal("phase.started"),
  phase: z.string().min(1).max(40),
  label: z.string().max(120),
  progress: z.object({
    current: z.number().int().min(0).max(32),
    total: z.number().int().min(1).max(32),
  }).strict().optional(),
}).strict();
const phaseCompletedSchema = z.object({
  type: z.literal("phase.completed"),
  phase: z.string().min(1).max(40),
  durationMs: z.number().int().min(0),
}).strict();
const thinkPlanSchema = z.object({
  type: z.literal("think.plan"),
  steps: z.array(z.object({
    id: z.string().min(1).max(80),
    title: z.string().min(1).max(120),
  }).strict()).min(1).max(12),
}).strict();
const thinkStepSchema = z.object({
  type: z.literal("think.step"),
  id: z.string().min(1).max(80),
  status: z.enum(["running", "done"]),
  text: z.string().min(1).max(400).optional(),
}).strict();
const sessionTitleSchema = z.object({
  type: z.literal("session.title"),
  title: z.string().trim().min(1).max(48),
}).strict();
const runCompletedSchema = z.object({
  type: z.literal("run.completed"),
  receipt: agentExecutionReceiptSchema.optional(),
  responseKind: z.literal("smalltalk").optional(),
}).strict();
// A failure is the case the receipt is most needed for, so it carries the same
// allowlisted receipt a completed run does. It stays optional because the
// receipt is built from live state that a hard failure may leave unparseable,
// and losing the whole failure event would be worse than losing its receipt.
// `request_rejected` is a failure the route used to answer with a JSON body
// and an HTTP status before the stream existed (credits, plan, subject,
// session full…). Since the stream opens first (BUG-1074) it travels here
// instead: `status` keeps the HTTP semantics the client acts on (402 opens
// billing, 4xx means the reservation never committed), `reason` is the JSON
// body's `code` when it had one (session_full, pricing_configuration_unavailable).
const runFailedSchema = z.object({
  type: z.literal("run.failed"),
  code: z.enum(["runtime_contract_incomplete", "calculation_failed", "empty_answer", "answer_truncated", "cancelled", "request_rejected"]),
  message: z.string().max(200),
  status: z.number().int().min(400).max(599).optional(),
  reason: z.string().max(80).optional(),
  receipt: agentExecutionReceiptSchema.optional(),
}).strict();

export const consultationAgentPublicEventSchema = z.discriminatedUnion("type", [
  runStartedSchema, skillStartedSchema, skillCompletedSchema, toolStartedSchema, activitySchema,
  toolCompletedSchema, toolFailedSchema, answerDeltaSchema, thinkingDeltaSchema,
  thinkingSectionEventSchema, phaseStartedSchema, phaseCompletedSchema, thinkPlanSchema,
  thinkStepSchema, sessionTitleSchema, runCompletedSchema, runFailedSchema,
]).superRefine((event, ctx) => {
  if (event.type === "run.completed"
    && (event.responseKind === "smalltalk" ? event.receipt !== undefined : event.receipt === undefined)) {
    ctx.addIssue({ code: z.ZodIssueCode.custom, message: "completion_requires_receipt_or_smalltalk" });
  }
  if (event.type === "run.failed") {
    if ((event.status !== undefined || event.reason !== undefined) && event.code !== "request_rejected") {
      ctx.addIssue({ code: z.ZodIssueCode.custom, message: "status_and_reason_only_on_request_rejected" });
    }
    if (event.code === "request_rejected" && event.status === undefined) {
      ctx.addIssue({ code: z.ZodIssueCode.custom, message: "request_rejected_requires_status" });
    }
  }
});
export type ConsultationAgentPublicEvent = z.infer<typeof consultationAgentPublicEventSchema>;

export function createNdjsonParser(onEvent: (event: ConsultationAgentPublicEvent) => void) {
  let buffer = "";
  function consume(value: string, final: boolean) {
    buffer += value;
    const lines = buffer.split("\n");
    buffer = lines.pop() ?? "";
    for (const line of lines) {
      if (line.trim()) onEvent(consultationAgentPublicEventSchema.parse(JSON.parse(line)));
    }
    if (final && buffer.trim()) {
      onEvent(consultationAgentPublicEventSchema.parse(JSON.parse(buffer)));
      buffer = "";
    }
  }
  return Object.freeze({
    push: (value: string) => consume(value, false),
    finish: (value = "") => consume(value, true),
  });
}