import {
  defineStage,
  WorkflowBuilder,
  createWorkflowRuntime,
  createPrismaWorkflowPersistence,
  createPrismaJobQueue,
  createPrismaAICallLogger,
} from "@bratsos/workflow-engine";
import { z } from "zod";

// 1. Define a stage
const processStage = defineStage({
  id: "process",
  name: "Process Data",
  schemas: {
    input: z.object({ data: z.string() }),
    output: z.object({ result: z.string() }),
    config: z.object({ verbose: z.boolean().default(false) }),
  },
  async execute(ctx) {
    const processed = ctx.input.data.toUpperCase();
    return { output: { result: processed } };
  },
});

// 2. Build a workflow
const workflow = new WorkflowBuilder(
  "my-workflow",
  "My Workflow",
  "Processes data",
  z.object({ data: z.string() }),
  z.object({ result: z.string() })
)
  .pipe(processStage)
  .build();

// 3. Create runtime and execute
const runtime = createWorkflowRuntime({
  persistence: createPrismaWorkflowPersistence(prisma),
  jobQueue: createPrismaJobQueue(prisma),
  aiCallLogger: createPrismaAICallLogger(prisma),
  registry: { getWorkflow: (id) => (id === "my-workflow" ? workflow : null) },
});

await runtime.start();
const { workflowRunId } = await runtime.createRun({
  workflowId: "my-workflow",
  input: { data: "hello" },
});

const myStage = defineStage({
  id: "my-stage",           // Unique identifier
  name: "My Stage",         // Display name
  description: "Optional",  // Description
  dependencies: ["prev"],   // Required previous stages

  schemas: {
    input: InputSchema,     // Zod schema or "none"
    output: OutputSchema,   // Zod schema
    config: ConfigSchema,   // Zod schema with defaults
  },

  async execute(ctx) {
    // Access input, config, workflow context
    const { input, config, workflowContext } = ctx;

    // Get output from previous stages
    const prevOutput = ctx.require("prev");      // Throws if missing
    const optOutput = ctx.optional("other");     // Returns undefined if missing

    // Access services
    await ctx.log("INFO", "Processing...");
    await ctx.storage.save("key", data);

    return {
      output: { ... },
      customMetrics: { itemsProcessed: 10 },
      artifacts: { rawData: data },
    };
  },
});

const batchStage = defineAsyncBatchStage({
  id: "batch-process",
  name: "Batch Process",
  mode: "async-batch",  // Required

  schemas: {
    input: "none",
    output: OutputSchema,
    config: ConfigSchema,
  },

  async execute(ctx) {
    // Check if resuming from suspension
    if (ctx.resumeState) {
      return { output: ctx.resumeState.cachedResult };
    }

    // Submit batch and suspend
    const batchId = await submitBatchJob(ctx.input);

    return {
      suspended: true,
      state: {
        batchId,
        submittedAt: new Date().toISOString(),
        pollInterval: 60000,
        maxWaitTime: 3600000,
      },
      pollConfig: {
        pollInterval: 60000,
        maxWaitTime: 3600000,
        nextPollAt: new Date(Date.now() + 60000),
      },
    };
  },

  async checkCompletion(suspendedState, ctx) {
    const status = await checkBatchStatus(suspendedState.batchId);

    if (status === "completed") {
      const results = await getBatchResults(suspendedState.batchId);
      return { ready: true, output: { results } };
    }

    if (status === "failed") {
      return { ready: false, error: "Batch failed" };
    }

    return { ready: false, nextCheckIn: 60000 };
  },
});

const workflow = new WorkflowBuilder(
  "workflow-id",
  "Workflow Name",
  "Description",
  InputSchema,
  OutputSchema
)
  .pipe(stage1)                    // Sequential
  .pipe(stage2)
  .parallel([stage3a, stage3b])    // Parallel execution
  .pipe(stage4)
  .build();

// Workflow utilities
workflow.getStageIds();            // ["stage1", "stage2", ...]
workflow.getExecutionPlan();       // Grouped by execution order
workflow.getDefaultConfig();       // Default config for all stages
workflow.validateConfig(config);   // Validate config object

// In a stage - use the standard convention
const ai = runtime.createAIHelper(`workflow.${ctx.workflowRunId}.stage.${ctx.stageId}`);

// Later, query costs by prefix:
const workflowCost = await aiLogger.getStats(`workflow.${workflowRunId}`);  // All stages
const stageCost = await aiLogger.getStats(`workflow.${workflowRunId}.stage.extraction`);  // One stage

// Text generation
const { text, cost } = await ai.generateText("gemini-2.5-flash", prompt, {
  temperature: 0.7,
  maxTokens: 1000,
});

// Structured output
const { object } = await ai.generateObject(
  "gemini-2.5-flash",
  prompt,
  z.object({ items: z.array(z.string()) })
);

// Embeddings
const { embedding, embeddings } = await ai.embed(
  "text-embedding-004",
  ["text1", "text2"],
  { dimensions: 768 }
);

// Batch operations (50% cost savings)
const batch = ai.batch("claude-sonnet-4-20250514", "anthropic");
const handle = await batch.submit([
  { id: "req1", prompt: "..." },
  { id: "req2", prompt: "...", schema: OutputSchema },
]);

// Check status and get results
const status = await batch.getStatus(handle.id);
const results = await batch.getResults(handle.id);

// Get aggregated stats
const stats = await ai.getStats();
console.log(`Cost: $${stats.totalCost.toFixed(4)}, Tokens: ${stats.totalInputTokens + stats.totalOutputTokens}`);

// Required enums
enum Status {
  PENDING
  RUNNING
  SUSPENDED
  COMPLETED
  FAILED
  CANCELLED
  SKIPPED
}

enum LogLevel {
  DEBUG
  INFO
  WARN
  ERROR
}

enum ArtifactType {
  STAGE_OUTPUT
  ARTIFACT
  METADATA
}

// ✅ REQUIRED: WorkflowRun
model WorkflowRun {
  id           String    @id @default(cuid())
  createdAt    DateTime  @default(now())
  updatedAt    DateTime  @updatedAt
  workflowId   String
  workflowName String
  workflowType String
  status       Status    @default(PENDING)
  startedAt    DateTime?
  completedAt  DateTime?
  duration     Int?                          // ⚠️ Must be "duration", not "durationMs"
  input        Json
  output       Json?
  config       Json      @default("{}")
  totalCost    Float     @default(0)
  totalTokens  Int       @default(0)
  priority     Int       @default(5)

  stages       WorkflowStage[]
  logs         WorkflowLog[]
  artifacts    WorkflowArtifact[]

  @@index([status])
  @@map("workflow_runs")
}

// ✅ REQUIRED: WorkflowStage
model WorkflowStage {
  id              String    @id @default(cuid())
  createdAt       DateTime  @default(now())
  updatedAt       DateTime  @updatedAt
  workflowRunId   String
  stageId         String
  stageName       String
  stageNumber     Int
  executionGroup  Int
  status          Status    @default(PENDING)
  startedAt       DateTime?
  completedAt     DateTime?
  duration        Int?                        // ⚠️ Must be "duration", not "durationMs"
  inputData       Json?
  outputData      Json?
  config          Json?
  suspendedState  Json?
  resumeData      Json?
  nextPollAt      DateTime?
  pollInterval    Int?
  maxWaitUntil    DateTime?
  metrics         Json?
  embeddingInfo   Json?
  errorMessage    String?

  workflowRun     WorkflowRun        @relation(fields: [workflowRunId], references: [id], onDelete: Cascade)
  logs            WorkflowLog[]
  artifacts       WorkflowArtifact[]

  @@unique([workflowRunId, stageId])
  @@map("workflow_stages")
}

// ✅ REQUIRED: WorkflowLog (missing = "Cannot read 'create'" error)
model WorkflowLog {
  id              String         @id @default(cuid())
  createdAt       DateTime       @default(now())
  workflowRunId   String?
  workflowStageId String?
  level           LogLevel
  message         String
  metadata        Json?

  workflowRun     WorkflowRun?   @relation(fields: [workflowRunId], references: [id], onDelete: Cascade)
  workflowStage   WorkflowStage? @relation(fields: [workflowStageId], references: [id], onDelete: Cascade)

  @@index([workflowRunId])
  @@map("workflow_logs")
}

// ✅ REQUIRED: WorkflowArtifact (missing = "Cannot read 'upsert'" error)
model WorkflowArtifact {
  id              String       @id @default(cuid())
  createdAt       DateTime     @default(now())
  updatedAt       DateTime     @updatedAt
  workflowRunId   String
  workflowStageId String?
  key             String
  type            ArtifactType
  data            Json
  size            Int
  metadata        Json?

  workflowRun     WorkflowRun   @relation(fields: [workflowRunId], references: [id], onDelete: Cascade)
  workflowStage   WorkflowStage? @relation(fields: [workflowStageId], references: [id], onDelete: Cascade)

  @@unique([workflowRunId, key])
  @@map("workflow_artifacts")
}

// ✅ REQUIRED: JobQueue
model JobQueue {
  id            String    @id @default(cuid())
  createdAt     DateTime  @default(now())
  updatedAt     DateTime  @updatedAt
  workflowRunId String
  stageId       String
  status        Status    @default(PENDING)
  priority      Int       @default(5)
  workerId      String?
  lockedAt      DateTime?
  startedAt     DateTime?
  completedAt   DateTime?
  attempt       Int       @default(0)
  maxAttempts   Int       @default(3)
  lastError     String?
  nextPollAt    DateTime?
  payload       Json      @default("{}")

  @@index([status, nextPollAt])
  @@map("job_queue")
}

import {
  createPrismaWorkflowPersistence,
  createPrismaJobQueue,
  createPrismaAICallLogger,
} from "@bratsos/workflow-engine/persistence/prisma";

// PostgreSQL (default)
const persistence = createPrismaWorkflowPersistence(prisma);
const jobQueue = createPrismaJobQueue(prisma);

// SQLite - MUST pass databaseType option
const persistence = createPrismaWorkflowPersistence(prisma, { databaseType: "sqlite" });
const jobQueue = createPrismaJobQueue(prisma, { databaseType: "sqlite" });

const aiCallLogger = createPrismaAICallLogger(prisma);

const runtime = createWorkflowRuntime({
  persistence,
  jobQueue,
  aiCallLogger,
  registry: {
    getWorkflow: (id) => workflowMap[id] ?? null,
  },

  // Optional configuration
  pollIntervalMs: 10000,      // Orchestration poll interval
  jobPollIntervalMs: 1000,    // Job dequeue interval
  staleJobThresholdMs: 60000, // Stale job timeout
  workerId: "worker-1",       // Custom worker ID
});

// Lifecycle
await runtime.start();        // Start processing
runtime.stop();               // Graceful shutdown

// Manual operations
await runtime.createRun({ workflowId, input });
await runtime.transitionWorkflow(runId);
await runtime.pollSuspendedStages();

import {
  TestWorkflowPersistence,
  TestJobQueue,
  MockAIHelper,
} from "@bratsos/workflow-engine/testing";

const persistence = new TestWorkflowPersistence();
const jobQueue = new TestJobQueue();
const mockAI = new MockAIHelper();

// Configure mock responses
mockAI.mockGenerateText("Expected response");
mockAI.mockGenerateObject({ items: ["a", "b"] });

// Test stage execution
const result = await myStage.execute({
  input: { data: "test" },
  config: { verbose: true },
  workflowContext: {},
  workflowRunId: "test-run",
  stageId: "my-stage",
  log: async () => {},
  storage: persistence.createStorage("test-run"),
});