// Talking to a LLeMbas instance: finding it, signing this machine in by device code (RFC 8628),
// reading the account's models and what else it offers, and signing out. The `webui` connection
// (webui.ts) uses the same calls at every start.
//
// The instance's address is all a person types. `/.well-known/lembas.json` says whether it is
// LLeMbas at all, which versions it speaks, and where to log in, so nobody is asked to approve
// anything for a server that only looks like one.
import { existsSync, readFileSync } from "node:fs"
import { hostname } from "node:os"
import { VERSION } from "../version.ts"

/** The device protocol this build speaks (LLeMbas `services/device_link.py:PROTOCOL`): 3, the
 *  same number as the link's `_lembas/*` extension (acp/agent.ts LEMBAS_PROTOCOL). */
export const PROTOCOL = 3
/** The ones it can log in to and link with: 3 alone. LLeMbas CLI 2.0 needs a LLeMbas 2.0 server
 *  (its tree, its memory, its models only); a 1.x server speaks 1 and 2, and is refused with a word
 *  on what to update (harness/acp.md *Versions*). */
export const PROTOCOLS = [3]

/** What this CLI says to a server that does not speak its protocol — at login, and at a dial: a
 *  1.x server (1 and 2) needs updating, a newer one needs a newer CLI. */
export function protocolMismatch(base: string, d: Pick<Discovery, "protocol" | "protocols" | "version">): string {
  const theirs = Array.isArray(d.protocols) && d.protocols.length ? d.protocols : [d.protocol]
  const said = `${base} runs LLeMbas ${d.version ?? "(an unknown version)"}, which speaks link protocol ${theirs.join(" and ")}`
  if (theirs.every((p) => p > PROTOCOL)) return `${said}; this LLeMbas CLI ${VERSION} speaks ${PROTOCOL}. Update LLeMbas CLI (lembas update).`
  return `${said}; this LLeMbas CLI ${VERSION} needs a LLeMbas 2.0 server (protocol ${PROTOCOL}). Update the server, or use LLeMbas CLI 1.x with it.`
}

export class LembasError extends Error {}
/** The server and this CLI speak no link protocol in common: no login, and no dial gets further. */
export class ProtocolError extends LembasError {}

export interface Discovery {
  service: "lembas"
  version: string
  /** The instance's own name (Example), where it gives one. */
  name?: string
  /** What a device calls it: its name as a connection name (example). */
  connection?: string
  protocol: number
  /** Every link protocol it speaks (`protocol: 1, protocols: [1, 2]` — `protocol` stays the
   *  lowest, so a client that compares it for equality still logs in). */
  protocols?: number[]
  harness_spec: string
  base_url: string
  api: { openai: string; instance?: string }
  login: { device: { code: string; token: string; verify: string } }
}

/** One entry of /v1/models: OpenAI's fields, then the harness spec's model metadata. */
export interface ServedModel {
  id: string
  name?: string
  family?: string
  context?: number
  max_output?: number
  temperature?: number
  top_p?: number
  efforts?: string[]
  effort?: string
  /** Spec 3.2.0: the server turns thinking off for enable_thinking false. */
  thinking_switch?: boolean
  vision?: boolean
  tools?: boolean
  notes?: string
  capacity?: { group?: string; single_session?: boolean; parallel?: number }
  default?: boolean
  /** The provider's slug — the part of `id` before its `/` — and what kind of model it is. */
  provider?: string
  /** "chat", "embedding", "stt", "tts" or "image"; absent from an older instance, which
   *  lists chat models only. */
  kind?: string
}

export interface Tls {
  ca?: string
}

/** Where a CA bundle that includes the system's own (and so an estate's internal CA, once it is
 *  installed system-wide) usually is. Bun trusts its built-in roots, not these, so a private CA
 *  needs naming — and this is the name to try before asking. */
export const SYSTEM_BUNDLES = ["/etc/ssl/certs/ca-certificates.crt", "/etc/pki/tls/certs/ca-bundle.crt", "/etc/ssl/cert.pem"]

function tlsInit(tls: Tls | undefined): { tls?: { ca: string } } {
  if (!tls?.ca) return {}
  try {
    return { tls: { ca: readFileSync(tls.ca, "utf8") } }
  } catch {
    throw new LembasError(`cannot read the CA file ${tls.ca}`)
  }
}

function untrusted(e: unknown): boolean {
  const code = String((e as { code?: string }).code ?? "")
  return /CERT|SIGNATURE|SELF_SIGNED|UNABLE_TO_GET_ISSUER/.test(code)
}

async function call(url: string, init: RequestInit, tls: Tls | undefined, timeoutMs = 15_000): Promise<Response> {
  try {
    return await fetch(url, { ...init, ...tlsInit(tls), signal: AbortSignal.timeout(timeoutMs) } as RequestInit)
  } catch (e) {
    if (untrusted(e)) throw new LembasError(`the TLS certificate of ${new URL(url).host} is not trusted. Pass --ca <file> with the CA that signed it.`)
    if ((e as Error).name === "TimeoutError") throw new LembasError(`${new URL(url).host} did not answer in time`)
    throw new LembasError(`cannot reach ${new URL(url).host}: ${(e as Error).message}`)
  }
}

/** What a person typed, as a base address: a scheme added (https unless they said http), the path
 *  and a trailing slash dropped. */
export function normalise(address: string): string {
  let a = address.trim()
  if (!a) throw new LembasError("an instance address is needed, e.g. ai.example.org")
  if (!/^[a-z][a-z0-9+.-]*:\/\//i.test(a)) a = `https://${a}`
  let u: URL
  try {
    u = new URL(a)
  } catch {
    throw new LembasError(`${address} is not an address`)
  }
  if (u.protocol !== "https:" && u.protocol !== "http:") throw new LembasError(`${address}: only http and https`)
  return `${u.protocol}//${u.host}`
}

/** Find the instance. With no CA given and an untrusted certificate, the system's own bundle is
 *  tried once — it is what an internal CA installed on this machine is in — and returned, so the
 *  connection written afterwards uses it too. */
/** The protocols both this CLI and the instance speak: `protocols` where it says, else `protocol`. */
export function spoken(d: Pick<Discovery, "protocol" | "protocols">): number[] {
  const theirs = Array.isArray(d.protocols) && d.protocols.length ? d.protocols : [d.protocol]
  return PROTOCOLS.filter((p) => theirs.includes(p))
}

/** The `_lembas` protocol to say in the link's hello: 3, the only one this build speaks. An
 *  instance that does not take it closes the link with 4400. Features stay gated by capabilities. */
export function linkProtocol(_protocols?: number[]): number {
  return PROTOCOL
}

export async function discover(address: string, tls?: Tls): Promise<{ discovery: Discovery; tls?: Tls }> {
  const base = normalise(address)
  const url = `${base}/.well-known/lembas.json`
  let res: Response
  let used = tls
  try {
    res = await call(url, { headers: { accept: "application/json" } }, tls)
  } catch (e) {
    const bundle = SYSTEM_BUNDLES.find((f) => existsSync(f))
    if (tls?.ca || !bundle || !(e instanceof LembasError) || !/not trusted/.test(e.message)) throw e
    used = { ca: bundle }
    res = await call(url, { headers: { accept: "application/json" } }, used)
  }
  if (res.status === 404) throw new LembasError(`${base} answers, but is not a LLeMbas instance (no /.well-known/lembas.json)`)
  if (!res.ok) throw new LembasError(`${base}: ${res.status} ${res.statusText}`)
  let body: Discovery
  try {
    body = (await res.json()) as Discovery
  } catch {
    throw new LembasError(`${base} answered /.well-known/lembas.json with something that is not JSON`)
  }
  if (body?.service !== "lembas") throw new LembasError(`${base} is not a LLeMbas instance`)
  if (!spoken(body).length) throw new ProtocolError(protocolMismatch(base, body))
  if (!body.login?.device?.code) throw new LembasError(`${base} offers no device login`)
  return { discovery: body, tls: used }
}

export interface DeviceStart {
  device_code: string
  user_code: string
  verification_uri: string
  verification_uri_complete?: string
  expires_in: number
  interval: number
}

export async function startDevice(d: Discovery, tls: Tls | undefined, scopes = ["models", "memory", "link"]): Promise<DeviceStart> {
  const res = await call(
    d.login.device.code,
    {
      method: "POST",
      headers: { "content-type": "application/json", accept: "application/json" },
      body: JSON.stringify({
        client_id: `lembas-cli ${VERSION}`,
        // The machine's name alone: the program is said by client_id, and the web UI shows this
        // where room is short.
        device_name: hostname(),
        device_host: hostname(),
        device_platform: `${process.platform}-${process.arch}`,
        scope: scopes.join(" "),
      }),
    },
    tls,
  )
  const body = (await res.json().catch(() => ({}))) as DeviceStart & { error?: string; error_description?: string }
  if (!res.ok) throw new LembasError(body.error_description ?? body.error ?? `${res.status} ${res.statusText}`)
  return body
}

export interface Granted {
  access_token: string
  scope: string
  account: { email?: string; name?: string }
  base_url?: string
}

/** Poll until the person decides, the code runs out, or `signal` stops it. */
export async function awaitToken(d: Discovery, start: DeviceStart, tls: Tls | undefined, opts: { signal?: AbortSignal; sleep?: (ms: number) => Promise<void> } = {}): Promise<Granted> {
  const sleep = opts.sleep ?? ((ms: number) => new Promise<void>((r) => setTimeout(r, ms)))
  let interval = Math.max(1, start.interval || 5)
  const deadline = Date.now() + Math.max(1, start.expires_in) * 1000
  while (Date.now() < deadline) {
    await sleep(interval * 1000)
    if (opts.signal?.aborted) throw new LembasError("cancelled")
    const res = await call(
      d.login.device.token,
      {
        method: "POST",
        headers: { "content-type": "application/json", accept: "application/json" },
        body: JSON.stringify({ grant_type: "urn:ietf:params:oauth:grant-type:device_code", device_code: start.device_code }),
      },
      tls,
    )
    const body = (await res.json().catch(() => ({}))) as Granted & { error?: string }
    if (res.ok && body.access_token) return body
    switch (body.error) {
      case "authorization_pending":
        continue
      case "slow_down":
        interval += 5
        continue
      case "access_denied":
        throw new LembasError("the sign-in was denied on the instance")
      case "expired_token":
        throw new LembasError("the code expired before it was approved — run login again")
      default:
        throw new LembasError(`the instance refused the sign-in: ${body.error ?? `${res.status} ${res.statusText}`}`)
    }
  }
  throw new LembasError("the code expired before it was approved — run login again")
}

/** The scopes a token carries, as the instance granted them (`models memory link`). */
export function scopesOf(scope: string | undefined): string[] {
  return (scope ?? "").split(/\s+/).filter(Boolean)
}

/** `lembas login --token-stdin`: prove a token handed over by another program — GET
 *  /api/v1/instance with it, which says the account and what the token may do (`scope`) — and
 *  return it as a device login's poll would have. A token that cannot link is refused: it would
 *  give this machine models but no agent chats, and the program handing it over is after the
 *  agent. Nothing is written by this. */
export async function checkToken(d: Discovery, token: string, tls: Tls | undefined): Promise<Granted> {
  const base = d.base_url
  const host = new URL(base).host
  const res = await call(`${base}/api/v1/instance`, { headers: { authorization: `Bearer ${token}`, accept: "application/json" } }, tls)
  const older = `${base} runs LLeMbas ${d.version}, which cannot say what a token may do: lembas login --token-stdin needs a newer LLeMbas 2.0 server. Update the server, or sign in with lembas login ${base}`
  if (res.status === 401) throw new LembasError(`${host} does not accept this token: it is revoked, expired, or not one of this instance's`)
  if (res.status === 403) throw new LembasError(`${host} refused this token for its models: it needs the models scope, and the account the “Use the API” permission`)
  if (res.status === 404) throw new LembasError(older)
  if (!res.ok) throw new LembasError(`checking the token: ${res.status} ${res.statusText}`)
  const body = (await res.json().catch(() => undefined)) as (InstanceInfo & { scope?: unknown; account?: { email?: string; name?: string } }) | undefined
  if (!body || typeof body !== "object" || typeof body.scope !== "string") throw new LembasError(older)
  const scopes = scopesOf(body.scope)
  if (!scopes.includes("link"))
    throw new LembasError(`this token cannot link this machine to ${host}: it carries ${scopes.length ? scopes.join(" ") : "no scope"}, and LLeMbas CLI's needs link (models memory link). Sign in with lembas login ${base} instead`)
  const account = body.account && typeof body.account === "object" ? { email: body.account.email, name: body.account.name } : {}
  return { access_token: token, scope: scopes.join(" "), account, base_url: base }
}

export async function models(d: { api: { openai: string } }, token: string, tls: Tls | undefined, timeoutMs?: number): Promise<ServedModel[]> {
  const res = await call(`${d.api.openai}/models`, { headers: { authorization: `Bearer ${token}`, accept: "application/json" } }, tls, timeoutMs)
  if (res.status === 401) throw new LembasError("the instance no longer accepts this sign-in — run login again")
  if (res.status === 403) throw new LembasError("this account may not use the instance's models through the API (the “Use the API” permission)")
  if (!res.ok) throw new LembasError(`listing the models: ${res.status} ${res.statusText}`)
  const body = (await res.json()) as { data?: ServedModel[] }
  return Array.isArray(body.data) ? body.data : []
}

/** GET /api/v1/instance: the instance by name, the account's default model, and
 *  which of voice and web search it may use through the API. */
export interface InstanceInfo {
  name?: string
  /** The account the token belongs to. */
  account?: { email?: string; name?: string }
  /** What the token presented may do, space-separated, as its login granted it. */
  scope?: string
  connection?: string
  version?: string
  default_model?: string | null
  services?: { stt?: boolean; tts?: boolean; search?: boolean; fetch?: boolean }
}

/** The instance's own description; undefined from an instance that does not have it (a 404). */
export async function instanceInfo(base: string, token: string, tls: Tls | undefined, timeoutMs?: number): Promise<InstanceInfo | undefined> {
  const res = await call(`${base}/api/v1/instance`, { headers: { authorization: `Bearer ${token}`, accept: "application/json" } }, tls, timeoutMs)
  if (res.status === 404) return undefined
  if (res.status === 401) throw new LembasError("the instance no longer accepts this sign-in — run login again")
  if (!res.ok) throw new LembasError(`reading the instance: ${res.status} ${res.statusText}`)
  const body = (await res.json().catch(() => undefined)) as InstanceInfo | undefined
  return body && typeof body === "object" ? body : undefined
}

/** Revoke this machine's token on the instance. Best-effort: signing out locally goes ahead even
 *  when the instance cannot be reached, and says so. */
/** DELETE /api/v1/token: the token revokes itself. `replacedBy`, a login again's new token: an
 *  instance that knows it hands the old device's agent chats to the new one first (one that does
 *  not ignores the body and only revokes). */
export async function signOut(baseUrl: string, token: string, tls: Tls | undefined, replacedBy?: string): Promise<boolean> {
  try {
    const headers: Record<string, string> = { authorization: `Bearer ${token}` }
    if (replacedBy) headers["content-type"] = "application/json"
    const body = replacedBy ? JSON.stringify({ replaced_by: replacedBy }) : undefined
    const res = await call(`${baseUrl}/api/v1/token`, { method: "DELETE", headers, body }, tls)
    return res.ok || res.status === 401
  } catch {
    return false
  }
}

/** GET /v1/usage, as the server's 2.0 answers it: counting and money, no credits. `plan`,
 *  `balance`, `balance_money`, the credit `cost` fields and `device.cost_money` are null and
 *  `windows` empty there; they are not read. */
export interface Usage {
  month: {
    from: string
    tokens_in: number
    tokens_out: number
    /** What the month cost in `currency` — null when it is in more than one currency. */
    cost_money?: number | null
    /** The month's cost per currency, `{ USD: 1.25, EUR: 0.4 }`: never added together. */
    cost_by_currency?: Record<string, number>
  }
  device: { tokens_in: number; tokens_out: number }
  /** The currency `month.cost_money` is in. */
  currency?: string
  /** An administrator of the instance. */
  admin?: boolean
  unit?: string
}

/** The usage, or undefined from an instance that does not have the endpoint. */
export async function usage(base: string, token: string, tls: Tls | undefined, timeoutMs = 5000): Promise<Usage | undefined> {
  const res = await call(`${base}/v1/usage`, { headers: { authorization: `Bearer ${token}`, accept: "application/json" } }, tls, timeoutMs)
  if (res.status === 404) return undefined
  if (res.status === 401) throw new LembasError("the instance no longer accepts this sign-in — run login again")
  if (!res.ok) throw new LembasError(`reading the usage: ${res.status} ${res.statusText}`)
  const body = (await res.json().catch(() => undefined)) as Usage | undefined
  return body && typeof body === "object" ? body : undefined
}

/** GET/PUT /v1/me/personalization: how the account wants to be helped, the same
 *  three things a person says about themselves here (personality, personality_custom, instructions).
 *  `available`: the administrator allows it; `enabled`: and the person has it on. */
export interface Personalization {
  enabled: boolean
  available: boolean
  personality: string
  personality_custom: string
  instructions: string
}

function asPersonalization(b: unknown): Personalization | undefined {
  if (!b || typeof b !== "object") return undefined
  const o = b as Record<string, unknown>
  const str = (v: unknown) => (typeof v === "string" ? v : "")
  return { enabled: o.enabled === true, available: o.available === true, personality: str(o.personality), personality_custom: str(o.personality_custom), instructions: str(o.instructions) }
}

/** The account's personalization, or undefined from an instance without the endpoint. */
export async function personalization(base: string, token: string, tls: Tls | undefined, timeoutMs = 5000): Promise<Personalization | undefined> {
  const res = await call(`${base}/v1/me/personalization`, { headers: { authorization: `Bearer ${token}`, accept: "application/json" } }, tls, timeoutMs)
  if (res.status === 404) return undefined
  if (res.status === 401) throw new LembasError("the instance no longer accepts this sign-in — run login again")
  if (!res.ok) throw new LembasError(`reading the personalization: ${res.status} ${res.statusText}`)
  return asPersonalization(await res.json().catch(() => undefined))
}

/** Write it back (the CLI's /settings, logged in). 403: not available, or a token that is not a
 *  device's; 422: a field of the wrong type. */
export async function savePersonalization(base: string, token: string, tls: Tls | undefined, value: Partial<Omit<Personalization, "available">>, timeoutMs = 10_000): Promise<Personalization | undefined> {
  const res = await call(
    `${base}/v1/me/personalization`,
    { method: "PUT", headers: { authorization: `Bearer ${token}`, accept: "application/json", "content-type": "application/json" }, body: JSON.stringify(value) },
    tls,
    timeoutMs,
  )
  // 403 is two things (not available, or not a device token): the instance's own words, where it
  // says which; 422 a field it would not take.
  if (res.status === 403 || res.status === 422) {
    const detail = await res
      .json()
      .then((b: any) => (typeof b?.detail === "string" ? b.detail : typeof b?.error === "string" ? b.error : ""))
      .catch(() => "")
    if (res.status === 422) throw new LembasError(`the instance would not take that personalization${detail ? `: ${detail}` : ""}`)
    throw new LembasError(detail ? `the instance did not allow it: ${detail}` : "not allowed: the instance refused to save the personalization")
  }
  if (res.status === 401) throw new LembasError("the instance no longer accepts this sign-in — run login again")
  if (!res.ok) throw new LembasError(`saving the personalization: ${res.status} ${res.statusText}`)
  return asPersonalization(await res.json().catch(() => undefined))
}