// The app admits only the SDK's root to a plugin, so the skill helpers come from there.
import { ACTIVATE_SKILL, runSkillTool, skillsPrompt, skillTools, type Skills } from '@jaspers-ai/sdk'
import type { Agent } from './agents.ts'
import {
  calculationId,
  passagesText,
  systemPrompt,
  type AgentMessage,
  type Holding,
  type Passage,
  type Search,
  type SkillUse,
} from './chat.ts'
import { CATEGORIES, CATEGORY_NAMES, readsText, tickerOf, type Category } from './sec.ts'
import { stepOf, thoughtOf } from './working.ts'
import { repairText, uncited, unknownRefs } from './cite.ts'
import { errorText } from './errors.ts'
import type { Hit } from './research.ts'
import { askedOf, mergeSearch } from './search.ts'
import { calculate } from './calc.ts'

// One agent answering one message: call the model with the conversation and its tools, search_prism
// and calculate and, for a web agent, search_web and fetch_page, run its searches through the agent's
// pinned filters, and keep
// going until it answers without searching, however many searches that takes. An agent given skills
// sees their names and descriptions and loads one's instructions, or reads one's files, when a
// question calls for it, through the app's skills. Each search hands the
// model at most three passages, and only ones the store's reranker judged relevant. The model and
// the store arrive through the context, so the loop is tested with a scripted model.

export interface ToolDefinition {
  name: string
  description: string
  parameters: Record<string, unknown>
}

export interface ToolCall {
  id: string
  name: string
  input: Record<string, unknown>
}

export type Turn =
  | { role: 'user'; text: string }
  | { role: 'assistant'; text: string; toolCalls: ToolCall[]; raw: unknown[] | null }
  | { role: 'tool'; results: { callId: string; output: string; isError: boolean }[] }

export interface Completion {
  text: string
  toolCalls: ToolCall[]
  usage: { input: number; output: number }
  raw: unknown[] | null
}

/** What a run reaches: the model, the store's search, a page read into the store, and the app's skills. */
export interface RunnerContext {
  llm: {
    complete(req: {
      system: string
      turns: Turn[]
      tools?: ToolDefinition[]
      model?: string
      signal?: AbortSignal
    }): Promise<Completion>
  }
  search(args: Record<string, unknown>, signal: AbortSignal): Promise<Hit[]>
  /** Reads a public page into the store and says the address it ended at and the document it is. */
  fetchPage(url: string, signal: AbortSignal): Promise<{ url: string; doc: string }>
  /** The companies the store holds and the years of their filings. */
  holdings(signal: AbortSignal): Promise<Holding[]>
  /**
   * Reads the kinds of a company's filings for the years into the store and waits until they are embedded, saying
   * how far it is as it goes; answers what was read, and throws when it cannot.
   */
  addCompany(
    ticker: string,
    years: number[],
    categories: Category[],
    signal: AbortSignal,
    progress: (step: string) => Promise<void>,
  ): Promise<string>
  /**
   * Before a search names a company the store holds: when SEC last listed its filings more than 15
   * minutes ago, reads in what it filed since and waits until it is embedded, saying how far it is.
   */
  keepCurrent(ticker: string, signal: AbortSignal, progress: (step: string) => Promise<void>): Promise<void>
  /** The web searched with the user's search provider; throws when none is set. */
  searchWeb(query: string, signal: AbortSignal): Promise<{ title: string; url: string; snippet: string }[]>
  skills: Skills
}

const SEARCH_PRISM: ToolDefinition = {
  name: 'search_prism',
  description:
    'Search the research store by meaning and keywords within the filters, reranked for relevance. Returns at most 3 passages, each labeled with where it is from and a ref to cite; none when nothing is relevant enough.',
  parameters: {
    type: 'object',
    properties: {
      query: { type: 'string', description: 'What to find, in a few specific words.' },
      ticker: { type: 'string', description: 'One company, like FLWS.' },
      document: {
        type: 'string',
        description: 'One document by its id, like doc:notes.pdf: a file the user attached, or a page read.',
      },
      form: {
        type: 'string',
        description: 'One form, like 10-K, 10-Q, 8-K, DEF 14A, or 4; its amendments are included.',
      },
      fiscal_year: { type: 'integer', description: "The company's fiscal year, as its 10-K or 10-Q says it." },
      fiscal_period: { type: 'string', description: 'FY for a 10-K, Q1, Q2, or Q3 for a 10-Q.' },
      section: {
        type: 'string',
        description:
          "A section by its item or name: Item 1A or Risk Factors, Item 7 or Management's Discussion and Analysis, Item 2.02, EX-99.1.",
      },
      filed_after: { type: 'string', description: 'YYYY-MM-DD, filed on or after.' },
      filed_before: { type: 'string', description: 'YYYY-MM-DD, filed before.' },
    },
    required: ['query'],
  },
}

const SEARCH_WEB: ToolDefinition = {
  name: 'search_web',
  description:
    'Search the web for pages: their titles, addresses, and a line from each. Read one with fetch_page to cite it; a search result alone is not a passage.',
  parameters: {
    type: 'object',
    properties: { query: { type: 'string', description: 'What to find, in a few specific words.' } },
    required: ['query'],
  },
}

const ADD_COMPANY: ToolDefinition = {
  name: 'add_company',
  description:
    "Read the kinds of a public company's SEC filings a question needs, filed in the years from to to, into the store. Read only what the question needs (risk factors are reports, not insiders' trades); another kind or year can be added when a later question needs it. Takes a minute or more; the filings are searchable once it answers.",
  parameters: {
    type: 'object',
    properties: {
      ticker: {
        type: 'string',
        description: 'Its ticker, like GT; for one SEC no longer lists by ticker, TICKER:CIK, like DFS:1393612.',
      },
      from: { type: 'integer', description: 'The first year filed, like 2024.' },
      to: { type: 'integer', description: 'The last year filed, like 2025.' },
      categories: {
        type: 'array',
        items: { type: 'string', enum: CATEGORY_NAMES },
        description: `The kinds of filing to read, each with its amendments: ${CATEGORY_NAMES.map((c) => `${c}, ${CATEGORIES[c].holds} (${CATEGORIES[c].forms.join(', ')})`).join('; ')}.`,
      },
    },
    required: ['ticker', 'from', 'to', 'categories'],
  },
}

const CALCULATE: ToolDefinition = {
  name: 'calculate',
  description:
    'Work out arithmetic exactly. One line each, `name = expression`, a line using the names above it. Numbers without commas, + - * / ^, parentheses, and abs, sqrt, ln, log10, exp, round(x, decimals), min, max, sum, avg, npv(rate, cf0, cf1, ...) with cf0 at time zero, irr(cf0, cf1, ...). Answers a passage to cite for the figures it works out.',
  parameters: {
    type: 'object',
    properties: {
      expressions: {
        type: 'string',
        description: 'Lines like `rev = 2994105`, `cost = 1427546 - 67032`, `margin = (rev - cost) / rev`.',
      },
    },
    required: ['expressions'],
  },
}

const FETCH_PAGE: ToolDefinition = {
  name: 'fetch_page',
  description:
    'Read a public web page into the research store and get the 3 passages of it most relevant to query. The page is kept, so search_prism finds it later.',
  parameters: {
    type: 'object',
    properties: {
      url: { type: 'string', description: 'A full https address.' },
      query: { type: 'string', description: 'What to find on the page, in a few specific words.' },
    },
    required: ['url', 'query'],
  },
}

export async function runAgent(
  ctx: RunnerContext,
  agent: Agent,
  message: AgentMessage,
  firstTurn: string,
  save: (message: AgentMessage) => Promise<void>,
  signal: AbortSignal,
): Promise<AgentMessage> {
  const turns: Turn[] = [{ role: 'user', text: firstTurn }]
  let current: AgentMessage = { ...message, status: 'working' }
  let refs = 0
  /** The answer goes back once to cite what it did not; what it still does not is said under it. */
  let repaired = false
  const nextRef = (): string => `p${++refs}`
  /** The model's latest note, kept on the working line until it writes another. */
  let thought: string | null = null
  const working = async (step: string, company?: string): Promise<void> => {
    current = { ...current, working: { step, thought, ...(company ? { company } : {}) } }
    await save(current)
  }
  try {
    // Only the skills it was given and the app has: one the user removed since is left out.
    const skills = agent.skills.length > 0 ? await ctx.skills.catalog({ names: agent.skills }) : []
    const system = [systemPrompt(agent, await ctx.holdings(signal)), skillsPrompt(skills)].filter(Boolean).join('\n\n')
    const offered = [
      ...(agent.web ? [SEARCH_PRISM, FETCH_PAGE, SEARCH_WEB] : [SEARCH_PRISM]),
      ADD_COMPANY,
      CALCULATE,
      ...skillTools(skills.map((s) => s.name)),
    ]
    for (;;) {
      await working(repaired ? 'Adding citations…' : 'Thinking…')
      const reply = await ctx.llm.complete({ system, turns, tools: offered, model: agent.model ?? undefined, signal })
      current = {
        ...current,
        usage: { input: current.usage.input + reply.usage.input, output: current.usage.output + reply.usage.output },
      }
      turns.push({ role: 'assistant', text: reply.text, toolCalls: reply.toolCalls, raw: reply.raw })
      if (reply.toolCalls.length === 0) {
        const text = reply.text.trim()
        const given = current.searches.flatMap((s) => s.passages.map((p) => p.ref))
        const missing = uncited(text)
        const unknown = unknownRefs(text, new Set(given))
        if (!repaired && (missing.length > 0 || unknown.length > 0)) {
          repaired = true
          turns.push({ role: 'user', text: repairText(missing, unknown, given) })
          current = { ...current, text }
          await save(current)
          continue
        }
        const audit = unknown.length > 0 ? { unknown } : null
        current = { ...current, status: 'done', text, audit }
        break
      }
      thought = thoughtOf(reply.text, reply.raw) ?? thought
      const results: { callId: string; output: string; isError: boolean }[] = []
      for (const call of reply.toolCalls) {
        await working(stepOf(call))
        if (!offered.some((tool) => tool.name === call.name)) {
          results.push({
            callId: call.id,
            output: `There is no tool ${call.name}; the tools are ${offered.map((tool) => tool.name).join(' and ')}.`,
            isError: true,
          })
          continue
        }
        const skill = await runSkillTool(ctx.skills, call, turns)
        if (skill) {
          const use: SkillUse = {
            skill: String(call.input[call.name === ACTIVATE_SKILL ? 'name' : 'skill'] ?? ''),
            file: call.name === ACTIVATE_SKILL ? null : String(call.input['path'] ?? ''),
            error: skill.isError ? skill.output : null,
          }
          current = { ...current, skills: [...(current.skills ?? []), use] }
          results.push(skill)
          continue
        }
        const query = typeof call.input['query'] === 'string' ? call.input['query'].trim() : ''
        const { search, output } =
          call.name === CALCULATE.name
            ? runCalculation(String(call.input['expressions'] ?? ''), nextRef)
            : call.name === ADD_COMPANY.name
              ? await runAddCompany(ctx, call.input, signal, working)
              : call.name === SEARCH_WEB.name
                ? await runWebSearch(ctx, query, signal)
                : withPassages(
                    call.name === FETCH_PAGE.name
                      ? await runFetch(ctx, agent, String(call.input['url'] ?? ''), query, nextRef, signal)
                      : await runSearch(ctx, agent, query, call.input, nextRef, signal, working),
                  )
        current = { ...current, searches: [...current.searches, search] }
        results.push({ callId: call.id, output: search.error ?? output, isError: search.error !== null })
      }
      turns.push({ role: 'tool', results })
      await save(current)
    }
  } catch (err) {
    current = signal.aborted
      ? { ...current, status: 'stopped', error: errorText(signal.reason, 'Stopped.') }
      : { ...current, status: 'error', error: errorText(err, 'failed') }
  }
  current = { ...current, working: null }
  await save(current)
  return current
}

/** A search of the store or a page, as the model reads it: its passages. */
function withPassages(search: Search): { search: Search; output: string } {
  return { search, output: passagesText(search.passages) }
}

/** The lines worked out, as one passage to cite; a calculation with no value at all is an error. */
function runCalculation(expressions: string, nextRef: () => string): { search: Search; output: string } {
  const worked = calculate(expressions)
  const query = expressions.trim()
  if (worked.lines.every((l) => l.error !== null)) {
    return { search: { query, tool: 'calculate', passages: [], error: worked.text }, output: worked.text }
  }
  const ref = nextRef()
  const passage: Passage = { ref, id: calculationId(ref), label: 'Calculation', url: null, score: 1, text: worked.text }
  return { search: { query, tool: 'calculate', passages: [passage], error: null }, output: passagesText([passage]) }
}

/** A company's filings read into the store and embedded for the years asked, before the model searches them. */
async function runAddCompany(
  ctx: RunnerContext,
  input: Record<string, unknown>,
  signal: AbortSignal,
  progress: (step: string, company?: string) => Promise<void>,
): Promise<{ search: Search; output: string }> {
  const ticker = String(input['ticker'] ?? '').trim()
  const from = Number(input['from'])
  const to = Number(input['to'])
  const asked = Array.isArray(input['categories']) ? input['categories'].map(String) : []
  const categories = CATEGORY_NAMES.filter((c) => asked.includes(c))
  const years =
    Number.isInteger(from) && Number.isInteger(to) && from <= to
      ? Array.from({ length: to - from + 1 }, (_, i) => from + i)
      : []
  const query = `${ticker} ${readsText(Object.fromEntries(categories.map((c) => [c, years]))) || `${from}–${to}`}`
  let named: string | undefined
  const failed = (error: string) => ({
    search: { query, tool: 'add_company' as const, ...(named ? { ticker: named } : {}), passages: [], error },
    output: '',
  })
  try {
    named = tickerOf(ticker)
  } catch (err) {
    return failed(errorText(err, 'that is not a ticker'))
  }
  if (years.length === 0) return failed('from and to are years, from first.')
  const unknown = asked.filter((c) => !(CATEGORY_NAMES as string[]).includes(c))
  if (categories.length === 0 || unknown.length > 0) {
    return failed(`categories are one or more of ${CATEGORY_NAMES.join(', ')}.`)
  }
  try {
    // The working line names the company while it is read in, so the view lists it from the start.
    const company = named
    await progress(`Reading ${company} filings`, company)
    const said = await ctx.addCompany(ticker, years, categories, signal, (step) => progress(step, company))
    return { search: { query, tool: 'add_company', ticker: named, passages: [], error: null }, output: said }
  } catch (err) {
    return failed(errorText(err, 'the filings could not be read'))
  }
}

/** Pages to read, not passages: the model fetches one to cite it. */
async function runWebSearch(
  ctx: RunnerContext,
  query: string,
  signal: AbortSignal,
): Promise<{ search: Search; output: string }> {
  if (!query) return { search: { query, tool: 'search_web', passages: [], error: 'query is empty.' }, output: '' }
  try {
    const found = await ctx.searchWeb(query, signal)
    const output =
      found.length === 0
        ? 'The web search found nothing. Change the words.'
        : found.map((r) => `${r.title}\n${r.url}\n${r.snippet}`).join('\n\n')
    const results = found.map((r) => ({ title: r.title, url: r.url }))
    return { search: { query, tool: 'search_web', results, passages: [], error: null }, output }
  } catch (err) {
    return {
      search: { query, tool: 'search_web', passages: [], error: errorText(err, 'the web search failed') },
      output: '',
    }
  }
}

/** A page read in, then searched alone: the agent's other holds are about what it had, not this page. */
async function runFetch(
  ctx: RunnerContext,
  agent: Agent,
  url: string,
  query: string,
  nextRef: () => string,
  signal: AbortSignal,
): Promise<Search> {
  if (!query) return { query, url, passages: [], error: 'query is empty.' }
  let page: { url: string; doc: string }
  try {
    page = await ctx.fetchPage(url.trim(), signal)
  } catch (err) {
    return { query, url, passages: [], error: errorText(err, 'the page could not be read') }
  }
  const search = await held(ctx, { query, documents: [page.doc] }, query, nextRef, signal)
  return { ...search, url: page.url }
}

async function runSearch(
  ctx: RunnerContext,
  agent: Agent,
  query: string,
  input: Record<string, unknown>,
  nextRef: () => string,
  signal: AbortSignal,
  progress: (step: string, company?: string) => Promise<void>,
): Promise<Search> {
  if (!query) return { query, passages: [], error: 'query is empty.' }
  const merged = mergeSearch(query, agent.search, askedOf(input))
  if ('refused' in merged) return { query, passages: [], error: merged.refused }
  const tickers = merged.args['tickers']
  for (const ticker of Array.isArray(tickers) ? tickers.map(String) : []) {
    try {
      await ctx.keepCurrent(ticker, signal, (step) => progress(step, ticker))
    } catch (err) {
      if (signal.aborted) throw err
      // SEC out of reach is no reason to refuse a search of what the store already holds.
      console.warn(`${ticker}: ${errorText(err, 'could not be checked for new filings')}`)
    }
  }
  return held(ctx, merged.args, query, nextRef, signal)
}

/** One search of the store, its hits as passages with refs that run on across the answer. */
async function held(
  ctx: RunnerContext,
  args: Record<string, unknown>,
  query: string,
  nextRef: () => string,
  signal: AbortSignal,
): Promise<Search> {
  const { query: _words, ...filters } = args
  try {
    const hits = await ctx.search(args, signal)
    const passages: Passage[] = hits.map((hit) => ({
      ref: nextRef(),
      id: hit.id,
      label: hit.label,
      url: hit.url,
      score: hit.score,
      text: hit.text,
    }))
    return { query, filters, passages, error: null }
  } catch (err) {
    return { query, filters, passages: [], error: errorText(err, 'search failed') }
  }
}