import { defineSource } from '@jaspers-ai/sdk'
import { z } from 'zod'
import { HOSTS, baseField, flatten, get, list, upper } from './client'

// Reference data: tickers, one ticker's details, dividends, splits, exchanges, conditions, market hours.

const limit = (d: number) => z.number().int().min(1).max(10000).default(d).describe('Most rows to answer.')
const page = (max: number) => Math.min(max, 1000)
const opt = (s?: string) => (s ? upper(s) : undefined)

const tickers = defineSource({
  description: 'Search or list tickers: ticker, name, market, type, primary exchange, active, currency, CIK, FIGI.',
  hosts: HOSTS,
  input: z.object({
    search: z.string().optional().describe('Words in the ticker or company name, like apple.'),
    ticker: z.string().optional().describe('One exact ticker.'),
    market: z.enum(['stocks', 'crypto', 'fx', 'otc', 'indices']).optional().describe('The market.'),
    type: z.string().optional().describe('A ticker type like CS (common stock), ETF, ADRC, INDEX.'),
    active: z.boolean().default(true).describe('Only tickers trading now.'),
    limit: limit(100),
    base: baseField,
  }),
  async run(args, ctx) {
    const res = await list(
      ctx,
      args.base,
      '/v3/reference/tickers',
      { search: args.search, ticker: opt(args.ticker), market: args.market, type: args.type, active: args.active, limit: page(args.limit) },
      args.limit,
    )
    const rows = res.rows.map((r) => ({
      ticker: r.ticker,
      name: r.name ?? null,
      market: r.market ?? null,
      locale: r.locale ?? null,
      type: r.type ?? null,
      primary_exchange: r.primary_exchange ?? null,
      active: r.active ?? null,
      currency: r.currency_name ?? null,
      cik: r.cik ?? null,
      composite_figi: r.composite_figi ?? null,
      last_updated: r.last_updated_utc ?? null,
    }))
    return { rows, more: res.more }
  },
})

const tickerDetails = defineSource({
  description: 'One ticker overview: name, description, market cap, employees, list date, SIC, shares outstanding, address, homepage, logo.',
  hosts: HOSTS,
  input: z.object({
    ticker: z.string().min(1).describe('The ticker, like AAPL.'),
    date: z.string().optional().describe('As of this date, YYYY-MM-DD. Leave out for now.'),
    base: baseField,
  }),
  async run(args, ctx) {
    const body = await get(ctx, args.base, `/v3/reference/tickers/${upper(args.ticker)}`, { date: args.date })
    return body?.results ? [flatten(body.results)] : []
  },
})

const dividends = defineSource({
  description: 'Cash dividends, newest ex-date first: ticker, cash amount, currency, declaration, ex-dividend, record, pay dates, frequency, type.',
  hosts: HOSTS,
  input: z.object({
    ticker: z.string().optional().describe('A ticker, like AAPL. Leave out for all.'),
    from: z.string().optional().describe('Earliest ex-dividend date, YYYY-MM-DD.'),
    to: z.string().optional().describe('Latest ex-dividend date, YYYY-MM-DD.'),
    limit: limit(100),
    base: baseField,
  }),
  async run(args, ctx) {
    const res = await list(
      ctx,
      args.base,
      '/v3/reference/dividends',
      {
        ticker: opt(args.ticker),
        'ex_dividend_date.gte': args.from,
        'ex_dividend_date.lte': args.to,
        sort: 'ex_dividend_date',
        order: 'desc',
        limit: page(args.limit),
      },
      args.limit,
    )
    const rows = res.rows.map((r) => ({
      ticker: r.ticker,
      ex_dividend_date: r.ex_dividend_date ?? null,
      cash_amount: r.cash_amount ?? null,
      currency: r.currency ?? null,
      declaration_date: r.declaration_date ?? null,
      record_date: r.record_date ?? null,
      pay_date: r.pay_date ?? null,
      frequency: r.frequency ?? null,
      dividend_type: r.dividend_type ?? null,
    }))
    return { rows, more: res.more }
  },
})

const splits = defineSource({
  description: 'Stock splits, newest first: ticker, execution date, split from, split to, ratio.',
  hosts: HOSTS,
  input: z.object({
    ticker: z.string().optional().describe('A ticker, like AAPL. Leave out for all.'),
    from: z.string().optional().describe('Earliest execution date, YYYY-MM-DD.'),
    to: z.string().optional().describe('Latest execution date, YYYY-MM-DD.'),
    limit: limit(100),
    base: baseField,
  }),
  async run(args, ctx) {
    const res = await list(
      ctx,
      args.base,
      '/v3/reference/splits',
      {
        ticker: opt(args.ticker),
        'execution_date.gte': args.from,
        'execution_date.lte': args.to,
        sort: 'execution_date',
        order: 'desc',
        limit: page(args.limit),
      },
      args.limit,
    )
    const rows = res.rows.map((r) => ({
      ticker: r.ticker,
      execution_date: r.execution_date ?? null,
      split_from: r.split_from ?? null,
      split_to: r.split_to ?? null,
      ratio: r.split_from ? r.split_to / r.split_from : null,
    }))
    return { rows, more: res.more }
  },
})

const exchanges = defineSource({
  description: 'Exchanges Massive knows: id, name, acronym, MIC, type, asset class, locale, participant id.',
  hosts: HOSTS,
  input: z.object({
    asset_class: z.enum(['stocks', 'options', 'crypto', 'fx']).optional().describe('Only this asset class.'),
    base: baseField,
  }),
  async run(args, ctx) {
    const body = await get(ctx, args.base, '/v3/reference/exchanges', { asset_class: args.asset_class })
    return ((body?.results ?? []) as any[]).map((r) => flatten(r))
  },
})

const conditions = defineSource({
  description: 'Trade and quote condition codes: id, name, type, asset class, data types, SIP mapping, description.',
  hosts: HOSTS,
  input: z.object({
    asset_class: z.enum(['stocks', 'options', 'crypto', 'fx']).optional().describe('Only this asset class.'),
    data_type: z.enum(['trade', 'bbo', 'nbbo']).optional().describe('Only conditions for this data type.'),
    limit: limit(1000),
    base: baseField,
  }),
  async run(args, ctx) {
    const res = await list(
      ctx,
      args.base,
      '/v3/reference/conditions',
      { asset_class: args.asset_class, data_type: args.data_type, limit: page(args.limit) },
      args.limit,
    )
    return { rows: res.rows.map((r) => flatten(r)), more: res.more }
  },
})

const marketStatus = defineSource({
  description: 'Whether US markets are open now: market, early hours, after hours, each exchange, currencies, index groups, server time. One row.',
  hosts: HOSTS,
  input: z.object({ base: baseField }),
  async run(args, ctx) {
    return [flatten(await get(ctx, args.base, '/v1/marketstatus/now'))]
  },
})

const marketHolidays = defineSource({
  description: 'Upcoming market holidays and early closes: date, exchange, name, status, open and close times.',
  hosts: HOSTS,
  input: z.object({ base: baseField }),
  async run(args, ctx) {
    const body = await get(ctx, args.base, '/v1/marketstatus/upcoming')
    const items = Array.isArray(body) ? body : (body?.results ?? [])
    return (items as any[]).map((r) => flatten(r))
  },
})

export const referenceSources = {
  tickers,
  'ticker-details': tickerDetails,
  dividends,
  splits,
  exchanges,
  conditions,
  'market-status': marketStatus,
  'market-holidays': marketHolidays,
}