// eslint-disable consistent-type-definitions
import { z } from '@bpinternal/zui'
import pLimit from 'p-limit'
import { ZaiContext } from '../context'
import { Response } from '../response'
import { getTokenizer } from '../tokenizer'
import { fastHash, stringify } from '../utils'
import { Zai } from '../zai'
import { PROMPT_INPUT_BUFFER, PROMPT_OUTPUT_BUFFER } from './constants'

export type Options = {
  /** The maximum number of tokens per item */
  tokensPerItem?: number
}

const _Options = z.object({
  tokensPerItem: z
    .number()
    .min(1)
    .max(100_000)
    .optional()
    .describe('The maximum number of tokens per item')
    .default(250),
})

// Evaluation criteria generated by LLM
type SortingCriteria = Record<
  string,
  {
    description: string
    labels: string[] // Ordered array of labels from FIRST to LAST, e.g., ['very_slow', 'slow', 'medium', 'fast', 'very_fast']
  }
>

declare module '@botpress/zai' {
  interface Zai {
    /**
     * Sorts array items based on natural language sorting criteria.
     *
     * This operation intelligently orders items according to your instructions, understanding
     * complex sorting logic like priority, quality, chronology, or any custom criteria.
     * Perfect for ranking, organizing, and prioritizing lists based on subjective or
     * multi-faceted criteria.
     *
     * @param input - Array of items to sort
     * @param instructions - Natural language description of how to sort (e.g., "by priority", "newest first")
     * @param options - Configuration for tokens per item
     * @returns Response resolving to the sorted array
     *
     * @example Sort by price
     * ```typescript
     * const products = [
     *   { name: 'Laptop', price: 999 },
     *   { name: 'Mouse', price: 29 },
     *   { name: 'Keyboard', price: 79 }
     * ]
     *
     * const sorted = await zai.sort(products, 'from least expensive to most expensive')
     * // Result: [Mouse ($29), Keyboard ($79), Laptop ($999)]
     * ```
     *
     * @example Sort by priority/urgency
     * ```typescript
     * const tasks = [
     *   "Update documentation",
     *   "Fix critical security bug",
     *   "Add new feature",
     *   "System is down - all users affected"
     * ]
     *
     * const prioritized = await zai.sort(tasks, 'by urgency and impact, most urgent first')
     * // Result: ["System is down...", "Fix critical security bug", "Add new feature", "Update documentation"]
     * ```
     *
     * @example Sort by quality/rating
     * ```typescript
     * const reviews = [
     *   "Product is okay",
     *   "Absolutely amazing! Best purchase ever!",
     *   "Terrible, broke immediately",
     *   "Good quality for the price"
     * ]
     *
     * const sorted = await zai.sort(reviews, 'by sentiment, most positive first')
     * // Result: [amazing review, good review, okay review, terrible review]
     * ```
     *
     * @example Sort by complexity
     * ```typescript
     * const problems = [
     *   "Fix typo in README",
     *   "Redesign entire authentication system",
     *   "Update a dependency version",
     *   "Implement new microservice architecture"
     * ]
     *
     * const sorted = await zai.sort(problems, 'by complexity, simplest first')
     * ```
     *
     * @example Sort by relevance to query
     * ```typescript
     * const documents = [
     *   "Article about cats",
     *   "Article about dogs and their training",
     *   "Article about dog breeds",
     *   "Article about fish"
     * ]
     *
     * const sorted = await zai.sort(
     *   documents,
     *   'by relevance to "dog training", most relevant first'
     * )
     * // Result: [dogs and training, dog breeds, cats, fish]
     * ```
     *
     * @example Sort candidates by fit
     * ```typescript
     * const candidates = [
     *   { name: 'Alice', experience: 5, skills: ['React', 'Node'] },
     *   { name: 'Bob', experience: 10, skills: ['Python', 'ML'] },
     *   { name: 'Charlie', experience: 3, skills: ['React', 'TypeScript'] }
     * ]
     *
     * const sorted = await zai.sort(
     *   candidates,
     *   'by fit for a senior React developer position, best fit first'
     * )
     * ```
     *
     * @example Sort chronologically
     * ```typescript
     * const events = [
     *   "Started the project last month",
     *   "Will launch next week",
     *   "Met with client yesterday",
     *   "Planning meeting tomorrow"
     * ]
     *
     * const chronological = await zai.sort(events, 'in chronological order')
     * // Understands relative time expressions
     * ```
     *
     * @example With token limit per item
     * ```typescript
     * const sorted = await zai.sort(
     *   longDocuments,
     *   'by relevance to climate change research',
     *   { tokensPerItem: 500 } // Allow 500 tokens per document
     * )
     * ```
     */
    sort<T>(input: Array<T>, instructions: string, options?: Options): Response<Array<T>, Array<T>>
  }
}

const END = '■END■'

const sort = async <T>(
  input: Array<T>,
  instructions: string,
  _options: Options | undefined,
  ctx: ZaiContext
): Promise<Array<T>> => {
  ctx.controller.signal.throwIfAborted()

  const options = _Options.parse(_options ?? {})
  const tokenizer = await getTokenizer()
  const model = await ctx.getModel()

  const taskId = ctx.taskId
  const taskType = 'zai.sort'

  // Handle empty or single element arrays
  if (input.length === 0) {
    return []
  }

  if (input.length === 1) {
    return input
  }

  const TOKENS_TOTAL_MAX = model.input.maxTokens - PROMPT_INPUT_BUFFER - PROMPT_OUTPUT_BUFFER

  // Phase 1: Generate sorting criteria from instructions + sample items
  const sampleSize = Math.min(5, input.length)
  const sampleItems = input.slice(0, sampleSize)
  const sampleItemsText = sampleItems.map((item, idx) => `■${idx}: ${stringify(item, false)}`).join('\n')

  // Get examples from adapter for criteria generation
  const criteriaInputStr = JSON.stringify({ instructions, sampleItems })
  const criteriaExamples =
    taskId && ctx.adapter
      ? await ctx.adapter.getExamples<string, Array<T>>({
          input: criteriaInputStr.slice(0, 1000),
          taskType,
          taskId,
        })
      : []

  // Format examples for few-shot learning in criteria generation
  const criteriaExampleMessages: Array<{ type: 'text'; role: 'user' | 'assistant'; content: string }> = []

  for (const example of criteriaExamples.slice(0, 2)) {
    try {
      const exampleInput = JSON.parse(example.input)
      const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput]

      criteriaExampleMessages.push({
        type: 'text',
        role: 'user',
        content: `Expert Example - Analyze sorting instruction: "${instructions}"\n\nSample items:\n${exampleItems
          .slice(0, 3)
          .map((el, i) => `■${i}: ${stringify(el, false).slice(0, 200)}`)
          .join('\n')}\n\nGenerate sorting criteria.`,
      })

      // Try to infer criteria from the example if available
      if (example.explanation) {
        criteriaExampleMessages.push({
          type: 'text',
          role: 'assistant',
          content: `${example.explanation}\n${END}`,
        })
      }
    } catch {
      // Skip malformed examples
    }
  }

  const generateCriteriaPrompt = `Analyze this sorting instruction: "${instructions}"

Sample items to be sorted:
${sampleItemsText}

Create 1-3 sorting criteria with ordered label arrays (3-10 labels each).

**CRITICAL RULES**:
1. Labels are single words, lowercase, no spaces, use underscores
2. Labels are ordered from FIRST to LAST in sorted result
3. If instruction says "from X to Y": first label represents X, last label represents Y
4. If instruction says "prioritize" or "highest/lowest priority":
   - First label = HIGHEST priority (top of todo list)
   - Last label = LOWEST priority (bottom of todo list)

Examples:

"from slowest to fastest" → first=slowest, last=fastest
■speed■
very_slow;slow;medium;fast;very_fast
■END■

"from most dangerous to least dangerous" → first=most dangerous, last=least dangerous
■danger■
extremely_dangerous;very_dangerous;dangerous;moderate;slightly_dangerous;harmless
■END■

"from least urgent (spam) to most urgent (bills)" → first=spam, last=bills
■urgency■
spam;promotional;normal;important;urgent;critical
■END■

"prioritize: highest priority=open old tickets; lowest priority=closed" → first=high priority, last=low priority
■status■
open_old;open_recent;closed
■age■
oldest;old;recent;new
■END■

Output format:
■criterion_name■
label1;label2;label3;label4
■END■

Use 3-10 labels per criterion. Labels should be intuitive and match the domain.
Keep criterion names short (1-2 words, lowercase, underscores).
`

  const { extracted: sortingCriteria } = await ctx.generateContent({
    systemPrompt: `You are creating sorting criteria with ordered label arrays.

CRITICAL: Output ordered labels from FIRST to LAST position in sorted result.
- Labels are single words, lowercase, underscores only
- 3-10 labels per criterion
- Order matters: first label = appears first, last label = appears last`,
    messages: [
      ...criteriaExampleMessages,
      {
        type: 'text',
        role: 'user',
        content: generateCriteriaPrompt,
      },
    ],
    transform: (text) => {
      const criteria: SortingCriteria = {}
      const criterionRegex = /■([^■]+)■\s*([^\n■]+)/g
      let match: RegExpExecArray | null

      while ((match = criterionRegex.exec(text)) !== null) {
        const name = (match[1] ?? '').trim().toLowerCase()
        const labelsStr = (match[2] ?? '').trim()

        if (!name || name === 'end') continue

        // Parse semicolon-separated labels
        const labels = labelsStr
          .split(';')
          .map((l) => l.trim().toLowerCase().replace(/\s+/g, '_'))
          .filter((l) => l.length > 0 && l.length < 50)

        if (labels.length >= 3 && labels.length <= 10) {
          criteria[name] = {
            description: `${labels.length} ordered labels`,
            labels,
          }
        }
      }

      if (Object.keys(criteria).length === 0) {
        throw new Error(`Failed to parse sorting criteria. LLM output: ${text.slice(0, 500)}`)
      }

      return criteria
    },
  })

  const criteriaKeys = Object.keys(sortingCriteria)
  if (criteriaKeys.length === 0) {
    throw new Error('No sorting criteria generated')
  }

  // Phase 2: Chunk items and score them in parallel
  const TOKENS_CRITERIA_MAX = Math.floor(TOKENS_TOTAL_MAX * 0.2)
  const TOKENS_ITEMS_MAX = TOKENS_TOTAL_MAX - TOKENS_CRITERIA_MAX

  const MAX_ITEMS_PER_CHUNK = 50

  // Prepare elements with indices
  const elements = input.map((element, idx) => ({
    element,
    index: idx,
    stringified: stringify(element, false),
  }))

  // Chunk elements
  const chunks: Array<typeof elements> = []
  let currentChunk: typeof elements = []
  let currentTokens = 0

  for (const elem of elements) {
    const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem)
    const elemTokens = tokenizer.count(truncated)

    if (
      (currentTokens + elemTokens > TOKENS_ITEMS_MAX || currentChunk.length >= MAX_ITEMS_PER_CHUNK) &&
      currentChunk.length > 0
    ) {
      chunks.push(currentChunk)
      currentChunk = []
      currentTokens = 0
    }

    currentChunk.push(elem)
    currentTokens += elemTokens
  }

  if (currentChunk.length > 0) {
    chunks.push(currentChunk)
  }

  // Phase 3: Score each chunk
  type ItemScore = {
    elementIndex: number
    scores: Record<string, number>
    totalScore: number
  }

  const scoreChunk = async (chunk: typeof elements): Promise<ItemScore[]> => {
    ctx.controller.signal.throwIfAborted()

    const chunkSize = chunk.length
    const chunkInputStr = JSON.stringify(chunk.map((c) => c.element))

    // Get examples from adapter for active learning
    const examples =
      taskId && ctx.adapter
        ? await ctx.adapter.getExamples<string, ItemScore[]>({
            input: chunkInputStr.slice(0, 1000),
            taskType,
            taskId,
          })
        : []

    // Check for exact match (cache hit)
    const key = fastHash(
      stringify({
        taskId,
        taskType,
        input: chunkInputStr,
        instructions,
      })
    )

    const exactMatch = examples.find((x) => x.key === key)
    if (exactMatch && exactMatch.output) {
      return exactMatch.output
    }

    const elementsText = chunk
      .map((elem, i) => {
        const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem)
        return `■${i}: ${truncated}■`
      })
      .join('\n')

    const criteriaText = criteriaKeys
      .map((key) => {
        const criterion = sortingCriteria[key]
        const labelsText = criterion.labels.join(';')
        return `**${key}**: ${labelsText}`
      })
      .join('\n')

    // Format examples for few-shot learning
    const exampleMessages: Array<{ type: 'text'; role: 'user' | 'assistant'; content: string }> = []

    for (const example of examples.slice(0, 3)) {
      try {
        const exampleInput = JSON.parse(example.input)
        const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput]

        exampleMessages.push({
          type: 'text',
          role: 'user',
          content: `Expert Example - Items to score:\n${exampleItems.map((el, i) => `■${i}: ${stringify(el, false).slice(0, 200)}■`).join('\n')}\n\nScore each item.`,
        })

        const exampleOutput = example.output
        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
          const formattedScores = exampleOutput
            .map((score) => {
              const pairs = criteriaKeys.map((key) => `${key}=${score.scores[key] ?? 0}`).join(';')
              return `■${score.elementIndex}:${pairs}■`
            })
            .join('\n')

          exampleMessages.push({
            type: 'text',
            role: 'assistant',
            content: `${formattedScores}\n${END}`,
          })

          if (example.explanation) {
            exampleMessages.push({
              type: 'text',
              role: 'assistant',
              content: `Reasoning: ${example.explanation}`,
            })
          }
        }
      } catch {
        // Skip malformed examples
      }
    }

    const { extracted } = await ctx.generateContent({
      systemPrompt: `You are ranking items for sorting using ordered label arrays.

${criteriaText}

Instructions: "${instructions}"

SCORING RULES:
- For each item and each criterion, assign ONE label from the ordered list
- Labels are ordered: first label = appears FIRST in sorted result, last label = appears LAST
- Choose the label that best describes each item

Output format:
■0:criterion1=label;criterion2=label■
■1:criterion1=label;criterion2=label■
${END}

IMPORTANT:
- Rank every item (■0 to ■${chunkSize - 1})
- Use exact criterion names: ${criteriaKeys.join(', ')}
- Use exact labels from the lists above (lowercase, underscores)
- Use semicolons (;) between criteria
- Use equals (=) between criterion and label`,
      stopSequences: [END],
      messages: [
        ...exampleMessages,
        {
          type: 'text',
          role: 'user',
          content: `Items to rank (■0 to ■${chunkSize - 1}):\n${elementsText}\n\nRank each item using the labeled scales.\nOutput format: ■index:criterion1=label;criterion2=label■\n${END}`,
        },
      ],
      transform: (text) => {
        const results: ItemScore[] = []
        const regex = /■(\d+):([^■]+)■/g
        let match: RegExpExecArray | null

        while ((match = regex.exec(text)) !== null) {
          const idx = parseInt(match[1] ?? '', 10)
          const labelsStr = match[2] ?? ''

          if (isNaN(idx) || idx < 0 || idx >= chunkSize) continue

          const scores: Record<string, number> = {}
          let total = 0

          const pairs = labelsStr.split(';').filter((x) => x.trim().length > 0)
          for (const pair of pairs) {
            const [criterion, labelStr] = pair.split('=').map((x) => x.trim().toLowerCase().replace(/\s+/g, '_'))
            if (!criterion || !labelStr) continue

            // Find the label index in the ordered array (index = score)
            const labels = sortingCriteria[criterion]?.labels ?? []
            const labelIndex = labels.findIndex((l) => l === labelStr)

            if (labelIndex >= 0) {
              // Use index as score (0 = first, higher = later)
              scores[criterion] = labelIndex
              total += labelIndex
            } else {
              // If label not found, use middle value
              const middleIndex = labels.length > 0 ? Math.floor(labels.length / 2) : 5
              scores[criterion] = middleIndex
              total += middleIndex
            }
          }

          results[idx] = {
            elementIndex: chunk[idx].index,
            scores,
            totalScore: total,
          }
        }

        // Fill in missing results with middle scores
        for (let i = 0; i < chunkSize; i++) {
          if (!results[i]) {
            const scores: Record<string, number> = {}
            let total = 0

            for (const key of criteriaKeys) {
              const labels = sortingCriteria[key]?.labels ?? []
              const middleIndex = labels.length > 0 ? Math.floor(labels.length / 2) : 5
              scores[key] = middleIndex
              total += middleIndex
            }

            results[i] = {
              elementIndex: chunk[i].index,
              scores,
              totalScore: total,
            }
          }
        }

        return results
      },
    })

    return extracted
  }

  // Process all chunks in parallel
  const limit = pLimit(10)
  const chunkPromises = chunks.map((chunk) => limit(() => scoreChunk(chunk)))
  const allScores = await Promise.all(chunkPromises)

  // Phase 4: Merge scores from all chunks
  // Build a map of elementIndex -> accumulated scores
  const scoreMap = new Map<number, { scores: Record<string, number>; totalScore: number; tieBreakOrder?: number }>()

  for (const chunkScores of allScores) {
    for (const itemScore of chunkScores) {
      const existing = scoreMap.get(itemScore.elementIndex)

      if (existing) {
        // Average the scores
        for (const key of criteriaKeys) {
          existing.scores[key] = (existing.scores[key] + (itemScore.scores[key] ?? 0)) / 2
        }
        existing.totalScore = (existing.totalScore + itemScore.totalScore) / 2
      } else {
        scoreMap.set(itemScore.elementIndex, {
          scores: { ...itemScore.scores },
          totalScore: itemScore.totalScore,
        })
      }
    }
  }

  // Verify we have scores for all elements
  if (scoreMap.size !== input.length) {
    throw new Error(`Score map size mismatch: expected ${input.length}, got ${scoreMap.size}`)
  }

  // Phase 5: Identify ties
  const scoreGroups = new Map<number, number[]>()
  for (const [index, scoreData] of scoreMap.entries()) {
    const roundedScore = Math.round(scoreData.totalScore * 100) // Round to avoid floating point issues
    const group = scoreGroups.get(roundedScore) ?? []
    group.push(index)
    scoreGroups.set(roundedScore, group)
  }

  // Find groups with more than one item (ties)
  const tiedGroups = Array.from(scoreGroups.values()).filter((group) => group.length > 1)

  // Phase 6: Tie-breaking - process all tied groups IN PARALLEL
  if (tiedGroups.length > 0) {
    const tieBreakLimit = pLimit(10)

    await Promise.all(
      tiedGroups.map((tiedIndices) =>
        tieBreakLimit(async () => {
          if (tiedIndices.length <= 1) return

          const tiedElements = tiedIndices.map((idx) => elements[idx])

          // Re-score just these items for tie-breaking
          const tieBreakText = tiedElements
            .map((elem, i) => {
              const truncated = tokenizer.truncate(elem.stringified, options.tokensPerItem)
              return `■${i}: ${truncated}■`
            })
            .join('\n')

          // Get examples from adapter for tie-breaking
          const tieBreakInputStr = JSON.stringify(tiedElements.map((e) => e.element))
          const tieBreakExamples =
            taskId && ctx.adapter
              ? await ctx.adapter.getExamples<string, Array<T>>({
                  input: tieBreakInputStr.slice(0, 1000),
                  taskType,
                  taskId,
                })
              : []

          // Format examples for few-shot learning in tie-breaking
          const tieBreakExampleMessages: Array<{ type: 'text'; role: 'user' | 'assistant'; content: string }> = []

          for (const example of tieBreakExamples.slice(0, 3)) {
            try {
              const exampleInput = JSON.parse(example.input)
              const exampleItems = Array.isArray(exampleInput) ? exampleInput : [exampleInput]

              tieBreakExampleMessages.push({
                type: 'text',
                role: 'user',
                content: `Expert Example - Items to order:\n${exampleItems.map((el, i) => `■${i}: ${stringify(el, false).slice(0, 200)}■`).join('\n')}\n\nOrder them from first to last.`,
              })

              if (Array.isArray(example.output) && example.output.length > 0) {
                const formattedOrder = example.output.map((_, i) => `■${i}■`).join('\n')

                tieBreakExampleMessages.push({
                  type: 'text',
                  role: 'assistant',
                  content: `${formattedOrder}\n${END}`,
                })

                if (example.explanation) {
                  tieBreakExampleMessages.push({
                    type: 'text',
                    role: 'assistant',
                    content: `Reasoning: ${example.explanation}`,
                  })
                }
              }
            } catch {
              // Skip malformed examples
            }
          }

          const { extracted: tieBreakOrder } = await ctx.generateContent({
            systemPrompt: `You are breaking a tie between items with identical total scores.

Instructions: ${instructions}

Criteria:
${criteriaKeys
  .map((key) => {
    const labels = sortingCriteria[key].labels.join(';')
    return `- ${key}: ${labels}`
  })
  .join('\n')}

Order these ${tiedElements.length} items from FIRST to LAST based on the instructions.
Earlier labels in each criterion should come FIRST.

Output format:
■original_index■
■original_index■
${END}

Output the indices in the order they should appear (first item at top).`,
            stopSequences: [END],
            messages: [
              ...tieBreakExampleMessages,
              {
                type: 'text',
                role: 'user',
                content: `Items with identical scores (need tie-breaking):\n${tieBreakText}\n\nOrder them from first to last.\nOutput format: ■index■ (one per line)\n${END}`,
              },
            ],
            transform: (text) => {
              const order: number[] = []
              const regex = /■(\d+)■/g
              let match: RegExpExecArray | null

              while ((match = regex.exec(text)) !== null) {
                const idx = parseInt(match[1] ?? '', 10)
                if (!isNaN(idx) && idx >= 0 && idx < tiedElements.length) {
                  order.push(idx)
                }
              }

              // If not all items were ordered, append missing ones
              for (let i = 0; i < tiedElements.length; i++) {
                if (!order.includes(i)) {
                  order.push(i)
                }
              }

              return order
            },
          })

          // Update scoreMap with tie-break order
          for (let i = 0; i < tieBreakOrder.length; i++) {
            const elementIndex = tiedElements[tieBreakOrder[i]].index
            const scoreData = scoreMap.get(elementIndex)
            if (scoreData) {
              scoreData.tieBreakOrder = i
            }
          }
        })
      )
    )
  }

  // Phase 7: Sort by total score, then by tie-break order
  const sorted = Array.from(scoreMap.entries())
    .sort((a, b) => {
      // First sort by total score
      const scoreDiff = a[1].totalScore - b[1].totalScore
      if (scoreDiff !== 0) return scoreDiff

      // If scores are equal, use tie-break order
      const orderA = a[1].tieBreakOrder ?? 0
      const orderB = b[1].tieBreakOrder ?? 0
      return orderA - orderB
    })
    .map(([index]) => elements[index].element)

  const result = sorted

  // Save example for active learning
  if (taskId && ctx.adapter && !ctx.controller.signal.aborted) {
    const key = fastHash(
      stringify({
        taskId,
        taskType,
        input: JSON.stringify(input),
        instructions,
      })
    )

    await ctx.adapter.saveExample({
      key,
      taskType,
      taskId,
      input: JSON.stringify(input),
      output: result,
      instructions,
      metadata: {
        cost: { input: 0, output: 0 },
        latency: 0,
        model: ctx.modelId,
        tokens: { input: 0, output: 0 },
      },
    })
  }

  return result
}

Zai.prototype.sort = function <T>(
  this: Zai,
  input: Array<T>,
  instructions: string,
  _options?: Options
): Response<Array<T>, Array<T>> {
  const context = new ZaiContext({
    client: this.client,
    modelId: this.Model,
    taskId: this.taskId,
    taskType: 'zai.sort',
    adapter: this.adapter,
    memoizer: this._resolveMemoizer(),
  })

  return new Response<Array<T>, Array<T>>(
    context,
    sort(input, instructions, _options, context),
    (result) => result // Simplified form is just the sorted array
  )
}
