/**
 * AI 响应解析与 Zod 校验
 *
 * 职责：
 * - JSON 提取与修复（从 AI 返回的原始文本中提取合法 JSON）
 * - Zod schema 定义（题目、试卷、结构等）
 * - 共享类型导出
 */

import { z } from "zod"
import { createAiChatCompletion } from "@/shared/lib/ai"
import { env } from "@/env.mjs"

// ---------------------------------------------------------------------------
// Zod schemas
// ---------------------------------------------------------------------------

const AiSubQuestionSchema = z.object({
  id: z.string().min(1).optional(),
  text: z.string().min(1),
  answer: z.string().min(1).optional(),
  score: z.coerce.number().int().min(0).optional(),
})

const AiQuestionContentSchema = z.object({
  text: z.string().min(1),
  options: z
    .array(
      z.object({
        id: z.string().min(1).optional(),
        text: z.string().min(1),
        isCorrect: z.boolean().optional(),
      })
    )
    .optional(),
  subQuestions: z.array(AiSubQuestionSchema).optional(),
})

export const AiQuestionSchema = z.object({
  type: z.enum(["single_choice", "multiple_choice", "text", "judgment"]),
  difficulty: z.coerce.number().int().min(1).max(5).optional(),
  score: z.coerce.number().int().min(0).optional(),
  content: AiQuestionContentSchema,
})

export const AiInsertQuestionSchema = z.object({
  id: z.string().min(1),
  type: z.enum(["single_choice", "multiple_choice", "text", "judgment"]),
  difficulty: z.coerce.number().int().min(1).max(5),
  score: z.coerce.number().int().min(0),
  content: AiQuestionContentSchema.extend({
    options: z
      .array(
        z.object({
          id: z.string().min(1),
          text: z.string().min(1),
          isCorrect: z.boolean().optional(),
        })
      )
      .optional(),
    subQuestions: z.array(
      AiSubQuestionSchema.extend({
        id: z.string().min(1),
      })
    ).optional(),
  }),
})

const AiSectionSchema = z.object({
  title: z.string().min(1),
  questions: z.array(AiQuestionSchema).min(1),
})

export const AiExamResponseSchema = z.object({
  title: z.string().optional(),
  questions: z.array(AiQuestionSchema).optional(),
  sections: z.array(AiSectionSchema).optional(),
})

const AiStructureQuestionSchema = z.object({
  text: z.string().min(1),
  score: z.coerce.number().int().min(0).optional(),
})

const AiStructureSectionSchema = z.object({
  title: z.string().min(1),
  questions: z.array(AiStructureQuestionSchema).min(1),
})

export const AiStructureResponseSchema = z.object({
  title: z.string().optional(),
  sections: z.array(AiStructureSectionSchema).optional(),
  questions: z.array(AiStructureQuestionSchema).optional(),
})

const AiSourceValidationSchema = z.object({
  valid: z.boolean(),
  reason: z.string().optional(),
})

// ---------------------------------------------------------------------------
// Structure node schema (recursive)
// ---------------------------------------------------------------------------

export type AiGeneratedStructureNode = {
  id: string
  type: "section" | "group" | "question"
  title?: string
  instruction?: string
  level?: number
  questionId?: string
  score?: number
  children?: AiGeneratedStructureNode[]
}

export const AiGeneratedStructureNodeSchema: z.ZodType<AiGeneratedStructureNode> = z.lazy(() => z.object({
  id: z.string().min(1),
  type: z.enum(["section", "group", "question"]),
  title: z.string().optional(),
  instruction: z.string().optional(),
  level: z.coerce.number().int().min(0).optional(),
  questionId: z.string().optional(),
  score: z.coerce.number().int().min(0).optional(),
  children: z.array(AiGeneratedStructureNodeSchema).optional(),
}))

export const AiGeneratedStructureSchema = z.array(AiGeneratedStructureNodeSchema)

// ---------------------------------------------------------------------------
// JSON extraction & repair
// ---------------------------------------------------------------------------

const sanitizeJsonCandidate = (value: string): string => value
  .replace(/\[\s*\.\.\.\s*\]/g, "[]")
  .replace(/\{\s*\.\.\.\s*\}/g, "{}")
  .trim()

const tryParseJson = (value: string): unknown | null => {
  const sanitized = sanitizeJsonCandidate(value)
  if (!sanitized) return null
  try {
    return JSON.parse(sanitized)
  } catch {
    return null
  }
}

const extractBalancedJsonSegment = (value: string): string | null => {
  const startBrace = value.indexOf("{")
  const startBracket = value.indexOf("[")
  const start =
    startBrace === -1
      ? startBracket
      : startBracket === -1
        ? startBrace
        : Math.min(startBrace, startBracket)
  if (start === -1) return null
  const opening = value[start]
  const closing = opening === "{" ? "}" : "]"
  let depth = 0
  let inString = false
  let escaped = false
  for (let i = start; i < value.length; i += 1) {
    const char = value[i]
    if (inString) {
      if (escaped) {
        escaped = false
      } else if (char === "\\") {
        escaped = true
      } else if (char === "\"") {
        inString = false
      }
      continue
    }
    if (char === "\"") {
      inString = true
      continue
    }
    if (char === opening) {
      depth += 1
      continue
    }
    if (char === closing) {
      depth -= 1
      if (depth === 0) {
        return value.slice(start, i + 1)
      }
    }
  }
  return null
}

const extractJson = (raw: string): unknown => {
  const trimmed = raw.trim()
  const candidates: string[] = []
  const fencedMatches = [...trimmed.matchAll(/```(?:json)?\s*([\s\S]*?)```/ig)]
  if (fencedMatches.length > 0) {
    candidates.push(...fencedMatches.map((match) => (match[1] ?? "").trim()))
  }
  candidates.push(trimmed)
  for (const candidate of candidates) {
    const direct = tryParseJson(candidate)
    if (direct !== null) return direct
    const segment = extractBalancedJsonSegment(candidate)
    if (!segment) continue
    const parsed = tryParseJson(segment)
    if (parsed !== null) return parsed
  }
  throw new Error("Invalid AI response")
}

const AI_JSON_REPAIR_PROMPT = [
  "You are a JSON repair engine.",
  "Fix the provided invalid JSON into valid JSON only.",
  "Keep the original structure and values as much as possible.",
  "Do not use placeholders such as ... or [...].",
  "Return JSON only without markdown.",
].join("\n")

const repairJson = async (raw: string, providerId?: string) => {
  const aiResult = await createAiChatCompletion({
    model: String(env.AI_MODEL ?? "gpt-4o-mini"),
    providerId,
    messages: [
      { role: "system" as const, content: AI_JSON_REPAIR_PROMPT },
      { role: "user" as const, content: raw },
    ],
    temperature: 0,
    maxTokens: 4000,
  })
  return extractJson(aiResult.content)
}

export const parseAiResponse = async (raw: string, providerId?: string) => {
  try {
    return extractJson(raw)
  } catch {
    return repairJson(raw, providerId)
  }
}

// ---------------------------------------------------------------------------
// Shared types
// ---------------------------------------------------------------------------

export type QuestionContentResult = {
  text: string
  options?: Array<{ id: string; text: string; isCorrect: boolean }>
  subQuestions?: Array<{ id: string; text: string; answer?: string; score?: number }>
}

export type AiPreviewQuestion = {
  id: string
  type: z.infer<typeof AiQuestionSchema>["type"]
  difficulty: number
  score: number
  content: QuestionContentResult
}

export type AiPreviewData = {
  title: string
  rawOutput?: string
  sections?: Array<{
    id: string
    title: string
    questions: AiPreviewQuestion[]
  }>
  questions?: AiPreviewQuestion[]
}

export type AiRewriteQuestionData = {
  type: z.infer<typeof AiQuestionSchema>["type"]
  difficulty: number
  score: number
  content: QuestionContentResult
}

export type AiGeneratedQuestion = {
  id: string
  type: z.infer<typeof AiQuestionSchema>["type"]
  difficulty: number
  score: number
  content: QuestionContentResult
}

// ---------------------------------------------------------------------------
// Pure transformation functions
// ---------------------------------------------------------------------------

export const normalizeScores = (scores: number[], totalScore: number): number[] => {
  if (scores.length === 0) return []
  const sum = scores.reduce((acc, s) => acc + s, 0)
  if (sum <= 0) {
    const base = Math.floor(totalScore / scores.length)
    const remainder = totalScore - base * scores.length
    return scores.map((_, idx) => base + (idx < remainder ? 1 : 0))
  }
  const scaled = scores.map((s) => Math.max(0, Math.round((s / sum) * totalScore)))
  let diff = totalScore - scaled.reduce((acc, s) => acc + s, 0)
  let i = 0
  while (diff !== 0 && i < scaled.length * 2) {
    const idx = i % scaled.length
    if (diff > 0) {
      scaled[idx] += 1
      diff -= 1
    } else if (scaled[idx] > 0) {
      scaled[idx] -= 1
      diff += 1
    }
    i += 1
  }
  return scaled
}

export const buildQuestionContent = (q: z.infer<typeof AiQuestionSchema>): QuestionContentResult => {
  const base = { text: q.content.text }
  const subQuestions = Array.isArray(q.content.subQuestions)
    ? q.content.subQuestions.map((item, index) => ({
      id: item.id ?? String(index + 1),
      text: item.text,
      answer: item.answer,
      score: item.score,
    }))
    : []
  if (q.type === "single_choice" || q.type === "multiple_choice") {
    const options = (q.content.options ?? []).map((opt, idx) => ({
      id: opt.id ?? String.fromCharCode(65 + idx),
      text: opt.text,
      isCorrect: opt.isCorrect ?? false,
    }))
    if (options.length > 0 && subQuestions.length > 0) return { ...base, options, subQuestions }
    if (options.length > 0) return { ...base, options }
    if (subQuestions.length > 0) return { ...base, subQuestions }
    return base
  }
  if (subQuestions.length > 0) return { ...base, subQuestions }
  return base
}

// ---------------------------------------------------------------------------
// Prompts
// ---------------------------------------------------------------------------

export const AI_EXAM_SYSTEM_PROMPT = [
  "You are an exam parsing engine.",
  "Parse the provided exam text and output JSON only.",
  "Allowed question types: single_choice, multiple_choice, judgment, text.",
  "Preserve the original order and sectioning if present.",
  "Escape double quotes inside string values.",
  "Output schema:",
  "{",
  '  "sections": [',
  '    { "title": "Section Title", "questions": [',
  '      { "type": "single_choice", "difficulty": 1, "score": 5, "content": { "text": "...", "options": [ { "id": "A", "text": "...", "isCorrect": true } ] } }',
  "    ] }",
  "  ]",
  "}",
  "For grouped blanks or one prompt with multiple small questions, keep one parent question and place each child item into content.subQuestions.",
  'content.subQuestions item schema: { "id": "1", "text": "lǎn duò（ ）", "answer": "懒惰", "score": 1 }',
  "If you do not need sections, return { \"questions\": [] } or include real question items.",
  "Never output placeholders like ..., [...], or {...}.",
  "Return JSON only without markdown.",
].join("\n")

export const AI_REWRITE_QUESTION_SYSTEM_PROMPT = [
  "You are a question rewriting engine.",
  "Rewrite exactly one question based on teacher instruction.",
  "Return JSON only without markdown.",
  "Allowed question types: single_choice, multiple_choice, judgment, text.",
  "Output schema:",
  "{",
  '  "type": "single_choice | multiple_choice | judgment | text",',
  '  "difficulty": 1,',
  '  "score": 5,',
  '  "content": { "text": "...", "options": [ { "id": "A", "text": "...", "isCorrect": true } ], "subQuestions": [ { "id": "1", "text": "...", "answer": "...", "score": 1 } ] }',
  "}",
  "For judgment/text, options can be omitted. Keep subQuestions when original question has multiple child items.",
  "Never output placeholders like ..., [...], or {...}.",
].join("\n")

export const AI_EXAM_STRUCTURE_SYSTEM_PROMPT = [
  "You are an exam splitter engine.",
  "Split the provided exam text into ordered question units quickly.",
  "Do not deeply analyze choices or answers in this step.",
  "Keep original sectioning and question order.",
  "If one stem contains multiple numbered sub-items, keep them in one question unit and include all sub-items in the same text.",
  "Do not split one parent question into several child-only units.",
  "Output JSON only.",
  "Output schema:",
  "{",
  '  "title": "Optional title",',
  '  "sections": [',
  '    { "title": "Section Title", "questions": [',
  '      { "text": "Original full question text", "score": 5 }',
  "    ] }",
  "  ]",
  "}",
  "If no sections, return:",
  '{ "questions": [ { "text": "Original full question text", "score": 5 } ] }',
  "Never output placeholders like ..., [...], or {...}.",
].join("\n")

export const AI_EXAM_SOURCE_VALIDATION_SYSTEM_PROMPT = [
  "You are an exam text validator.",
  "Judge whether the input text is readable and likely a normal exam/question text.",
  "Reject garbled text, random symbols, severely disordered fragments, or meaningless content.",
  "Do not require strict section formatting. Focus only on readability and whether it resembles exam questions.",
  "Return JSON only without markdown.",
  "Output schema:",
  '{ "valid": true, "reason": "short reason" }',
].join("\n")

export const AI_QUESTION_DETAIL_SYSTEM_PROMPT = [
  "You are an exam question detail parser.",
  "Given one split question text, output one structured question JSON only.",
  "Allowed question types: single_choice, multiple_choice, judgment, text.",
  "For one stem with multiple child sub-items, keep one parent content.text and place child items in content.subQuestions.",
  "Use exact key name content.subQuestions (camelCase).",
  "Output schema:",
  "{",
  '  "type": "single_choice | multiple_choice | judgment | text",',
  '  "difficulty": 1,',
  '  "score": 5,',
  '  "content": { "text": "...", "options": [ { "id": "A", "text": "...", "isCorrect": true } ], "subQuestions": [ { "id": "1", "text": "...", "answer": "...", "score": 1 } ] }',
  "}",
  "For judgment/text, options can be omitted.",
  "Never output placeholders like ..., [...], or {...}.",
].join("\n")

export { AiSourceValidationSchema }
