requires "pocketflow >= 0.2.0"

import asyncdispatch, json
import pocketflow

let step1 = newNode(
  exec = proc(ctx: PfContext, params: JsonNode, prepRes: JsonNode): Future[JsonNode] {.async, closure, gcsafe.} =
    ctx["message"] = %"Hello from PocketFlow!"
    return %"done"
)

let step2 = newNode(
  exec = proc(ctx: PfContext, params: JsonNode, prepRes: JsonNode): Future[JsonNode] {.async, closure, gcsafe.} =
    echo ctx["message"].getStr()
    return %"complete"
)

# Chain nodes with >> operator
discard step1 >> step2

let flow = newFlow(step1)
let ctx = newPfContext()
waitFor flow.internalRun(ctx)

import asyncdispatch, json, os
import pocketflow

let llm = newLlmClient(
  provider = OpenAI,
  apiKey = getEnv("OPENAI_API_KEY"),
  model = "gpt-4"
)

let chatNode = newNode(
  exec = proc(ctx: PfContext, params: JsonNode, prepRes: JsonNode): Future[JsonNode] {.async, closure, gcsafe.} =
    let response = await llm.chat(ctx["prompt"].getStr())
    ctx["response"] = %response
    return %"done"
)

let flow = newFlow(chatNode)
let ctx = newPfContext()
ctx["prompt"] = %"Tell me a fun fact!"
waitFor flow.internalRun(ctx)
echo ctx["response"].getStr()

import pocketflow

# Chunk document
let document = readFile("document.txt")
let opts = newChunkingOptions(strategy = Sentences, chunkSize = 500)
var chunks = chunkDocument(document, opts)

# Search with embeddings
let queryEmbedding = @[0.1, 0.5, 0.3]
let topChunks = findTopK(queryEmbedding, chunks, k = 3)

# Rerank results
let reranked = rerankChunks(chunks, "search query")
echo "Top result: ", topChunks[0].chunk.text

let node = newNode(
  prep = proc(ctx: PfContext, params: JsonNode): Future[JsonNode] {.async, closure, gcsafe.} =
    return %42
  ,
  exec = proc(ctx: PfContext, params: JsonNode, prepRes: JsonNode): Future[JsonNode] {.async, closure, gcsafe.} =
    let value = prepRes.getInt()
    return %(value * 2)
  ,
  post = proc(ctx: PfContext, params: JsonNode, prepRes: JsonNode, execRes: JsonNode): Future[string] {.async, closure, gcsafe.} =
    return "next_step"
  ,
  maxRetries = 3,
  waitMs = 1000
)

# Linear flow with >> operator
discard node1 >> node2 >> node3
let flow = newFlow(node1)

# Branching flow with - operator
discard router - "success" >> successNode
discard router - "error" >> errorNode
let branchFlow = newFlow(router)

let batchNode = newBatchNode(
  prep = proc(ctx: PfContext, params: JsonNode): Future[JsonNode] {.async, closure, gcsafe.} =
    return %[%"item1", %"item2", %"item3"]
  ,
  execItem = proc(ctx: PfContext, params: JsonNode, item: JsonNode): Future[JsonNode] {.async, closure, gcsafe.} =
    return %(item.getStr().toUpperAscii())
  ,
  post = proc(ctx: PfContext, params: JsonNode, prepRes: JsonNode, execRes: JsonNode): Future[string] {.async, closure, gcsafe.} =
    ctx["results"] = execRes
    return "done"
)

# Parallel processing with concurrency limit
let parallelNode = newParallelBatchNode(
  prep = ...,
  execItem = ...,
  maxConcurrency = 5
)

let llm = newLlmClient(
  provider = OpenAI,
  apiKey = "sk-...",
  model = "gpt-4o-mini"
)

let llm = newLlmClient(
  provider = Anthropic,
  apiKey = "sk-ant-...",
  model = "claude-3-5-sonnet-20241022"
)

let llm = newLlmClient(
  provider = Google,
  apiKey = "...",
  model = "gemini-1.5-flash"
)

let llm = newLlmClient(
  provider = Ollama,
  model = "llama3"
)

let tracker = newCostTracker()
let llm = newLlmClient(
  provider = OpenAI,
  apiKey = apiKey,
  costTracker = tracker
)

# Use the client...
let response = await llm.chat("Hello")

# Get summary
let summary = tracker.getSummary()
echo "Total cost: $", summary["total_cost_usd"].getFloat()

let cache = newCache()
let llm = newLlmClient(
  provider = OpenAI,
  apiKey = apiKey,
  cache = cache
)

# First call hits API
let response1 = await llm.chat("Same prompt")

# Second call returns cached result
let response2 = await llm.chat("Same prompt")

# Clear cache when needed
cache.clear()

let conditional = newConditionalNode(
  condition = proc(ctx: PfContext, params: JsonNode): Future[bool] {.async, closure, gcsafe.} =
    return ctx["score"].getInt() > 80
  ,
  trueNode = highScoreNode,
  falseNode = lowScoreNode
)

let loop = newLoopNode(
  items = proc(ctx: PfContext, params: JsonNode): Future[JsonNode] {.async, closure, gcsafe.} =
    return %[%1, %2, %3, %4, %5]
  ,
  body = processItemNode,
  maxIterations = 100,
  aggregateResults = true
)
# Access current item via ctx["__loop_item__"]
# Access current index via ctx["__loop_index__"]

let timeout = newTimeoutNode(
  innerNode = slowOperationNode,
  timeoutMs = 5000  # 5 seconds
)

let mapNode = newMapNode(
  mapFunc = proc(ctx: PfContext, item: JsonNode): Future[JsonNode] {.async, closure, gcsafe.} =
    return %(item.getInt() * 2)
  ,
  maxConcurrency = 3
)
# Set items via ctx["__map_items__"]
# Results stored in ctx["__map_results__"]

let store = newStateStore(".pocketflow_state")

# Capture and save state
let state = captureState(ctx, "my_flow", %*{"version": "1.0"})
saveState(store, state)

# Load and restore state
let loadedState = loadState(store, "my_flow")
let restoredCtx = newPfContext()
restoreContext(restoredCtx, loadedState)

# Structured logging
logStructured(Info, "Processing started", [("user_id", "123")])

# Metrics
recordMetric("requests_processed", 1.0, [("status", "success")])

# Tracing with spans
let span = newSpan("expensive_operation")
# ... do work ...
finish(span)
echo "Duration: ", getDurationMs(span), "ms"

# Get metrics summary
echo getMetricsSummary().pretty()