import { definePlugin, defineSource } from '@jaspers-ai/sdk'
import { z } from 'zod'
import { getJson, HOST, searchQuery } from './client'
import { detailRow, summaryRow, type Study } from './rows'

// ClinicalTrials.gov API v2: search studies and read one study, no key.

const PHASES = ['EARLY_PHASE1', 'PHASE1', 'PHASE2', 'PHASE3', 'PHASE4', 'NA'] as const
const STATUSES = [
  'RECRUITING',
  'NOT_YET_RECRUITING',
  'ACTIVE_NOT_RECRUITING',
  'ENROLLING_BY_INVITATION',
  'COMPLETED',
  'SUSPENDED',
  'TERMINATED',
  'WITHDRAWN',
  'UNKNOWN',
] as const
const date = z
  .string()
  .regex(/^\d{4}(-\d{2}(-\d{2})?)?$/)
  .optional()

const searchStudies = defineSource({
  description:
    'Search ClinicalTrials.gov studies by condition, intervention/drug, sponsor, free text, phase, overall status, and start or primary completion date range. One row per study: nctId, briefTitle, leadSponsor, collaborators, phases, overallStatus, startDate, primaryCompletionDate (with its ACTUAL/ESTIMATED type), completionDate, conditions, interventions, enrollment, studyType, lastUpdatePostDate, url. meta has totalCount and nextPageToken (pass it back as pageToken for the next page). For catalysts: sponsor, phases [PHASE3], statuses [RECRUITING, ACTIVE_NOT_RECRUITING], primaryCompletionFrom today, sort PrimaryCompletionDate:asc.',
  hosts: [HOST],
  input: z.object({
    condition: z.string().optional().describe('Condition or disease (query.cond), like "non-small cell lung cancer".'),
    intervention: z.string().optional().describe('Intervention or drug name (query.intr), like "pembrolizumab".'),
    sponsor: z.string().optional().describe('Sponsor or collaborator (query.spons), like "Pfizer".'),
    term: z.string().optional().describe('Free-text search over the whole record (query.term).'),
    phases: z.array(z.enum(PHASES)).optional().describe('Phases to match, any of them: PHASE1, PHASE2, PHASE3, PHASE4, EARLY_PHASE1, NA.'),
    statuses: z.array(z.enum(STATUSES)).optional().describe('Overall statuses to match (filter.overallStatus).'),
    startFrom: date.describe('Earliest start date, YYYY-MM-DD (or YYYY-MM).'),
    startTo: date.describe('Latest start date, YYYY-MM-DD.'),
    primaryCompletionFrom: date.describe('Earliest primary completion date, YYYY-MM-DD.'),
    primaryCompletionTo: date.describe('Latest primary completion date, YYYY-MM-DD.'),
    sort: z
      .string()
      .optional()
      .describe('Sort as Field:asc or Field:desc, like LastUpdatePostDate:desc, StartDate:desc, PrimaryCompletionDate:asc; a bare field sorts ascending. Default is relevance.'),
    pageSize: z.number().int().min(1).max(1000).default(50).describe('Studies per page, up to 1000.'),
    pageToken: z.string().optional().describe('meta.nextPageToken from the previous page.'),
  }),
  async run(args, ctx) {
    const page = (await getJson(ctx, '/studies', searchQuery(args))) as {
      studies?: Study[]
      totalCount?: number
      nextPageToken?: string
    }
    return {
      rows: (page.studies ?? []).map(summaryRow),
      totalCount: page.totalCount ?? null,
      nextPageToken: page.nextPageToken ?? null,
    }
  },
})

const getStudy = defineSource({
  description:
    "One ClinicalTrials.gov study by NCT ID: the search row's fields plus officialTitle, briefSummary, primaryOutcomes, arms, interventions, eligibility (sex, ages, healthy volunteers, criteria), locationsCount, hasResults, studyFirstPostDate, resultsFirstPostDate.",
  hosts: [HOST],
  input: z.object({
    nctId: z.string().regex(/^NCT\d{8}$/i).describe('The NCT ID, like NCT04368728.'),
  }),
  async run(args, ctx) {
    const study = (await getJson(ctx, `/studies/${args.nctId.toUpperCase()}`, new URLSearchParams({ format: 'json' }))) as Study
    return [detailRow(study)]
  },
})

export default definePlugin({
  id: 'clinicaltrials',
  sources: { 'search-studies': searchStudies, 'get-study': getStudy },
})