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

/** @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)
}
```
*/

/** 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). */
export type LlmDefaults = {
  model?: string;
  provider?: string;
  temperature?: number;
  reasoningEffort?: "low" | "medium" | "high";
  maxTokens?: number;
  maxToolResultChars?: number;
  maxToolCallRounds?: number;
  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.

  @param name - The model name, e.g. "gpt-4o-mini" or "claude-opus-4-8"
  """
  _setLlmOptions({ model: name })
}

export def envVarFor(provider: string): string {
  """
  Return the environment variable that holds the API key for a recognized
  provider, or "" for an unrecognized name. Recognized: "anthropic"
  (ANTHROPIC_API_KEY), "google" (GEMINI_API_KEY), "openai" (OPENAI_API_KEY),
  "openrouter" (OPENROUTER_API_KEY), "litellm" (LITELLM_API_KEY). Note that
  "litellm" also requires a base URL (LITELLM_BASE_URL). This returns only
  the API-key var.

  @param provider - The provider name to map
  """
  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, 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)
      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}")
}

export def registerProviderModule(path: string) {
  """
  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)
  """
  _registerProviderModule(path)
}

// 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)
}

export def modelSupportsInput(model: string, modality: string): boolean | null {
  """
  Whether a model accepts a given input modality ("image" or "pdf").
  Tri-state: true / false when the model catalog says so, null when the
  model or its modality data is unknown. Treat null as "do not gate",
  the same rule llm() applies at send time.

  @param model - The model name (e.g. "gpt-4o-mini")
  @param modality - "image" or "pdf"
  """
  return _modelSupportsInput(model, modality)
}


// Accumulation details (kept out of the docstring, which is also the LLM tool
// description): when called more than once, the newly loaded file wins over
// previously loaded files and over the built-in catalog on provider+name
// collisions, deep-merging fields so a partial entry augments an existing one
// instead of replacing it; `hostedTools` are merged too. The returned count is
// the number of models in this file, not the running total registered. `path`
// is resolved relative to the working directory, or used as-is if absolute.
export def loadModelData(path: string): Result<number, string> {
  """
  Load model data from a JSON file (the shape `agency models refresh`
  prints) and register it so `llm()` and the model catalog recognize
  those models. Multiple loads accumulate. Returns the number of models
  loaded, or a failure if the file cannot be read.

  @param path - Path to a model-data JSON file
  """
  const res = _loadModelData(path)
  if (res.ok) {
    return success(res.count)
  }
  return failure(res.error)
}