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

// Bars: aggregates for any asset class, previous close, the grouped market day, one day's open/close.

const ticker = z.string().min(1).describe('A ticker: AAPL for a stock, O:SPY251219C00600000 an option, I:SPX an index, C:EURUSD forex, X:BTCUSD crypto.')

const aggregates = defineSource({
  description:
    'Aggregate bars for any ticker (stocks, options O:, indices I:, forex C:, crypto X:) over a date range: time (ISO), open, high, low, close, volume, vwap, transactions. Newest first by default.',
  hosts: HOSTS,
  input: z.object({
    ticker,
    multiplier: z.number().int().min(1).default(1).describe('How many timespans make one bar, like 5 with minute for 5-minute bars.'),
    timespan: z.enum(['second', 'minute', 'hour', 'day', 'week', 'month', 'quarter', 'year']).default('day').describe('The bar size unit.'),
    from: z.string().optional().describe('Start date YYYY-MM-DD or a millisecond timestamp. Defaults to one year ago.'),
    to: z.string().optional().describe('End date YYYY-MM-DD or a millisecond timestamp. Defaults to today.'),
    adjusted: z.boolean().default(true).describe('Adjust for splits.'),
    sort: z.enum(['asc', 'desc']).default('desc').describe('desc for newest first, asc for oldest first.'),
    limit: z.number().int().min(1).max(50000).default(5000).describe('Most bars to answer.'),
    base: baseField,
  }),
  async run(args, ctx) {
    const t = upper(args.ticker)
    const from = args.from ?? daysAgo(365)
    const to = args.to ?? daysAgo(0)
    const path = `/v2/aggs/ticker/${t}/range/${args.multiplier}/${args.timespan}/${from}/${to}`
    const res = await list(ctx, args.base, path, { adjusted: args.adjusted, sort: args.sort, limit: args.limit }, args.limit)
    return { rows: res.rows.map((b) => bar(b, t)), ticker: t, from, to, more: res.more, pages: res.pages }
  },
})

const previousClose = defineSource({
  description: "The previous trading day's bar for a ticker (any asset class): open, high, low, close, volume, vwap.",
  hosts: HOSTS,
  input: z.object({ ticker, adjusted: z.boolean().default(true).describe('Adjust for splits.'), base: baseField }),
  async run(args, ctx) {
    const t = upper(args.ticker)
    const body = await get(ctx, args.base, `/v2/aggs/ticker/${t}/prev`, { adjusted: args.adjusted })
    return ((body?.results ?? []) as any[]).map((b) => bar(b, t))
  },
})

const groupedDaily = defineSource({
  description: 'Daily bars for every US stock on one date, highest volume first: ticker, open, high, low, close, volume, vwap, transactions.',
  hosts: HOSTS,
  input: z.object({
    date: z.string().describe('The trading date, YYYY-MM-DD.'),
    adjusted: z.boolean().default(true).describe('Adjust for splits.'),
    include_otc: z.boolean().default(false).describe('Include OTC securities.'),
    limit: z.number().int().min(1).max(15000).default(1000).describe('Most rows to answer, by volume.'),
    base: baseField,
  }),
  async run(args, ctx) {
    const body = await get(ctx, args.base, `/v2/aggs/grouped/locale/us/market/stocks/${args.date}`, {
      adjusted: args.adjusted,
      include_otc: args.include_otc,
    })
    const bars = ((body?.results ?? []) as any[]).sort((a, b) => (b.v ?? 0) - (a.v ?? 0))
    return { rows: bars.slice(0, args.limit).map((b) => bar(b)), date: args.date, total: bars.length }
  },
})

const dailyOpenClose = defineSource({
  description: 'Open, high, low, close, volume, pre-market and after-hours prices for one ticker on one date.',
  hosts: HOSTS,
  input: z.object({
    ticker,
    date: z.string().describe('The date, YYYY-MM-DD.'),
    adjusted: z.boolean().default(true).describe('Adjust for splits.'),
    base: baseField,
  }),
  async run(args, ctx) {
    const t = upper(args.ticker)
    const b = await get(ctx, args.base, `/v1/open-close/${t}/${args.date}`, { adjusted: args.adjusted })
    return [
      {
        ticker: b.symbol ?? t,
        date: b.from ?? args.date,
        open: b.open ?? null,
        high: b.high ?? null,
        low: b.low ?? null,
        close: b.close ?? null,
        volume: b.volume ?? null,
        preMarket: b.preMarket ?? null,
        afterHours: b.afterHours ?? null,
      },
    ]
  },
})

export const barSources = {
  aggregates,
  'previous-close': previousClose,
  'grouped-daily': groupedDaily,
  'daily-open-close': dailyOpenClose,
}