import { db } from "../db/index";
import { courses, uploads, topics, lessons } from "../db/schema";
import { eq } from "drizzle-orm";
import { askAI } from "./openrouter";
import { generateTTSToPath } from "./generateTTS";

function log(lessonId: string, msg: string) {
  console.log(`[branches:${lessonId.slice(0, 8)}] ${msg}`);
}

function parseJSON<T>(raw: string): T {
  let text = raw
    .replace(/<think>[\s\S]*?<\/think>/gi, "")
    .replace(/^\s*thought\s*\n/i, "")
    .trim();
  try {
    return JSON.parse(text);
  } catch {
    const cleaned = text.replace(/^```[a-z]*\n?/i, "").replace(/\n?```$/i, "").trim();
    return JSON.parse(cleaned);
  }
}

export async function generateBranches(topicId: string, lessonId: string): Promise<void> {
  try {
    await db.update(lessons).set({ branchStatus: "generating" }).where(eq(lessons.id, lessonId));

    const lesson = await db.query.lessons.findFirst({ where: eq(lessons.id, lessonId) });
    if (!lesson) throw new Error(`Lesson ${lessonId} not found`);

    const topic = await db.query.topics.findFirst({ where: eq(topics.id, topicId) });
    if (!topic) throw new Error(`Topic ${topicId} not found`);

    const course = await db.query.courses.findFirst({ where: eq(courses.id, topic.courseId) });
    if (!course) throw new Error(`Course ${topic.courseId} not found`);

    const lessonContent: { keyConcepts: string[]; analogiesUsed: string[]; steps: any[] } = JSON.parse(lesson.content);
    const courseSubject = course.subject;

    // load source material for branch prompts
    const topicRelevantFiles: string[] = (() => {
      try { return JSON.parse(topic.relevantFiles ?? "[]"); } catch { return []; }
    })();

    const uploadRows = await db.query.uploads.findMany({
      where: eq(uploads.courseId, topic.courseId),
    });

    const relevantUploads = topicRelevantFiles.length > 0
      ? uploadRows.filter((u) => topicRelevantFiles.includes(u.filename) && u.extractedText)
      : uploadRows.filter((u) => (u.type === "past_paper" || u.type === "lab_worksheet") && u.extractedText);

    const primaryTextForLesson = relevantUploads
      .map((u) => `--- ${u.filename} ---\n${u.extractedText}`)
      .join("\n\n");

    let costBranchAI = 0;
    let costBranchAudio = 0;

    let branchErrors = 0;
    let branchSuccesses = 0;

    for (let si = 0; si < lessonContent.steps.length; si++) {
      const step = lessonContent.steps[si] as any;
      if (step.type !== "question") continue;

      const OPTION_LABELS = ["A", "B", "C", "D"];
      const optionsText = (step.options as string[])
        .map((o: string, idx: number) => `${OPTION_LABELS[idx]}: ${o}`)
        .join("\n");

      const branchPrompt = `You are writing remediation branches for a question in a lesson about ${courseSubject}.

SOURCE MATERIAL FOR THIS TOPIC:
${primaryTextForLesson || "(none)"}

THE QUESTION:
${step.body}

THE OPTIONS:
${optionsText}

THE CORRECT ANSWER: ${step.answer}

YOUR TASK:
For each wrong answer option, write a branch that:
1. Opens by acknowledging the specific thinking behind that wrong answer — not generically ("that's wrong") but specifically ("It makes sense to think that, because X — but here is where that reasoning breaks down...")
2. Contains 1-2 short teaching steps that directly address the specific misconception behind that wrong answer
3. Ends with a confirming question that tests whether the student now understands the correct concept
4. Includes a warm, honest hard-stop message for if they fail the confirming question too

BRANCH TEACHING RULES:
- Each branch must feel like a patient teacher who has seen this exact mistake before and knows exactly how to fix it
- The opening must name the misconception specifically — never say "incorrect" or "wrong" — say "here is what that answer is actually describing..." or "that reasoning would be right if... but in this case..."
- Steps must be 2-3 sentences maximum — these are micro-lessons, not full explanations
- The confirming question must be simpler than the original question — it tests the core concept only
- The confirming question options must be short and scannable — under 12 words each
- The hard stop message must be warm, not discouraging — acknowledge that some concepts need more time, tell them specifically what to look up, and invite them back
- Never use technical terms that haven't been taught yet

Return only valid JSON, no markdown. Structure:
{
  "branches": {
    "{exact text of wrong option}": {
      "steps": [
        { "type": "concept", "title": "...", "body": "..." }
      ],
      "confirmQuestion": {
        "body": "...",
        "options": ["...", "...", "...", "..."],
        "answer": "full correct answer text",
        "explanation": "..."
      },
      "hardStop": "..."
    }
  }
}

Only generate branches for the 3 wrong options. Do not generate a branch for the correct answer.`;

      try {
        const branchResult = await askAI([{ role: "user", content: branchPrompt }]);
        costBranchAI += branchResult.cost;

        const parsed = parseJSON<{ branches: Record<string, any> }>(branchResult.text);
        step.branches = parsed.branches ?? {};

        log(lessonId, `  step ${si} branches generated — ${Object.keys(step.branches).length} wrong options`);

        const wrongOptions = (step.options as string[]).filter((o: string) => o !== step.answer);

        // batch TTS for all branches in this step, 4 at a time
        type TTSTask = () => Promise<void>;
        const ttsTasks: TTSTask[] = [];

        for (let bi = 0; bi < wrongOptions.length; bi++) {
          const wrongOpt = wrongOptions[bi];
          const branch = step.branches[wrongOpt];
          if (!branch) continue;

          for (let bsi = 0; bsi < (branch.steps ?? []).length; bsi++) {
            const bStep = branch.steps[bsi];
            const text = [bStep.body, bStep.callout].filter(Boolean).join(" ");
            if (!text.trim()) continue;

            const filename = `branch_${si}_${bi}_step_${bsi}.mp3`;
            ttsTasks.push(async () => {
              const r = await generateTTSToPath(text, lessonId, filename);
              if (r) {
                bStep.audioPath = r.audioPath;
                bStep.audioChunks = r.audioChunks;
                costBranchAudio += r.cost;
              }
            });
          }

          if (branch.confirmQuestion?.body?.trim()) {
            const cqBody = branch.confirmQuestion.body;
            const cqFile = `branch_${si}_${bi}_confirm_q.mp3`;
            ttsTasks.push(async () => {
              const r = await generateTTSToPath(cqBody, lessonId, cqFile);
              if (r) {
                branch.confirmQuestion.questionAudioPath = r.audioPath;
                branch.confirmQuestion.questionAudioChunks = r.audioChunks;
                costBranchAudio += r.cost;
              }
            });
          }

          if (Array.isArray(branch.confirmQuestion?.options)) {
            branch.confirmQuestion.optionAudioPaths = new Array(branch.confirmQuestion.options.length).fill(null);

            for (let oi = 0; oi < branch.confirmQuestion.options.length; oi++) {
              const optText = branch.confirmQuestion.options[oi];
              const oiCopy = oi;
              const optFile = `branch_${si}_${bi}_confirm_opt_${oi}.mp3`;

              if (optText?.trim()) {
                ttsTasks.push(async () => {
                  const r = await generateTTSToPath(optText, lessonId, optFile);
                  branch.confirmQuestion.optionAudioPaths[oiCopy] = r ? r.audioPath : null;
                  if (r) costBranchAudio += r.cost;
                });
              }
            }
          }

          if (branch.hardStop?.trim()) {
            const hsText = branch.hardStop;
            const hsFile = `branch_${si}_${bi}_hardstop.mp3`;
            ttsTasks.push(async () => {
              const r = await generateTTSToPath(hsText, lessonId, hsFile);
              if (r) {
                branch.hardStopAudioPath = r.audioPath;
                costBranchAudio += r.cost;
              }
            });
          }
        }

        await Promise.all(ttsTasks.map((fn) => fn()));

        log(lessonId, `  step ${si} branch TTS done`);
        branchSuccesses++;
      } catch (err: any) {
        console.error(`[branches] step ${si} failed for lesson ${lessonId}: ${err?.message ?? err}`);
        branchErrors++;
      }
    }

    const totalQuestionSteps = lessonContent.steps.filter((s: any) => s.type === "question").length;
    const branchStatus = branchErrors > 0 ? "error" : "ready";

    if (branchSuccesses > 0 || totalQuestionSteps === 0) {
      const existing = await db.query.lessons.findFirst({ where: eq(lessons.id, lessonId) });
      const prevCostAI = existing?.costAI ?? 0;
      const prevCostAudio = existing?.costAudio ?? 0;

      await db.update(lessons)
        .set({
          content: JSON.stringify(lessonContent),
          costBranchAI,
          costBranchAudio,
          costTotal: prevCostAI + prevCostAudio + costBranchAI + costBranchAudio,
          branchStatus,
        })
        .where(eq(lessons.id, lessonId));

      if (branchErrors > 0) {
        log(lessonId, `branches done with errors — ${branchSuccesses} ok, ${branchErrors} failed`);
      } else {
        log(lessonId, `✓ branches ready — AI $${costBranchAI.toFixed(4)}, audio $${costBranchAudio.toFixed(4)}`);
      }
    } else {
      await db.update(lessons).set({ branchStatus }).where(eq(lessons.id, lessonId));

      if (totalQuestionSteps === 0) {
        log(lessonId, "no question steps found, branch_status set to ready");
      } else {
        log(lessonId, "all branch steps failed, branch_status set to error");
      }
    }
  } catch (err: any) {
    console.error(`[branches:${lessonId.slice(0, 8)}] ✗ failed: ${err?.message ?? err}`);
    await db.update(lessons).set({ branchStatus: "error" }).where(eq(lessons.id, lessonId));
  }
}