import { z } from "zod";

/**
 * Declarative workflow schema — a Microsoft-agentic-workflow-style routed graph of typed
 * steps a runner walks unattended.
 *
 * This file is the STRICT grammar we accept — and it is a SECURITY boundary, not just a
 * convenience:
 *
 *   - `z.strictObject` everywhere: an unknown key is a HARD reject, so a malicious
 *     (or just malformed) file can't smuggle a field the runner half-understands.
 *   - `z.discriminatedUnion` on `type`: an unknown step type is a HARD reject — we
 *     never fall through to a permissive default. New step types are opt-in here.
 *   - Sub-workflow refs are LOCAL RELATIVE PATHS ONLY. A remote ref (registry/URL/GitHub)
 *     is rejected, because a remote ref is an unsigned code-fetch an untrusted source could
 *     redirect (= RCE). Composition stays within the user's own on-disk workflows dir.
 *
 * Jinja/template strings (prompts, `when:` conditions, `{{ ... }}` refs) are kept as
 * OPAQUE strings here — parsing/eval happens in the runner (expr.ts), not the schema. We
 * only validate structure and step types.
 */

// A step name: the node id in the route graph. Kept simple so it's safe to use as a
// map key and (for `agent` steps) an output directory segment.
const StepName = z
  .string()
  .min(1)
  .regex(/^[a-zA-Z0-9_-]{1,64}$/, "step names are 1-64 chars of [a-zA-Z0-9_-]");

// A route target: another step's name, or one of the terminal sentinels. `$end`
// finishes successfully; `$fail` finishes failed. (`$end` mirrors the common convention;
// `$fail` lets a route terminate as failed without a dedicated terminate step.)
const RouteTarget = z.union([StepName, z.literal("$end"), z.literal("$fail")]);

// One edge out of a step. Routes are evaluated in order; the first whose `when`
// (an opaque expr.ts condition) is truthy is taken. A `when`-less route is the
// unconditional fallthrough and MUST be last among a step's routes (enforced in the
// runner, not here, since it needs the whole list).
const Route = z.strictObject({
  to: RouteTarget,
  when: z.string().min(1).optional(),
});
const Routes = z.array(Route).default([]);

// A typed output field. We only carry the coarse type + a description — enough to
// drive strict JSON extraction from an `agent`/`script` step without a full JSON
// Schema. Minimal, reviewable, no surprises.
const OutputField = z.strictObject({
  type: z.enum(["string", "number", "boolean", "array", "object"]),
  description: z.string().optional(),
});
const OutputSchema = z.record(z.string(), OutputField);

// A tool ceiling entry: a builtin name ("read"), an exact MCP selector
// ("<server>__<tool>"), or a whole-server wildcard ("<server>__*"). The runner
// resolves these against its safe default when `tools` is omitted.
const ToolEntry = z.string().min(1);
const Tools = z.array(ToolEntry);

// ---- Step types (the discriminated union) --------------------------------------

// `agent` — an LLM turn run as a headless session under the host's auto-approve gate.
// `tools` is the HARD ceiling; omitted → the safe read/web set. This is the only
// step that reasons; every other type is deterministic plumbing.
const AgentStep = z.strictObject({
  name: StepName,
  type: z.literal("agent"),
  description: z.string().optional(),
  // "provider:model" override; falls back to the workflow/config default.
  model: z.string().optional(),
  // Opaque Jinja prompt: may reference {{ workflow.input.x }} / {{ step.output.y }}.
  prompt: z.string().min(1),
  tools: Tools.optional(),
  output: OutputSchema.optional(),
  routes: Routes,
});

// `script` — a shell command. ⚠️ This BYPASSES the agent permission gate: a script step
// is direct code execution. It is only reachable unattended after the runner's fail-closed
// posture check (dangerous → forced human_gate). The schema keeps it deliberately spartan:
// a fixed command + args, no shell string to inject into. `command` is the argv[0]; `args`
// are passed as-is (not through a shell).
const ScriptStep = z.strictObject({
  name: StepName,
  type: z.literal("script"),
  description: z.string().optional(),
  command: z.string().min(1),
  args: z.array(z.string()).default([]),
  // A small allow-list of env overrides; the process otherwise inherits the runner's.
  env: z.record(z.string(), z.string()).optional(),
  // Confined to the workflow's cwd; the runner defaults it to that.
  working_dir: z.string().optional(),
  timeout: z.number().int().positive().max(3600).optional(),
  output: OutputSchema.optional(),
  routes: Routes,
});

// `human_gate` — pause and ask a human to choose. Surfaces to the host UI as an approval
// and resolves ONLY through the host's authenticated approval path (never an unsigned
// channel). The chosen option name is exposed as `{{ <name>.choice }}` for routing.
const GateOption = z.strictObject({
  name: StepName,
  description: z.string().optional(),
});
const HumanGateStep = z.strictObject({
  name: StepName,
  type: z.literal("human_gate"),
  description: z.string().optional(),
  // Markdown prompt shown in the approval UI.
  prompt: z.string().min(1),
  options: z.array(GateOption).min(1),
  routes: Routes,
});

// `set` — a pure context transformation (no LLM, no tools). Either a single `value`
// or a `values` map of name → opaque Jinja expression. Free and side-effect-free.
const SetStep = z.strictObject({
  name: StepName,
  type: z.literal("set"),
  description: z.string().optional(),
  value: z.string().optional(),
  values: z.record(z.string(), z.string()).optional(),
  routes: Routes,
})
  .refine((s) => Boolean(s.value) !== Boolean(s.values), {
    message: "a `set` step needs exactly one of `value` or `values`",
    path: ["value"],
  });

// `wait` — sleep. Duration is seconds (number) or a suffixed string ("30s"/"5m"/"1h").
const WaitStep = z.strictObject({
  name: StepName,
  type: z.literal("wait"),
  description: z.string().optional(),
  duration: z.union([z.number().positive(), z.string().regex(/^\d+(\.\d+)?(ms|s|m|h)$/)]),
  reason: z.string().optional(),
  routes: Routes,
});

// A LOCAL relative path to another workflow file. This regex is the security control
// that rejects remote refs: no scheme, no "@" (blocks registry + GitHub refs), no leading
// "/" or "~" (must be relative), and no ".." segment (no escaping the workflows dir). The
// runner additionally resolves + confines it to the workflows directory.
const LocalWorkflowRef = z
  .string()
  .min(1)
  .regex(
    /^(?!.*\.\.)(?![/~])[a-zA-Z0-9_./-]+\.ya?ml$/,
    "sub-workflow must be a local relative *.yaml path (no @registry, url, /abs, ~, or .. refs)",
  );

// `workflow` — run another local workflow as a black-box sub-step.
const SubWorkflowStep = z.strictObject({
  name: StepName,
  type: z.literal("workflow"),
  description: z.string().optional(),
  workflow: LocalWorkflowRef,
  // name → opaque Jinja expression, mapping parent context into the child's input.
  input_mapping: z.record(z.string(), z.string()).optional(),
  output: OutputSchema.optional(),
  routes: Routes,
});

// `terminate` — end the workflow explicitly.
const TerminateStep = z.strictObject({
  name: StepName,
  type: z.literal("terminate"),
  description: z.string().optional(),
  status: z.enum(["success", "failed"]),
  reason: z.string().optional(),
});

// The discriminated union: an unknown `type` fails HARD (no permissive default).
export const Step = z.discriminatedUnion("type", [
  AgentStep,
  ScriptStep,
  HumanGateStep,
  SetStep,
  WaitStep,
  SubWorkflowStep,
  TerminateStep,
]);
export type Step = z.infer<typeof Step>;

// ---- Fan-out groups -------------------------------------------------------------

const FailureMode = z.enum(["fail_fast", "continue_on_error", "all_or_nothing"]);

// `parallel` — a static named group running several already-defined agent steps
// concurrently, bounded by the runner's global concurrency cap.
const ParallelGroup = z.strictObject({
  name: StepName,
  description: z.string().optional(),
  agents: z.array(StepName).min(1),
  failure_mode: FailureMode.default("fail_fast"),
  routes: Routes,
});

// `for_each` — a dynamic group: one gated agent per item of an opaque source ref.
const ForEachGroup = z.strictObject({
  name: StepName,
  description: z.string().optional(),
  // Opaque ref to an array in context, e.g. "planner.output.items".
  source: z.string().min(1),
  as: z.string().min(1).default("item"),
  agent: z.strictObject({
    model: z.string().optional(),
    prompt: z.string().min(1),
    tools: Tools.optional(),
    output: OutputSchema.optional(),
  }),
  // Bounded so a workflow can't outrun the runner's one-at-a-time discipline.
  max_concurrent: z.number().int().positive().max(16).default(4),
  failure_mode: FailureMode.default("fail_fast"),
  key_by: z.string().optional(),
  routes: Routes,
});

// ---- Workflow header + file -----------------------------------------------------

const InputField = z.strictObject({
  type: z.enum(["string", "number", "boolean", "array", "object"]),
  required: z.boolean().default(false),
  description: z.string().optional(),
  default: z.unknown().optional(),
});

const Limits = z.strictObject({
  // Hard cap on step transitions — the loop backstop for a cyclic route graph.
  max_iterations: z.number().int().positive().max(1000).default(50),
  timeout_seconds: z.number().int().positive().max(86_400).optional(),
});

const Header = z.strictObject({
  // Stable id ("w-" + mint time), the key for updates/removal.
  id: z.string(),
  // Human label, unique across workflows; the selector + output dir name.
  name: z.string().min(1),
  description: z.string().optional(),
  // The step the run starts at. Cross-checked against `steps` by validateWorkflow.
  entry_point: StepName,
  // "provider:model" default for agent steps that omit their own.
  default_model: z.string().optional(),
  limits: Limits.default({ max_iterations: 50 }),
});

// The on-disk shape of a workflow file. Steps and fan-out groups share one route
// graph keyed by name; the runner links them. Kept flat rather than nested, so the
// graph is inspectable at a glance.
export const Workflow = z.strictObject({
  workflow: Header,
  input: z.record(z.string(), InputField).optional(),
  steps: z.array(Step).default([]),
  parallel: z.array(ParallelGroup).default([]),
  for_each: z.array(ForEachGroup).default([]),
  // Workflow-level output: name → opaque Jinja expression over final step outputs.
  output: z.record(z.string(), z.string()).optional(),
});
export type Workflow = z.infer<typeof Workflow>;

// A time-ordered workflow id minted once at creation.
export function newWorkflowId(): string {
  return `w-${Date.now()}`;
}

/**
 * Structural validation the flat schema can't express on its own: unique step/group
 * names, a resolvable `entry_point`, and every route/`parallel.agents` target either a
 * defined name or a `$end`/`$fail` sentinel. Returns a list of human-readable errors
 * (empty = valid). The runner MUST call this after Workflow.parse — a dangling `to:`
 * is a graph bug, and we refuse to run a graph with a hole in it (fail-closed).
 */
export function validateWorkflow(wf: Workflow): string[] {
  const errors: string[] = [];

  const stepNames = wf.steps.map((s) => s.name);
  const groupNames = [...wf.parallel.map((g) => g.name), ...wf.for_each.map((g) => g.name)];
  const allNames = [...stepNames, ...groupNames];

  // Duplicate names across every step + group — the graph keys must be unique.
  const seen = new Set<string>();
  for (const n of allNames) {
    if (seen.has(n)) errors.push(`duplicate step/group name "${n}"`);
    seen.add(n);
  }
  const known = seen;

  if (!known.has(wf.workflow.entry_point)) {
    errors.push(`entry_point "${wf.workflow.entry_point}" is not a defined step or group`);
  }

  // Every route target and every parallel member must resolve.
  const isTarget = (t: string) => t === "$end" || t === "$fail" || known.has(t);
  const checkRoutes = (owner: string, routes: { to: string }[]) => {
    for (const r of routes) {
      if (!isTarget(r.to)) errors.push(`"${owner}" routes to unknown target "${r.to}"`);
    }
  };
  for (const s of wf.steps) {
    if ("routes" in s) checkRoutes(s.name, s.routes);
  }
  for (const g of wf.parallel) {
    checkRoutes(g.name, g.routes);
    for (const a of g.agents) {
      if (!known.has(a)) errors.push(`parallel group "${g.name}" references unknown step "${a}"`);
    }
  }
  for (const g of wf.for_each) checkRoutes(g.name, g.routes);

  return errors;
}