import { z } from 'zod';
import { dagNodeSchema, effortLevelSchema, thinkingConfigSchema } from './dag-node.js';

export const modelReasoningEffortSchema = z.enum(['low', 'medium', 'high']);

export type ModelReasoningEffort = z.infer<typeof modelReasoningEffortSchema>;

export const webSearchModeSchema = z.enum(['off', 'auto', 'on']);

export type WebSearchMode = z.infer<typeof webSearchModeSchema>;

export const workflowRequirementSchema = z.object({
  /** Human-readable name of the requirement. */
  name: z.string().min(1, 'requirement name must not be empty'),

  /** Short description of what the requirement enforces. */
  description: z.string().optional(),
});

export type WorkflowRequirement = z.infer<typeof workflowRequirementSchema>;

export const workflowWorktreePolicySchema = z.object({
  /** Whether worktree isolation is enabled for this workflow. */
  enabled: z.boolean().optional(),
});

export type WorkflowWorktreePolicy = z.infer<typeof workflowWorktreePolicySchema>;

export const sandboxConfigSchema = z.object({
  /** Whether sandboxing is enabled. */
  enabled: z.boolean().default(false),

  /** Docker image to use for sandboxing. */
  image: z.string().optional(),

  /** Memory limit in megabytes. */
  memory_mb: z.number().int().positive().optional(),

  /** CPU limit in cores. */
  cpu_cores: z.number().positive().optional(),

  /** Network access inside the sandbox. */
  network: z.boolean().default(false),
});

export type SandboxConfig = z.infer<typeof sandboxConfigSchema>;

export const providerOverridesSchema = z.record(
  z.string(),
  z.object({
    model: z.string().optional(),
    base_url: z.string().optional(),
    api_key_env: z.string().optional(),
  }),
);

export type ProviderOverrides = z.infer<typeof providerOverridesSchema>;

export const workflowBaseSchema = z.object({
  /** Human-readable workflow name. */
  name: z.string().min(1, 'workflow name must not be empty'),

  /** Short description of what the workflow does. */
  description: z.string().optional(),

  /** Default provider for all nodes. */
  provider: z.string().optional(),

  /** Default model for all nodes. */
  model: z.string().optional(),

  /** Reasoning effort level for the model. */
  modelReasoningEffort: modelReasoningEffortSchema.optional(),

  /** Web search mode for the model. */
  webSearchMode: webSearchModeSchema.optional(),

  /** Whether the workflow can pause for human input. */
  interactive: z.boolean().default(false),

  /** Default effort level for all nodes. */
  effort: effortLevelSchema.optional(),

  /** Default thinking configuration for all nodes. */
  thinking: thinkingConfigSchema.optional(),

  /** Fallback model when the primary model is unavailable. */
  fallbackModel: z.string().optional(),

  /** Tags for categorisation and filtering. */
  tags: z.array(z.string()).default([]),

  /** Requirements that must be met before the workflow can run. */
  requires: z.array(workflowRequirementSchema).default([]),

  /** Whether this workflow modifies the git checkout. */
  mutates_checkout: z.boolean().default(false),

  /** Whether to persist conversation sessions between runs. */
  persist_sessions: z.boolean().default(false),

  /** Per-provider configuration overrides. */
  provider_overrides: providerOverridesSchema.optional(),

  /** Sandbox configuration for the workflow. */
  sandbox: sandboxConfigSchema.optional(),

  /** Worktree policy for the workflow. */
  worktree: workflowWorktreePolicySchema.optional(),
});

export type WorkflowBase = z.infer<typeof workflowBaseSchema>;

export const workflowDefinitionSchema = workflowBaseSchema.extend({
  /** The DAG nodes that make up this workflow. */
  nodes: z.array(dagNodeSchema),
});

export type WorkflowDefinition = z.infer<typeof workflowDefinitionSchema>;

export const WorkflowSourceSchema = z.enum(['bundled', 'global', 'project']);

export type WorkflowSource = z.infer<typeof WorkflowSourceSchema>;

export const workflowExecutionResultSchema = z.discriminatedUnion('status', [
  z.object({
    status: z.literal('success'),
    output: z.string(),
    metadata: z.record(z.string(), z.unknown()).optional(),
  }),
  z.object({
    status: z.literal('fail'),
    error: z.string(),
    node_id: z.string().optional(),
    metadata: z.record(z.string(), z.unknown()).optional(),
  }),
  z.object({
    status: z.literal('paused'),
    node_id: z.string(),
    reason: z.string(),
    metadata: z.record(z.string(), z.unknown()).optional(),
  }),
]);

export type WorkflowExecutionResult = z.infer<typeof workflowExecutionResultSchema>;

export const workflowWithSourceSchema = z.object({
  definition: workflowDefinitionSchema,
  source: WorkflowSourceSchema,
  path: z.string().optional(),
});

export type WorkflowWithSource = z.infer<typeof workflowWithSourceSchema>;

export const workflowLoadErrorSchema = z.object({
  message: z.string(),
  path: z.string().optional(),
  cause: z.unknown().optional(),
});

export type WorkflowLoadError = z.infer<typeof workflowLoadErrorSchema>;

export const workflowLoadResultSchema = z.union([
  workflowWithSourceSchema,
  workflowLoadErrorSchema,
]);

export type WorkflowLoadResult = z.infer<typeof workflowLoadResultSchema>;

export const loadCommandResultSchema = z.discriminatedUnion('status', [
  z.object({
    status: z.literal('success'),
    definition: workflowDefinitionSchema,
    source: WorkflowSourceSchema,
    path: z.string().optional(),
  }),
  z.object({
    status: z.literal('fail'),
    reason: z.enum([
      'not_found',
      'parse_error',
      'validation_error',
      'permission_denied',
    ]),
    message: z.string(),
    path: z.string().optional(),
  }),
]);

export type LoadCommandResult = z.infer<typeof loadCommandResultSchema>;