// 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'

// Rating scale constants
const RATING_VALUES = {
  very_bad: 1,
  bad: 2,
  average: 3,
  good: 4,
  very_good: 5,
} as const

type RatingLabel = keyof typeof RATING_VALUES

// Evaluation criteria generated by LLM
type EvaluationCriteria = Record<
  string,
  {
    very_bad: string
    bad: string
    average: string
    good: string
    very_good: string
  }
>

export type RatingInstructions = string | Record<string, string>

export type Options = {
  /** The maximum number of tokens per item */
  tokensPerItem?: number
  /** The maximum number of items to rate per chunk */
  maxItemsPerChunk?: number
}

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

// Result types based on instructions type
export type RatingResult<T extends RatingInstructions> = T extends string
  ? {
      [key: string]: number // criteria scores
      total: number // sum of all criteria
    }
  : T extends Record<string, string>
    ? {
        [K in keyof T]: number // score for each criterion
      } & {
        total: number // sum of all criteria
      }
    : never

export type SimplifiedRatingResult<T extends RatingInstructions> = T extends string ? number : RatingResult<T>

declare module '@botpress/zai' {
  interface Zai {
    /**
     * Rates array items on a 1-5 scale based on single or multiple criteria.
     *
     * This operation evaluates each item and assigns numeric ratings (1-5) where:
     * - 1 = Very Bad, 2 = Bad, 3 = Average, 4 = Good, 5 = Very Good
     *
     * Supports both simple single-criterion rating (string instructions) and
     * multi-criteria rating (object with criterion → description mapping).
     *
     * @param input - Array of items to rate
     * @param instructions - Single criterion (string) or multiple criteria (object)
     * @param options - Configuration for chunking and tokens per item
     * @returns Response with ratings array (simplified to numbers for single criterion)
     *
     * @example Single criterion rating
     * ```typescript
     * const reviews = [
     *   "Amazing product! Best purchase ever!",
     *   "It's okay, nothing special",
     *   "Terrible quality, broke immediately"
     * ]
     *
     * const ratings = await zai.rate(reviews, 'Rate the sentiment')
     * // Result: [5, 3, 1] (simplified to numbers)
     *
     * // Get full details
     * const { output } = await zai.rate(reviews, 'Rate the sentiment').result()
     * // output[0]: { sentiment: 5, total: 5 }
     * ```
     *
     * @example Multi-criteria rating
     * ```typescript
     * const essays = [
     *   "... student essay text ...",
     *   "... another essay ...",
     * ]
     *
     * const ratings = await zai.rate(essays, {
     *   grammar: 'Rate the grammar and spelling',
     *   clarity: 'Rate how clear and well-organized the writing is',
     *   argumentation: 'Rate the strength of arguments and evidence'
     * })
     *
     * // Result: [
     * //   { grammar: 4, clarity: 5, argumentation: 3, total: 12 },
     * //   { grammar: 3, clarity: 4, argumentation: 4, total: 11 }
     * // ]
     * ```
     *
     * @example Rating customer support conversations
     * ```typescript
     * const conversations = [
     *   { agent: 'John', messages: [...], duration: 300 },
     *   { agent: 'Jane', messages: [...], duration: 180 }
     * ]
     *
     * const ratings = await zai.rate(conversations, {
     *   professionalism: 'How professional was the agent?',
     *   helpfulness: 'How helpful was the agent in solving the issue?',
     *   efficiency: 'How efficiently was the conversation handled?'
     * })
     *
     * // Calculate average scores
     * const avgProfessionalism = ratings.reduce((sum, r) => sum + r.professionalism, 0) / ratings.length
     * ```
     *
     * @example Rating code quality
     * ```typescript
     * const codeSamples = [
     *   "function foo() { return x + y }",
     *   "const calculateTotal = (items) => items.reduce((sum, item) => sum + item.price, 0)",
     * ]
     *
     * const quality = await zai.rate(codeSamples, {
     *   readability: 'How readable and clear is the code?',
     *   best_practices: 'Does it follow coding best practices?',
     *   documentation: 'Is the code well-documented?'
     * })
     * ```
     *
     * @example Product review rating
     * ```typescript
     * const products = [
     *   { name: 'Laptop', reviews: [...], avgStars: 4.2 },
     *   { name: 'Mouse', reviews: [...], avgStars: 3.8 }
     * ]
     *
     * const ratings = await zai.rate(
     *   products,
     *   'Rate overall product quality based on reviews and rating',
     *   {
     *     tokensPerItem: 500, // Allow more tokens for detailed reviews
     *     maxItemsPerChunk: 20 // Process 20 products per chunk
     *   }
     * )
     * ```
     *
     * @example Finding highest-rated items
     * ```typescript
     * const ratings = await zai.rate(items, {
     *   quality: 'Product quality',
     *   value: 'Value for money',
     *   design: 'Design and aesthetics'
     * })
     *
     * // Sort by total score
     * const sorted = items
     *   .map((item, idx) => ({ item, rating: ratings[idx] }))
     *   .sort((a, b) => b.rating.total - a.rating.total)
     *
     * console.log('Top rated:', sorted[0].item)
     * console.log('Score breakdown:', sorted[0].rating)
     * ```
     */
    rate<T, I extends RatingInstructions>(
      input: Array<T>,
      instructions: I,
      options?: Options
    ): Response<Array<RatingResult<I>>, Array<SimplifiedRatingResult<I>>>
  }
}

const END = '■END■'

const rate = async <T, I extends RatingInstructions>(
  input: Array<T>,
  instructions: I,
  _options: Options | undefined,
  ctx: ZaiContext
): Promise<Array<RatingResult<I>>> => {
  ctx.controller.signal.throwIfAborted()
  const options = _Options.parse(_options ?? {})
  const tokenizer = await getTokenizer()
  const model = await ctx.getModel()

  // Handle empty array
  if (input.length === 0) {
    return []
  }

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

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

  // Phase 1: Generate evaluation criteria
  const isStringInstructions = typeof instructions === 'string'
  const criteriaKeys: string[] = isStringInstructions
    ? [] // Will be generated by LLM
    : Object.keys(instructions as Record<string, string>)

  const generateCriteriaPrompt = isStringInstructions
    ? `Generate 3-5 evaluation criteria for: "${instructions}"

For each criterion, provide 5 labels (very_bad, bad, average, good, very_good) with brief descriptions.

Output format (JSON):
{
  "criterion1_name": {
    "very_bad": "description",
    "bad": "description",
    "average": "description",
    "good": "description",
    "very_good": "description"
  },
  "criterion2_name": { ... }
}

Keep criterion names short (1-2 words, lowercase, use underscores).
Keep descriptions brief (5-10 words each).`
    : `For these evaluation criteria, provide 5 labels (very_bad, bad, average, good, very_good) with brief descriptions for each:

${criteriaKeys.map((key) => `- ${key}: ${(instructions as Record<string, string>)[key]}`).join('\n')}

Output format (JSON):
{
  "${criteriaKeys[0]}": {
    "very_bad": "description",
    "bad": "description",
    "average": "description",
    "good": "description",
    "very_good": "description"
  }
  ${criteriaKeys.length > 1 ? '...' : ''}
}

Keep descriptions brief (5-10 words each).`

  const { extracted: evaluationCriteria } = await ctx.generateContent({
    systemPrompt: `You are creating evaluation criteria for rating items on a 1-5 scale.
Each criterion must have exactly 5 labels: very_bad (1), bad (2), average (3), good (4), very_good (5).
Output valid JSON only.`,
    messages: [
      {
        type: 'text',
        role: 'user',
        content: generateCriteriaPrompt,
      },
    ],
    transform: (text) => {
      // Extract JSON from markdown code blocks if present
      const jsonMatch = text.match(/```(?:json)?\s*(\{[\s\S]*?\})\s*```/) || text.match(/(\{[\s\S]*\})/)
      if (!jsonMatch) {
        throw new Error('Failed to parse evaluation criteria JSON')
      }
      return JSON.parse(jsonMatch[1]) as EvaluationCriteria
    },
  })

  // Extract final criteria keys
  const finalCriteriaKeys = Object.keys(evaluationCriteria)
  if (finalCriteriaKeys.length === 0) {
    throw new Error('No evaluation criteria generated')
  }

  // Phase 2: Chunk items
  const TOKENS_CRITERIA_MAX = Math.floor(TOKENS_TOTAL_MAX * 0.3)
  const TOKENS_ITEMS_MAX = TOKENS_TOTAL_MAX - TOKENS_CRITERIA_MAX

  let chunks: Array<typeof input> = []
  let currentChunk: typeof input = []
  let currentChunkTokens = 0

  for (const element of input) {
    const elementAsString = tokenizer.truncate(stringify(element, false), options.tokensPerItem)
    const elementTokens = tokenizer.count(elementAsString)

    if (currentChunkTokens + elementTokens > TOKENS_ITEMS_MAX || currentChunk.length >= options.maxItemsPerChunk) {
      if (currentChunk.length > 0) {
        chunks.push(currentChunk)
      }
      currentChunk = []
      currentChunkTokens = 0
    }

    currentChunk.push(element)
    currentChunkTokens += elementTokens
  }

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

  chunks = chunks.filter((x) => x.length > 0)

  // Phase 3: Rate each chunk in parallel
  type ChunkResult = {
    ratings: Array<Record<string, number>>
    meta: { cost: { input: number; output: number }; latency: number; tokens: { input: number; output: number } }
  }

  const rateChunk = async (chunk: typeof input): Promise<ChunkResult> => {
    ctx.controller.signal.throwIfAborted()

    // Get examples from adapter for active learning
    const chunkInputStr = JSON.stringify(chunk)
    const examples =
      taskId && ctx.adapter
        ? await ctx.adapter.getExamples<string, Array<Record<string, number>>>({
            input: chunkInputStr.slice(0, 1000), // Limit search string length
            taskType,
            taskId,
          })
        : []

    // Check for exact match (cache hit)
    const key = fastHash(
      stringify({
        taskId,
        taskType,
        input: chunkInputStr,
        instructions: stringify(instructions),
      })
    )
    const exactMatch = examples.find((x) => x.key === key)
    if (exactMatch && exactMatch.output) {
      // Return cached result with zero cost
      return {
        ratings: exactMatch.output,
        meta: { cost: { input: 0, output: 0 }, latency: 0, tokens: { input: 0, output: 0 } },
      }
    }

    const formatCriteria = () => {
      return finalCriteriaKeys
        .map((key) => {
          const labels = evaluationCriteria[key]
          return `**${key}**:
  - very_bad (1): ${labels?.very_bad}
  - bad (2): ${labels?.bad}
  - average (3): ${labels?.average}
  - good (4): ${labels?.good}
  - very_good (5): ${labels?.very_good}`
        })
        .join('\n\n')
    }

    const formatItems = (items: typeof chunk) => {
      return items
        .map((item, idx) => {
          const itemStr = tokenizer.truncate(stringify(item, false), options.tokensPerItem)
          return `■${idx}: ${itemStr}■`
        })
        .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, 5)) {
      // User message with input
      try {
        const exampleInput = JSON.parse(example.input)
        exampleMessages.push({
          type: 'text',
          role: 'user',
          content: `Expert Example - Items to rate:
${formatItems(Array.isArray(exampleInput) ? exampleInput : [exampleInput])}

Rate each item on all criteria.`,
        })

        // Assistant message with ratings
        const exampleOutput = example.output
        if (Array.isArray(exampleOutput) && exampleOutput.length > 0) {
          const formattedRatings = exampleOutput
            .map((rating, idx) => {
              const pairs = finalCriteriaKeys
                .map((key) => {
                  const value = rating[key]
                  if (typeof value === 'number') {
                    // Convert number back to label
                    const labelMap: Record<number, string> = {
                      1: 'very_bad',
                      2: 'bad',
                      3: 'average',
                      4: 'good',
                      5: 'very_good',
                    }
                    return `${key}=${labelMap[value] || 'average'}`
                  }
                  return null
                })
                .filter(Boolean)
                .join(';')
              return `■${idx}:${pairs}■`
            })
            .join('\n')

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

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

    const { extracted, meta } = await ctx.generateContent({
      systemPrompt: `You are rating items based on evaluation criteria.

Evaluation Criteria:
${formatCriteria()}

For each item, rate it on EACH criterion using one of these labels:
very_bad, bad, average, good, very_good

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

IMPORTANT:
- Rate every item (■0 to ■${chunk.length - 1})
- Use exact criterion names: ${finalCriteriaKeys.join(', ')}
- Use exact label names: very_bad, bad, average, good, very_good
- Use semicolons (;) between criteria
- Use equals (=) between criterion and label`,
      stopSequences: [END],
      messages: [
        ...exampleMessages,
        {
          type: 'text',
          role: 'user',
          content: `Items to rate (■0 to ■${chunk.length - 1}):
${formatItems(chunk)}

Rate each item on all criteria.
Output format: ■index:criterion1=label;criterion2=label■
${END}`,
        },
      ],
      transform: (text) => {
        const results: Array<Record<string, number>> = []

        // Parse ratings: ■0:affordability=good;quality=very_good■
        const regex = /■(\d+):([^■]+)■/g
        let match: RegExpExecArray | null

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

          if (isNaN(idx) || idx < 0 || idx >= chunk.length) {
            continue
          }

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

          // Parse criterion=label pairs
          const pairs = ratingsStr.split(';').filter((x) => x.trim().length > 0)
          for (const pair of pairs) {
            const [criterion, label] = pair.split('=').map((x) => x.trim())
            if (!criterion || !label) continue

            // Convert label to number
            const labelLower = label.toLowerCase().replace(/\s+/g, '_')
            const ratingValue = RATING_VALUES[labelLower as RatingLabel] ?? 3 // default to average

            itemRatings[criterion] = ratingValue
            total += ratingValue
          }

          itemRatings.total = total
          results[idx] = itemRatings
        }

        // Fill in missing results with defaults
        for (let i = 0; i < chunk.length; i++) {
          if (!results[i]) {
            const defaultRatings: Record<string, number> = {}
            let total = 0
            for (const key of finalCriteriaKeys) {
              defaultRatings[key] = 3 // average
              total += 3
            }
            defaultRatings.total = total
            results[i] = defaultRatings
          }
        }

        return results
      },
    })

    return { ratings: extracted, meta }
  }

  // Process chunks in parallel with p-limit
  const limit = pLimit(10)
  const chunkPromises = chunks.map((chunk) => limit(() => rateChunk(chunk)))

  const ratedChunks = await Promise.all(chunkPromises)

  // Phase 4: Flatten results and accumulate metadata
  const allRatings = ratedChunks.flatMap((result) => result.ratings) as Array<RatingResult<I>>

  // Accumulate metadata from all chunks
  const totalMeta = ratedChunks.reduce(
    (acc, result) => ({
      cost: {
        input: acc.cost.input + result.meta.cost.input,
        output: acc.cost.output + result.meta.cost.output,
      },
      latency: Math.max(acc.latency, result.meta.latency), // Use max latency
      tokens: {
        input: acc.tokens.input + result.meta.tokens.input,
        output: acc.tokens.output + result.meta.tokens.output,
      },
    }),
    {
      cost: { input: 0, output: 0 },
      latency: 0,
      tokens: { input: 0, output: 0 },
    }
  )

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

    await ctx.adapter.saveExample({
      key,
      taskType,
      taskId,
      input: JSON.stringify(input),
      output: allRatings,
      instructions: typeof instructions === 'string' ? instructions : JSON.stringify(instructions),
      metadata: {
        cost: {
          input: totalMeta.cost.input,
          output: totalMeta.cost.output,
        },
        latency: totalMeta.latency,
        model: ctx.modelId,
        tokens: {
          input: totalMeta.tokens.input,
          output: totalMeta.tokens.output,
        },
      },
    })
  }

  return allRatings
}

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

  return new Response<Array<RatingResult<I>>, Array<SimplifiedRatingResult<I>>>(
    context,
    rate(input, instructions, _options, context),
    (result) => {
      // If instructions is a string, simplify to just the total number
      if (typeof instructions === 'string') {
        return result.map((r) => r.total as SimplifiedRatingResult<I>)
      }
      // Otherwise return the full result (including total)
      return result as Array<SimplifiedRatingResult<I>>
    }
  )
}
