import { defineSource } from '@jaspers-ai/sdk'
import { z } from 'zod'
import type { BackendContext } from '@jaspers-ai/sdk'
import { DATA_HOST, alpaca, paged, type Query } from './client'
import { mid, num, str } from './rows'

/** Symbols as Alpaca takes them: upper case, comma separated. */
export function list(symbols: string[]): string {
  return symbols.map((s) => s.trim().toUpperCase()).filter(Boolean).join(',')
}

const symbols = z
  .array(z.string().min(1))
  .min(1)
  .max(100)
  .describe('Stock symbols, like ["AAPL","MSFT"].')
const feed = z
  .enum(['iex', 'sip', 'delayed_sip'])
  .default('iex')
  .describe('Data feed: iex (free plan), sip (all US exchanges; recent data needs a paid plan), delayed_sip (SIP, 15 min delayed).')
const start = z.string().optional().describe('Start, RFC 3339 or YYYY-MM-DD, like 2025-01-02 or 2025-01-02T14:30:00Z.')
const end = z.string().optional().describe('End, RFC 3339 or YYYY-MM-DD. Leave out for now.')
const sort = z.enum(['asc', 'desc']).default('asc').describe('Time order.')

type Fields = Record<string, unknown>
interface Snapshot {
  latestTrade?: Fields
  latestQuote?: Fields
  minuteBar?: Fields
  dailyBar?: Fields
  prevDailyBar?: Fields
}

/** The start Alpaca gets when none is given for bars: 30 days ago, as a date. */
export function monthAgo(): string {
  return new Date(Date.now() - 30 * 86400_000).toISOString().slice(0, 10)
}

/** A paged, symbol-keyed history flattened to rows by `pick`, with meta saying whether it was cut. */
export async function history(
  ctx: BackendContext,
  path: string,
  key: string,
  query: Query,
  limit: number,
  pick: (item: Record<string, unknown>) => Record<string, unknown>,
) {
  const { data, truncated } = await paged(ctx, path, query, key, limit, 10000)
  const rows: Record<string, unknown>[] = []
  for (const [symbol, items] of Object.entries(data))
    for (const item of items as Record<string, unknown>[]) rows.push({ symbol, ...pick(item) })
  return { rows: rows.slice(0, limit), truncated }
}

const bar = (b: Record<string, unknown>) => ({
  t: str(b.t),
  o: num(b.o),
  h: num(b.h),
  l: num(b.l),
  c: num(b.c),
  v: num(b.v),
  n: num(b.n),
  vw: num(b.vw),
})
export { bar as barRow }

export const stockBars = defineSource({
  description:
    'Historical OHLCV bars for US stocks from Alpaca. One row per bar: symbol, t (bar start, ISO), o, h, l, c, v (volume), n (trade count), vw (VWAP).',
  hosts: [DATA_HOST],
  input: z.object({
    symbols,
    timeframe: z.string().default('1Day').describe('Bar size: 1Min, 5Min, 15Min, 30Min, 1Hour, 4Hour, 1Day, 1Week, 1Month.'),
    start: start.describe('Start, RFC 3339 or YYYY-MM-DD. Defaults to 30 days ago.'),
    end,
    limit: z.number().int().min(1).max(10000).default(1000).describe('Most bars in all, across symbols.'),
    adjustment: z.enum(['raw', 'split', 'dividend', 'all']).default('raw').describe('Corporate action adjustment.'),
    feed,
    sort,
  }),
  async run(args, ctx) {
    const query = {
      symbols: list(args.symbols),
      timeframe: args.timeframe,
      start: args.start ?? monthAgo(),
      end: args.end,
      adjustment: args.adjustment,
      feed: args.feed,
      sort: args.sort,
    }
    return history(ctx, '/v2/stocks/bars', 'bars', query, args.limit, bar)
  },
})

export const stockQuotes = defineSource({
  description:
    'Historical NBBO quotes for US stocks from Alpaca. Rows: symbol, t, bp (bid), bs (bid size), ap (ask), as (ask size), bx, ax (exchanges), c (conditions). Without start Alpaca begins at the start of today.',
  hosts: [DATA_HOST],
  input: z.object({
    symbols,
    start,
    end,
    limit: z.number().int().min(1).max(10000).default(1000).describe('Most quotes in all, across symbols.'),
    feed,
    sort,
  }),
  async run(args, ctx) {
    const query = { symbols: list(args.symbols), start: args.start, end: args.end, feed: args.feed, sort: args.sort }
    return history(ctx, '/v2/stocks/quotes', 'quotes', query, args.limit, (q) => ({
      t: str(q.t),
      bp: num(q.bp),
      bs: num(q.bs),
      ap: num(q.ap),
      as: num(q.as),
      bx: str(q.bx),
      ax: str(q.ax),
      c: Array.isArray(q.c) ? q.c.join(',') : str(q.c),
    }))
  },
})

export const stockTrades = defineSource({
  description:
    'Historical trades for US stocks from Alpaca. Rows: symbol, t, p (price), s (size), x (exchange), c (conditions), i (trade id). Without start Alpaca begins at the start of today.',
  hosts: [DATA_HOST],
  input: z.object({
    symbols,
    start,
    end,
    limit: z.number().int().min(1).max(10000).default(1000).describe('Most trades in all, across symbols.'),
    feed,
    sort,
  }),
  async run(args, ctx) {
    const query = { symbols: list(args.symbols), start: args.start, end: args.end, feed: args.feed, sort: args.sort }
    return history(ctx, '/v2/stocks/trades', 'trades', query, args.limit, (t) => ({
      t: str(t.t),
      p: num(t.p),
      s: num(t.s),
      x: str(t.x),
      c: Array.isArray(t.c) ? t.c.join(',') : str(t.c),
      i: num(t.i),
    }))
  },
})

export const stockLatestQuote = defineSource({
  description:
    'Latest bid/ask quote for US stocks from Alpaca. One row per symbol: symbol, bid, bid_size, ask, ask_size, mid, bid_exchange, ask_exchange, time.',
  hosts: [DATA_HOST],
  input: z.object({ symbols, feed }),
  async run(args, ctx) {
    const answer = (await alpaca(ctx, {
      host: DATA_HOST,
      path: '/v2/stocks/quotes/latest',
      query: { symbols: list(args.symbols), feed: args.feed },
    })) as { quotes?: Record<string, Record<string, unknown>> } | null
    return Object.entries(answer?.quotes ?? {}).map(([symbol, q]) => {
      const bid = num(q.bp)
      const ask = num(q.ap)
      return {
        symbol,
        bid,
        bid_size: num(q.bs),
        ask,
        ask_size: num(q.as),
        mid: mid(bid, ask),
        bid_exchange: str(q.bx),
        ask_exchange: str(q.ax),
        time: str(q.t),
      }
    })
  },
})

export const stockLatestTrade = defineSource({
  description: 'Latest trade for US stocks from Alpaca. One row per symbol: symbol, price, size, exchange, conditions, time.',
  hosts: [DATA_HOST],
  input: z.object({ symbols, feed }),
  async run(args, ctx) {
    const answer = (await alpaca(ctx, {
      host: DATA_HOST,
      path: '/v2/stocks/trades/latest',
      query: { symbols: list(args.symbols), feed: args.feed },
    })) as { trades?: Record<string, Record<string, unknown>> } | null
    return Object.entries(answer?.trades ?? {}).map(([symbol, t]) => ({
      symbol,
      price: num(t.p),
      size: num(t.s),
      exchange: str(t.x),
      conditions: Array.isArray(t.c) ? t.c.join(',') : str(t.c),
      time: str(t.t),
    }))
  },
})

export const stockSnapshots = defineSource({
  description:
    'Snapshots of US stocks from Alpaca, one row per symbol: last trade price and time, bid/ask, the minute bar close, the daily bar (day_open, day_high, day_low, day_close, day_volume, day_vwap), prev_close, change (last - prev_close), and change_pct (in percent).',
  hosts: [DATA_HOST],
  input: z.object({ symbols, feed }),
  async run(args, ctx) {
    const answer = (await alpaca(ctx, {
      host: DATA_HOST,
      path: '/v2/stocks/snapshots',
      query: { symbols: list(args.symbols), feed: args.feed },
    })) as Record<string, Snapshot | null> | null
    // The multi-symbol endpoint answers a map keyed by symbol; tolerate a { snapshots } wrapper too.
    const map = ((answer as { snapshots?: Record<string, Snapshot | null> } | null)?.snapshots ?? answer ?? {}) as Record<
      string,
      Snapshot | null
    >
    return Object.entries(map).map(([symbol, s]) => {
      const trade = s?.latestTrade ?? {}
      const quote = s?.latestQuote ?? {}
      const day = s?.dailyBar ?? {}
      const prev = s?.prevDailyBar ?? {}
      const last = num(trade.p) ?? num(day.c)
      const prevClose = num(prev.c)
      const change = last !== null && prevClose ? Math.round((last - prevClose) * 10000) / 10000 : null
      return {
        symbol,
        last,
        last_size: num(trade.s),
        last_time: str(trade.t),
        bid: num(quote.bp),
        bid_size: num(quote.bs),
        ask: num(quote.ap),
        ask_size: num(quote.as),
        minute_close: num(s?.minuteBar?.c),
        minute_time: str(s?.minuteBar?.t),
        day_open: num(day.o),
        day_high: num(day.h),
        day_low: num(day.l),
        day_close: num(day.c),
        day_volume: num(day.v),
        day_vwap: num(day.vw),
        day_time: str(day.t),
        prev_close: prevClose,
        change,
        change_pct: change !== null && prevClose ? Math.round((change / prevClose) * 1_000_000) / 10000 : null,
      }
    })
  },
})