padhle/backend/src/routes/chat.js

import express from "express";
import { RECENT_LIMIT, streamChatResponse, summarizeConversation } from "../services/ai.js";
import { createSession, addMessage, getMessages, getSession, listSessions, setSummary } from "../services/db.js";
import { isOwnedSession, validateChatInput, validateChapter } from "./chatValidation.js";
import trialStore from "../services/anonTrial.js";
import { sseWrite } from "../services/sse.js";
import { logError } from "../services/logRedact.js";

/**
 * Router factory. `streamFn` is injected so integration tests can substitute a
 * fake AI streamer (defaults to the real one). Returns a configured router.
 */
export function createChatRouter({ streamFn = streamChatResponse } = {}) {
  const router = express.Router();

  /**
   * Anonymous chat handler — the 5-free-message trial.
   *
   * Flow (enforced server-side):
   *  1. Validate input first (same rules as signed-in).
   *  2. Resolve the trial: `chatId` continues an existing in-memory trial;
   *     anything else (missing/unknown/forged id) starts a fresh trial — the
   *     "refresh loses the conversation, budget resets" behavior. A forged id
   *     is never treated as a session id and never touches the database.
   *  3. If the trial budget is spent → 403 SIGNIN_REQUIRED before any AI call.
   *  4. Otherwise consume one user message, stream the reply, and emit
   *     `session` (trial id) + `limit` (remaining) SSE events.
   */
  async function handleAnonChat(req, res, { text, chatId, grade, subject, chapter }) {
    let trial = chatId ? trialStore.getTrial(chatId) : null;
    if (!trial) trial = trialStore.createTrial(); // unknown id → fresh trial

    if (trialStore.isExhausted(trial.id)) {
      return res
        .status(403)
        .json({ error: "Please sign in to continue", code: "SIGNIN_REQUIRED" });
    }

    // Consume one free message (only user messages count).
    trialStore.addMessage(trial.id, "user", text.trim());
    const remaining = trialStore.countRemaining(trial.id);

    res.setHeader("Content-Type", "text/event-stream");
    res.setHeader("Cache-Control", "no-cache");
    res.setHeader("Connection", "keep-alive");
    res.setHeader("X-Accel-Buffering", "no");

    // Same contract as the signed-in path: the client stores this id as
    // activeChat and echoes it back as chatId to continue the thread.
    sseWrite(res, { type: "session", id: trial.id });
    sseWrite(res, { type: "limit", remaining });

    const context = trialStore.getContextMessages(trial.id).map((m) => ({ role: m.role, text: m.text }));
    let assistantText = "";
    await streamFn(
      context,
      grade,
      subject,
      chapter,
      (chunk) => {
        assistantText += chunk;
        sseWrite(res, { type: "chunk", content: chunk });
      },
      () => {
        sseWrite(res, { type: "error", message: "AI service error" });
        res.end();
      },
      undefined // anonymous trials have no rolling summary
    );

    if (!res.writableEnded) {
      if (assistantText) trialStore.addMessage(trial.id, "assistant", assistantText);
      sseWrite(res, { type: "done" });
      res.end();
    }
  }

  router.post("/", async (req, res) => {
    const { text, chatId, grade, subject, chapter } = req.body;
    const validation = validateChatInput({ text, grade, subject });
    if (!validation.ok) return res.status(400).json({ error: validation.error });
    const chapterCheck = validateChapter(chapter);
    if (!chapterCheck.ok) return res.status(400).json({ error: chapterCheck.error });

    // Anonymous (no verified JWT) → free-trial path.
    if (!req.user) {
      try {
        return await handleAnonChat(req, res, { text, chatId, grade, subject, chapter });
      } catch (err) {
        logError("Anonymous chat error", err);
        if (!res.headersSent) res.status(500).json({ error: "Internal server error" });
        else {
          sseWrite(res, { type: "error", message: "Internal server error" });
          res.end();
        }
        return;
      }
    }

    const userId = req.user.uid;
    try {
      let session;
      if (chatId) {
        session = await getSession(chatId, req.token);
        if (!isOwnedSession(session, userId)) {
          return res.status(404).json({ error: "Session not found" });
        }
      } else {
        session = await createSession(userId, grade, subject, chapter, req.token);
      }

      const currentChatId = session.id;
      // Persist the user message BEFORE streaming so it is included in the AI
      // context (the model must see the current question) and survives even if
      // the AI call fails.
      await addMessage(currentChatId, "user", text.trim(), req.token);
      const chatMessages = await getMessages(currentChatId, req.token);
      res.setHeader("Content-Type", "text/event-stream");
      res.setHeader("Cache-Control", "no-cache");
      res.setHeader("Connection", "keep-alive");
      res.setHeader("X-Accel-Buffering", "no");

      // Tell the client which session this conversation belongs to, so follow-ups
      // reuse the same session and history accumulates.
      sseWrite(res, { type: "session", id: currentChatId });

      // Bounded context: fold older messages into a persistent summary, keep the
      // last RECENT_LIMIT raw.
      let summary = session.summary || "";
      if (chatMessages.length > RECENT_LIMIT) {
        const overflow = chatMessages.slice(0, chatMessages.length - RECENT_LIMIT);
        try {
          summary = await summarizeConversation(session.summary, overflow);
          await setSummary(currentChatId, summary, req.token);
        } catch (err) {
          logError("Summarize error", err);
        }
      }

      let assistantText = "";
      await streamFn(
        chatMessages,
        grade || session.grade,
        subject || session.subject,
        chapter || session.chapter,
        (chunk) => {
          assistantText += chunk;
          sseWrite(res, { type: "chunk", content: chunk });
        },
        () => {
          sseWrite(res, { type: "error", message: "AI service error" });
          res.end();
        },
        summary
      );

      if (!res.writableEnded) {
        sseWrite(res, { type: "done" });
        await addMessage(currentChatId, "assistant", assistantText, req.token);
      }
    } catch (err) {
      logError("Chat error", err);
      if (!res.headersSent) res.status(500).json({ error: "Internal server error" });
      else {
        sseWrite(res, { type: "error", message: "Internal server error" });
        res.end();
      }
    } finally {
      res.end();
    }
  });

  // Static route must precede /:chatId.
  router.get("/sessions", async (req, res) => {
    if (!req.user) return res.status(401).json({ error: "Authentication required" });
    try {
      res.json(await listSessions(req.user.uid, req.token));
    } catch (err) {
      logError("List sessions error", err);
      res.status(500).json({ error: "Failed to load sessions" });
    }
  });

  router.get("/:chatId", async (req, res) => {
    if (!req.user) return res.status(401).json({ error: "Authentication required" });
    try {
      const { chatId } = req.params;
      const session = await getSession(chatId, req.token);
      if (!isOwnedSession(session, req.user.uid)) {
        return res.status(404).json({ error: "Session not found" });
      }
      const msgHistory = await getMessages(chatId, req.token);
      res.json({
        session: {
          id: session.id,
          grade: session.grade,
          subject: session.subject,
          chapter: session.chapter,
          createdAt: session.created_at,
          updatedAt: session.updated_at,
        },
        messages: msgHistory,
      });
    } catch (err) {
      logError("Get chat error", err);
      res.status(500).json({ error: "Failed to load session" });
    }
  });

  return router;
}

export default createChatRouter();