import type { BackendContext } from '@jaspers-ai/sdk'
import { research, type JobStatus, type ServerStatus, type StoredCompany } from './research.ts'
import { companyOf, isCurrent, readsText, tickerOf, type Category, type Reads } from './sec.ts'

// Reading a company's filings into the store: the store does it, as a job on its worker, for
// everyone. add_company names the kinds of filing and the years; the job lists them at SEC, reads
// what the store lacks, and the embedding queue follows. This side starts the job and waits on it,
// saying how far it is: first until it is read, then until nothing of the company is left to embed.

/** How long between looks at a job or a company's embedding. */
const POLL_MS = 2_000
const OPEN = new Set(['queued', 'listing', 'reading'])

/** Waits a poll, or throws once the answer is stopped. */
async function pause(signal: AbortSignal, ms: number): Promise<void> {
  if (signal.aborted) throw signal.reason ?? new Error('Stopped.')
  await new Promise((resolve) => setTimeout(resolve, ms))
}

/** The named companies as the store holds them. */
export async function companies(
  ctx: BackendContext,
  tickers: string[],
  signal?: AbortSignal,
): Promise<StoredCompany[]> {
  if (tickers.length === 0) return []
  const status = await research<ServerStatus>(ctx, 'status', { tickers }, signal)
  return status.companies
}

/** Waits while a job is open, saying how far it is; the job as it ended. Throws with the store's words when it failed. */
async function untilRead(
  ctx: BackendContext,
  job: JobStatus,
  ticker: string,
  signal: AbortSignal,
  said: (step: string) => Promise<void>,
  pollMs: number,
): Promise<JobStatus> {
  let current = job
  while (OPEN.has(current.state)) {
    await said(
      current.state === 'reading'
        ? `Reading ${ticker} filings: ${current.done.toLocaleString('en-US')} of ${current.total.toLocaleString('en-US')}`
        : current.state === 'listing'
          ? `Listing ${ticker} filings`
          : `Waiting to read ${ticker}`,
    )
    await pause(signal, pollMs)
    current = await research<JobStatus>(ctx, 'job', { id: job.id }, signal)
  }
  if (current.state === 'failed') throw new Error(current.error ?? `${ticker} could not be read.`)
  if (current.state === 'unknown') throw new Error(`The store lost the job reading ${ticker}.`)
  return current
}

/** Waits until none of a company's chunks wait for an embedding, saying how many are left; the company as it stands. */
async function untilEmbedded(
  ctx: BackendContext,
  ticker: string,
  signal: AbortSignal,
  said: (step: string) => Promise<void>,
  pollMs: number,
): Promise<StoredCompany | null> {
  for (;;) {
    const [held] = await companies(ctx, [ticker], signal)
    if (!held || held.unembedded === 0) return held ?? null
    await said(`Embedding ${ticker}: ${held.unembedded.toLocaleString('en-US')} chunks to go`)
    await pause(signal, pollMs)
  }
}

/**
 * What an agent's add_company does: has the store read the kinds of a company's filings for the
 * years, and answers once they are read and embedded, so a search finds them by meaning as well as
 * by their words. A company another account read in is already there: the job ends at once.
 */
export async function readCompany(
  ctx: BackendContext,
  rawTicker: string,
  years: number[],
  categories: Category[],
  signal: AbortSignal,
  said: (step: string) => Promise<void>,
  pollMs: number = POLL_MS,
): Promise<string> {
  const { ticker, cik } = companyOf(rawTicker)
  const reads: Reads = Object.fromEntries(categories.map((c) => [c, years]))
  const started = await research<{ job: JobStatus }>(ctx, 'add_company', { ticker, cik, reads }, signal)
  const job = await untilRead(ctx, started.job, ticker, signal, said, pollMs)
  const held = await untilEmbedded(ctx, ticker, signal, said, pollMs)
  const failed = job.failed ? `, and ${job.failed} could not be read` : ''
  return `${held?.company ?? ticker}: ${job.done} of the ${job.total} filings, ${readsText(reads)}, are in the store and embedded${failed}. Search them now.`
}

/**
 * When this session last checked each company, ISO. The store's recently_as_of is the record, but a
 * check that fails (SEC out of reach) keeps no time there: without this every search would check again.
 */
const checked = new Map<string, string>()

/**
 * Before a conversation searches a company the store holds: when SEC last listed its filings more
 * than CURRENT_FOR_MS ago, has the store list them again and read what is new in the kinds it keeps
 * current, and waits until that is embedded. A company the store lacks is left alone: the agent adds it.
 */
export async function keepCurrent(
  ctx: BackendContext,
  rawTicker: string,
  signal: AbortSignal,
  said: (step: string) => Promise<void>,
  pollMs: number = POLL_MS,
): Promise<void> {
  const { ticker } = companyOf(rawTicker)
  const [held] = await companies(ctx, [ticker], signal)
  const now = new Date()
  if (!held || isCurrent(held.recently_as_of, now) || isCurrent(checked.get(ticker) ?? null, now)) return
  checked.set(ticker, now.toISOString())
  if (held.job && OPEN.has(held.job.state)) {
    await untilRead(ctx, held.job, ticker, signal, said, pollMs)
    await untilEmbedded(ctx, ticker, signal, said, pollMs)
    return
  }
  await said(`Checking ${ticker} for new filings`)
  const started = await research<{ job: JobStatus }>(ctx, 'refresh_company', { ticker }, signal)
  await untilRead(ctx, started.job, ticker, signal, said, pollMs)
  await untilEmbedded(ctx, ticker, signal, said, pollMs)
}

/** The chip's Refresh: what the company has filed since the store last listed it, in the kinds it keeps current. */
export async function refresh(ctx: BackendContext, rawTicker: string): Promise<string> {
  const ticker = tickerOf(rawTicker)
  await research<{ job: JobStatus }>(ctx, 'refresh_company', { ticker })
  return `Checking ${ticker}.`
}