import {
  _setLlmOptions,
  _registerProviderModule,
  _listHostedModels,
  _hostedModelInfo,
  _loadModelData,
  _modelSupportsInput,
} from "agency-lang/stdlib-lib/llm.js"
import { _realTarget } from "agency-lang/stdlib-lib/fs.js"
import { env } from "std::system"
import { dirname, basename } from "std::path"

/** @module
Choose which model and options your `llm()` calls use at runtime.

`setModel` and `setLlmOptions` set defaults for later `llm()` calls,
and `pickProvider` finds an available provider from your API-key
environment variables. Defaults are branch-scoped and survive
interrupt/resume. A per-call `llm(..., { ... })` option always wins.

```ts
import { setModel, setLlmOptions } from "std::llm"

node main() {
  setModel("claude-opus-4-8")
  setLlmOptions({ temperature: 0.2 })
  const answer: string = llm("Say hello")
  print(answer)
}
```
*/

/** How OpenRouter picks the company that runs the model. Provider names
   are OpenRouter's slugs, such as `"z-ai"` or `"fireworks"`. Mirrors
   smoltalk's `OpenRouterProviderPreferences`. */
export type OpenRouterProviderPreferences = {
  order?: string[];
  only?: string[];
  ignore?: string[];
  allowFallbacks?: boolean;
  requireParameters?: boolean;
  dataCollection?: "allow" | "deny";
  sort?: "price" | "throughput" | "latency";
  quantizations?: string[]
}

/** OpenRouter-only `llm()` options. Other providers ignore them. */
export type OpenRouterOptions = {
  provider?: OpenRouterProviderPreferences
}

/** When the tool loop warns the model that its tool rounds are running
   out. See `warnAtRoundsLeft` on `LlmDefaults`. */
export type WarnAtRoundsLeft = {
  remaining?: number;
  text?: string;
  finalWarning?: boolean;
  finalWarningText?: string
}

/** The base URL for each provider that takes one. Each field is
   optional. A provider left out uses the URL from `agency.json`, or its
   default. */
export type LlmBaseUrls = {
  openAiCompat?: string;
  openRouter?: string;
  deepInfra?: string;
  liteLlm?: string;
  mlx?: string;
  ollama?: string;
  typesafe?: string
}

/** Default options for `llm()` calls. Every field is optional; only the
   fields you pass are changed. `provider` is normally derived from the
   model name. Set it only when the name doesn't imply a provider (e.g.
   a custom or local model).

   `thinking` turns a model's thinking on or off, with an optional budget
   of tokens it may think for before it has to answer. On a local MLX
   model, `{ enabled: false }` skips the thinking entirely, which is the
   biggest saving there is on a small task.

   `replyLimits` changes how the MLX chat server treats a reply that goes
   in circles: `hedgeLimit` is how many second thoughts ("But wait",
   "Hmm") it allows and `repeatLimit` how many times one sentence may
   repeat, both within the last two thousand tokens, and `limitAnswers`
   watches the answer as well as the thinking. `0` turns a limit off,
   which a song with a chorus needs. Other providers ignore it.

   `openRouter.provider` chooses which companies OpenRouter may send the
   call to, for example `{ provider: { ignore: ["open-inference"] } }` to
   skip one, or `{ provider: { only: ["fireworks"], allowFallbacks: false } }`
   to use only one. Other providers ignore it. `rawAttributes` is spread
   into the provider's request body as given, for a request field Agency
   has no option for, such as `{ top_k: 5 }`.

   `warnAtRoundsLeft` warns the model before it runs out of tool rounds
   (`maxToolCallRounds`). When `remaining` rounds are left (default 3),
   the call sends `text` as a user message; when one round is left and
   `finalWarning` is true (the default), it sends `finalWarningText`. In
   either text, `{remaining}` becomes the number of rounds left and
   `{limit}` the cap. Each warning is sent once per `llm()` call and is
   removed from the thread when the call returns. `remaining: 0` turns
   the first warning off. The default texts are "Of the {limit} tool
   rounds for this request, {remaining} remain. Finish what you can with
   them and then write your answer." and "This is the last tool round for
   this request. The results of the tool calls you make now are the last
   you will see. After them, write your answer from what you have." A
   call's own `warnAtRoundsLeft` combines with this one field by field.

   `baseUrl` sets the base URLs to use for the providers it names, for
   example `{ mlx: "http://studio.local:8090/v1" }` to reach an MLX
   server on another machine. A later `setLlmOptions` call with `baseUrl`
   replaces this one whole. A provider it leaves out uses the URL from
   `agency.json`, or its default.

   Setting a field to `null` removes it, so the value `agency.json` gives
   it, if any, applies again. To turn OpenRouter routing off entirely,
   pass `openRouter: {}`. */
export type LlmDefaults = {
  model?: string;
  provider?: string;
  temperature?: number;
  reasoningEffort?: "low" | "medium" | "high";
  thinking?: { enabled: boolean, budgetTokens?: number };
  replyLimits?: { hedgeLimit?: number, repeatLimit?: number, limitAnswers?: boolean };
  maxTokens?: number;
  maxToolResultChars?: number;
  maxToolCallRounds?: number;
  maxRepeatedToolCalls?: number;
  warnAtRoundsLeft?: WarnAtRoundsLeft;
  openRouter?: OpenRouterOptions;
  baseUrl?: LlmBaseUrls;
  rawAttributes?: any;
  retries?: number;
  timeout?: number;
  validationRetries?: number
}

export def setLlmOptions(opts: LlmDefaults) {
  """
  Set default options for subsequent llm() calls. Only the fields you
  pass are changed. A per-call llm(..., { ... }) option overrides them.

  @param opts - The default LLM options to merge in
  """
  _setLlmOptions(opts)
}

export def setModel(name: string) {
  """
  Set the default model for subsequent llm() calls.
  """
  _setLlmOptions({ model: name })
}

/**
Recognizes anthropic, google, openai, openrouter, and litellm. LiteLLM also
needs LITELLM_BASE_URL; this helper returns only its API-key variable name.
*/
export def envVarFor(provider: string): string {
  """
  Return a provider's API-key environment variable name, or an empty string if unknown.
  """
  if (provider == "anthropic") {
    return "ANTHROPIC_API_KEY"
  }
  if (provider == "google") {
    return "GEMINI_API_KEY"
  }
  if (provider == "openai") {
    return "OPENAI_API_KEY"
  }
  if (provider == "openrouter") {
    return "OPENROUTER_API_KEY"
  }
  if (provider == "litellm") {
    return "LITELLM_API_KEY"
  }
  return ""
}

export def pickProvider(order: string[] = ["anthropic", "google", "openai"]): Result<string> {
  """
  Return the first provider in `order` whose API-key environment variable
  is set, or a failure if none are. Checkable providers are anthropic,
  google, openai, openrouter, and litellm; unrecognized names in `order`
  are skipped.

  @param order - Providers to check, highest preference first
  """
  let checked = ""
  for (provider in order) {
    const varName = envVarFor(provider)
    if (varName != "") {
      if (checked == "") {
        checked = varName
      } else {
        checked = "${checked}, ${varName}"
      }
      const value = env(varName) with approve
      if (value != null && value != "") {
        return success(provider)
      }
    }
  }
  if (checked == "") {
    return failure("pickProvider: no recognized providers given (expected one of anthropic, google, openai, openrouter, litellm)")
  }
  return failure("No LLM API key found. Set one of: ${checked}")
}

// Loading a provider module runs its JavaScript in this process, so the
// load is gated like any other effect. The payload names the real path.
@always(path)
effect std::llm::registerProvider { path: string }

@capabilities(fileRead)
export def registerProviderModule(path: string): Result raises <std::llm::registerProvider> {
  """
  Load a provider module by path at runtime and register its custom provider
  for llm() calls. The module must export register({ registerProvider }).

  @param path - Path to the provider module (.mjs/.js)
  """
  const real = try _realTarget(path)
  if (isFailure(real)) {
    return real
  }
  return interrupt std::llm::registerProvider("Load this provider module? It runs as code in this process.", {
    path: real.value
  })
  return try _registerProviderModule(real.value)
}

// The TS HostedModelInfo in lib/stdlib/llm.ts mirrors this field order, and
// callers JSON-compare it in tests. Do not reorder.
export type HostedModelInfo = {
  name: string;
  provider: string;
  openWeights: boolean;
  inputCost: number;
  outputCost: number;
  contextWindow: number;
  family: string
}

export def listHostedModels(): HostedModelInfo[] {
  """
  Return all known hosted text models (the built-in catalog plus any
  refreshed data) for model discovery and pickers.
  """
  return _listHostedModels()
}

export def hostedModelInfo(name: string): HostedModelInfo | null {
  """
  Metadata for one hosted model by name, or null if the name is unknown or
  is not a text model.

  @param name - The hosted model name (e.g. "gpt-4o-mini")
  """
  return _hostedModelInfo(name)
}

/**
Use this to decide whether to offer an image or PDF to a model. A null
answer means the catalog has no information; it does not prove incompatibility.
*/
export def modelSupportsInput(model: string, modality: string): boolean | null {
  """
  Check catalog support for an input modality; null means unknown.

  @param modality - "image" or "pdf".
  """
  return _modelSupportsInput(model, modality)
}


// `std::read` is also declared in the prelude. A module that raises an
// effect declares it itself. The payload is `{ dir, filename }`.
@alwaysUnder(dir)
effect std::read {
  dir: string;
  filename: string
}

// Resolve before approval and refuse a link planted at the approved path.
/**
Use the JSON produced by `agency models refresh` to add or update model
metadata. Later files win on matching provider and model names, merging
individual fields. Hosted-tool metadata also merges. The returned count is
for this file, not the accumulated catalog.
*/
@capabilities(fileRead)
export def loadModelData(path: string): Result<number> raises <std::read> {
  """
  Load a model-data JSON file into the catalog and return its model count.

  @param path - File produced by agency models refresh, relative to the working directory or absolute.
  """
  const real = try _realTarget(path)
  if (isFailure(real)) {
    return real
  }
  return interrupt std::read("Are you sure you want to read this file?", {
    dir: dirname(real.value),
    filename: basename(real.value)
  })
  const res = _loadModelData(real.value)
  if (res.ok) {
    return success(res.count)
  }
  return failure(res.error)
}