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

export const publicActivityPhaseSchema = z.enum([
  "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[];
}>;

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(),
}).strict();

const executionStepSchema = z.object({
  sequence: z.number().int().min(1).max(32),
  kind: z.enum(["skill", "tool", "validation"]),
  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(),
}).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();
const answerDeltaSchema = z.object({ type: z.literal("answer.delta"), text: z.string() }).strict();
const runCompletedSchema = z.object({ type: z.literal("run.completed"), receipt: agentExecutionReceiptSchema }).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.
const runFailedSchema = z.object({
  type: z.literal("run.failed"),
  code: z.enum(["runtime_contract_incomplete", "calculation_failed", "empty_answer", "cancelled"]),
  message: z.string().max(200),
  receipt: agentExecutionReceiptSchema.optional(),
}).strict();

export const consultationAgentPublicEventSchema = z.discriminatedUnion("type", [
  runStartedSchema, skillStartedSchema, skillCompletedSchema, toolStartedSchema, activitySchema,
  toolCompletedSchema, toolFailedSchema, answerDeltaSchema, runCompletedSchema, runFailedSchema,
]);
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),
  });
}