import { _generateImage, _generateImageLocal, _imageSources, _imageDestination } from "agency-lang/stdlib-lib/image.js"
import { _cropImage, _imageSize, _pasteImages } from "agency-lang/stdlib-lib/imageTools.js"
import { _realTarget } from "agency-lang/stdlib-lib/fs.js"
import { _localImageInputs } from "agency-lang/stdlib-lib/localImageInputs.js"
import { dirname, basename } from "std::path"

/** @module
  @summary Generate, edit, crop, and combine images.
  Generate images from a text prompt. `generateImage` uses a hosted
provider (OpenAI, Google, or an open-source model via LiteLLM / Together),
and can also edit input images. `generateImageLocal` uses a model on this
machine, served by `agency local serve --image`. Both return base64 you can
persist with `writeBinary()` or send onward with `std::thread`'s `image(...)`.

  ```ts
  import { generateImage, generateImageLocal } from "std::image"

  node main() {
    const r = generateImage("a red bicycle in the rain", size: "1024x1024")
    if (isFailure(r)) { print("failed: ${r.error}"); return }
    writeBinary("bike.png", r.value.base64)

    const local = generateImageLocal("a lighthouse in a storm", "z-image-turbo")
    if (isFailure(local)) { print("failed: ${local.error}"); return }
    writeBinary("lighthouse.png", local.value.base64)
  }
  ```
*/

// Sending a local image to a hosted provider for generateImage to edit.
// Its own effect, not std::readImage: reading a file on this machine and
// sending it to a third party are different permissions, and a policy
// that approves local reads must not approve uploads. The payload names
// the file and where it goes: the provider and model the default client
// resolves, and the base URL given for openai-compat or litellm ("" when
// none was). An "always" answer pins both the folder and the model.
@alwaysUnder(dir)
@always(model)
effect std::uploadImage { dir: string, filename: string, provider: string, model: string, baseUrl: string }

// Cutting a box out of one image and writing it as another. Both paths are
// real spellings, since one approval is for "read this, write that". An
// "always" answer pins both folders, as std::copy does, so a rule saved
// for one crop cannot read an image from anywhere else.
@alwaysUnder(dir, outDir)
effect std::cropImage { dir: string, filename: string, outDir: string, outFilename: string }

// Laying several images out on one canvas. Every input is listed the way
// applyPatch lists the files a patch touches. No "always" tag, as with
// applyPatch: a list of inputs cannot be pinned to one folder.
effect std::pasteImages { files: string[], outDir: string, outFilename: string }

/** A region of an image, normalized to 0..1 with the origin at the top
left, as std::ocr and std::vision return it. */
export type ImageBox = {
  x: number,
  y: number,
  width: number,
  height: number,
}

/** Width and height in pixels. */
export type ImageSize = {
  width: number,
  height: number,
}

/** Image bytes encoded as base64, with their MIME type. */
export type GeneratedImage = {
  base64: string;
  mimeType: string
}

export type ImageQuality = "low" | "medium" | "high" | "auto"

/**
Use a hosted image model for generation or editing. Each local input file
requires upload approval. URL and data-URI inputs do not read a local file.
Save the returned base64 with `writeBinary` when you need a file.
*/
@capabilities(fileRead, llm)
export def generateImage(
  prompt: string,
  model: string = "",
  provider: string = "",
  size: string = "",
  quality: ImageQuality = "auto",
  images: string[] = [],
  apiKey: string = "",
  baseUrl: string = "",
): Result<GeneratedImage> raises <std::uploadImage> {
  """
  Generate or edit an image with a hosted provider. Local input images are uploaded.

  @param prompt - What to generate or how to edit the inputs.
  @param model - Image model, or empty for the provider default.
  @param provider - Override the provider derived from the model name.
  @param size - Provider-supported size such as "1024x1024".
  @param images - Input images as local PNG/JPEG/GIF/WebP paths, HTTP(S) URLs, or data URIs.
  @param apiKey - Override the API key.
  @param baseUrl - Custom endpoint for openai-compat or litellm.
  """
  // Every path is checked before the first interrupt, so a bad one fails
  // here instead of after an approval.
  const sources = try _imageSources(images)
  if (isFailure(sources)) {
    return sources
  }
  const destination = _imageDestination(model, provider)
  for (source in sources.value) {
    if (source.local) {
      raise std::uploadImage("Send this image to ${destination.provider} (${destination.model} model)?", {
        dir: dirname(source.source),
        filename: basename(source.source),
        provider: destination.provider,
        model: destination.model,
        baseUrl: baseUrl,
      })
    }
  }
  return _generateImage(
    prompt,
    model,
    provider,
    size,
    quality,
    sources.value,
    apiKey,
    baseUrl,
  )
}

/** Generated image data and the seed used for reproducibility. */
export type LocalImage = {
  base64: string;
  mimeType: string;
  seed: number
}

/**
Start the model with `agency local serve --image <model>`. Save the returned
base64 with `writeBinary`.

For SDXL, a LoRA adapter changes the style or character. Put its
`.safetensors` file in `client.adaptersDir` and pass the filename without
that extension. A ControlNet uses a drawing to guide the composition. Put
it in `client.controlnetsDir` and pair `controlnet` with `controlImage`.
The drawing is scaled to fit and centered on black. Scribble ControlNets
expect white lines on black; invert a dark-on-white drawing.

FLUX.2 [klein] and Qwen-Image-Edit accept reference pictures in `images`.
Other models use `startImage` to redraw a picture. Its image is scaled to
cover the requested size, cropping any overflow. Leave `size` empty to
preserve the input's aspect ratio.

```ts
generateImageLocal("a vase of sunflowers", "z-image-turbo",
  startImage: "photo.png", mask: "vase-mask.png")
```

A mask selects the part to redraw: white changes, black stays, and grey
blends the two. Transparency counts as black. The mask must match the
start image's dimensions. Strength defaults can differ with a mask;
Z-Image Turbo uses 1.0 with one and 0.6 without one.
*/
@capabilities(fileRead, subprocess, env, llm)
export def generateImageLocal(
  prompt: string,
  model: string,
  size: string = "",
  steps: number | null = null,
  guidance: number | null = null,
  seed: number | null = null,
  negativePrompt: string = "",
  format: string = "png",
  lora: string = "",
  loraScale: number | null = null,
  controlnet: string = "",
  controlImage: string = "",
  controlScale: number | null = null,
  invertControlImage: boolean = false,
  images: string[] = [],
  startImage: string = "",
  strength: number | null = null,
  mask: string = "",
): Result<LocalImage> raises <std::readImage> {
  """
  Generate or edit an image with a running local model. Returns image data
  for saving with writeBinary.

  @param model - Catalog name, diffusers: URI, or model directory.
  @param size - WxH with each side a multiple of 16 from 256 to 2048 and at most 4 million pixels; empty uses the input aspect ratio or 1024x1024.
  @param steps - Refinement passes, or null for the model default.
  @param guidance - Prompt adherence; leave null for the model default since some models reject overrides.
  @param seed - Random seed, or null to choose one and return it with the image.
  @param negativePrompt - Unwanted content; leave empty unless the model supports it.
  @param format - Output format: png, jpeg, or webp.
  @param lora - SDXL adapter name in client.adaptersDir, without .safetensors.
  @param loraScale - Adapter strength from 0 to 2, or null for 1.
  @param controlnet - SDXL ControlNet folder in client.controlnetsDir, paired with controlImage.
  @param controlImage - Local conditioning drawing, paired with controlnet.
  @param controlScale - ControlNet strength from 0 to 2, or null for 1.
  @param invertControlImage - Invert dark lines on white for a scribble ControlNet.
  @param images - Up to four local references for FLUX.2 [klein] or Qwen-Image-Edit; Qwen-Image-Edit requires at least one.
  @param startImage - Local picture to redraw; unsupported by FLUX.2 [klein] and Qwen-Image-Edit.
  @param strength - Fraction of startImage to redraw, greater than 0 and at most 1, or null for the model default.
  @param mask - Local mask matching startImage dimensions, with white to redraw and black to keep.
  """
  // Every input path is checked before the first interrupt, so a call that
  // is going to fail asks for nothing.
  const inputs = try _localImageInputs(
    controlnet,
    controlImage,
    controlScale,
    invertControlImage,
    images,
    startImage,
    strength,
    mask,
  )
  if (isFailure(inputs)) {
    return inputs
  }
  for (file in inputs.value.files) {
    raise std::readImage(file.question, { dir: file.dir, filename: file.filename })
  }
  return _generateImageLocal(
    prompt,
    model,
    size,
    steps,
    guidance,
    seed,
    negativePrompt,
    format,
    lora,
    loraScale,
    inputs.value,
  )
}

@capabilities(fileRead, fileWrite, subprocess, env)
export def cropImage(
  path: string,
  box: ImageBox,
  outPath: string,
  pad: number = 0,
  square: boolean = false,
): Result<string> raises <std::cropImage> {
  """
  Crop an image into a new file without overwriting an existing output.

  @param path - Input PNG, JPEG, or WebP file.
  @param box - Region normalized to 0..1 from the top left.
  @param outPath - Output path whose extension selects the format.
  @param pad - Expand each side by this fraction of the box size.
  @param square - Pad to a square using the edge color.
  """
  const real = try _realTarget(path)
  if (isFailure(real)) {
    return real
  }
  const out = try _realTarget(outPath)
  if (isFailure(out)) {
    return out
  }
  return interrupt std::cropImage("Cut a region out of this image and write it as a new file?", {
    dir: dirname(real.value),
    filename: basename(real.value),
    outDir: dirname(out.value),
    outFilename: basename(out.value),
  })
  return try _cropImage(real.value, box, out.value, pad, square)
}

@capabilities(fileRead, subprocess, env, llm)
export idempotent def imageSize(path: string): Result<ImageSize> raises <std::readImage> {
  """
  The width and height of an image in pixels.

  @param path - The image: PNG, JPEG, WebP, or GIF
  """
  const real = try _realTarget(path)
  if (isFailure(real)) {
    return real
  }
  return interrupt std::readImage("Read this image's size?", {
    dir: dirname(real.value),
    filename: basename(real.value),
  })
  return try _imageSize(real.value)
}

/**
Use two columns for a before-and-after comparison, or more columns for a
contact sheet. Images keep their original sizes on a white canvas.
*/
@capabilities(fileRead, fileWrite, subprocess, env)
export def pasteImages(
  paths: string[],
  outPath: string,
  columns: number = 2,
): Result<string> raises <std::pasteImages> {
  """
  Arrange images in rows and write a new file without overwriting existing output.

  @param paths - Input images in reading order.
  @param outPath - Output path whose extension selects the format.
  @param columns - Number of images per row.
  """
  // Every path is checked before the interrupt, so a bad one fails here
  // instead of being approved and failing afterwards.
  const reals: string[] = []
  for (p in paths) {
    const real = try _realTarget(p)
    if (isFailure(real)) {
      return real
    }
    reals.push(real.value)
  }
  const out = try _realTarget(outPath)
  if (isFailure(out)) {
    return out
  }
  return interrupt std::pasteImages("Paste these images onto one canvas and write it as a new file?", {
    files: reals,
    outDir: dirname(out.value),
    outFilename: basename(out.value),
  })
  return try _pasteImages(reals, out.value, columns)
}