qbank/internal/handlers/upload.go

package handlers

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/go-chi/chi/v5"

	"qbank/internal/auth"
	"qbank/internal/db"
	"qbank/internal/llm"
	"qbank/internal/models"
	"qbank/internal/parse"
)

type UploadHandler struct {
	auth    *auth.Manager
	repo    *db.Repo
	llm     *llm.Client
	render  *Renderer
	dataDir string
}

func NewUploadHandler(a *auth.Manager, repo *db.Repo, llmClient *llm.Client, r *Renderer, dataDir string) *UploadHandler {
	return &UploadHandler{auth: a, repo: repo, llm: llmClient, render: r, dataDir: dataDir}
}

func (h *UploadHandler) UploadGet(w http.ResponseWriter, r *http.Request) {
	data := BaseData(h.auth, r)
	h.render.Render(w, http.StatusOK, "upload", data)
}

func (h *UploadHandler) UploadPost(w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, 20<<20) // 20 MB

	renderErr := func(msg string) {
		data := BaseData(h.auth, r)
		data["Error"] = msg
		h.render.Render(w, http.StatusUnprocessableEntity, "upload", data)
	}

	if !h.auth.CheckCSRF(r) {
		HTTPError(w, http.StatusForbidden)
		return
	}
	if err := r.ParseMultipartForm(20 << 20); err != nil {
		renderErr("File too large (max 20 MB).")
		return
	}

	file, header, err := r.FormFile("file")
	if err != nil {
		renderErr("No file selected.")
		return
	}
	defer file.Close()

	ext := strings.ToLower(filepath.Ext(header.Filename))
	if ext != ".pdf" && ext != ".docx" {
		renderErr("Unsupported file type. Upload a .pdf or .docx file.")
		return
	}

	// Read file once; use for both saving and extraction.
	data, err := io.ReadAll(file)
	if err != nil {
		renderErr("Failed to read uploaded file.")
		return
	}

	// Save to disk.
	uploadsDir := filepath.Join(h.dataDir, "uploads")
	if err := os.MkdirAll(uploadsDir, 0755); err != nil {
		renderErr("Server error: could not create uploads directory.")
		return
	}
	safeName := filepath.Base(header.Filename)
	savedPath := filepath.Join(uploadsDir, fmt.Sprintf("%d_%s", time.Now().UnixMilli(), safeName))
	if err := os.WriteFile(savedPath, data, 0644); err != nil {
		renderErr("Server error: could not save file.")
		return
	}

	// Extract text.
	var text string
	switch ext {
	case ".pdf":
		text, err = parse.ExtractPDF(bytes.NewReader(data))
		if errors.Is(err, parse.ErrScanPDF) {
			renderErr("This PDF appears to be image-only (scan-based). Please convert it to a text PDF first.")
			return
		}
	case ".docx":
		text, err = parse.ExtractDOCX(bytes.NewReader(data))
	}
	if err != nil {
		renderErr("Could not extract text from the document: " + err.Error())
		return
	}

	// Chunk and call LLM.
	chunks := parse.Chunk(text, 10_000)
	seen := make(map[string]bool)
	var draftQs []models.DraftQuestion

	for _, chunk := range chunks {
		qs, err := h.llm.ExtractQuestions(r.Context(), chunk)
		if err != nil {
			slog.Warn("llm chunk failed", "err", err)
			continue
		}
		for _, q := range qs {
			key := db.QuestionID(q.Question)
			if seen[key] {
				continue
			}
			seen[key] = true
			dq := models.DraftQuestion{Text: q.Question}
			for _, a := range q.Answers {
				dq.Answers = append(dq.Answers, models.DraftAnswer{Text: a.Text, IsCorrect: a.Correct})
			}
			draftQs = append(draftQs, dq)
		}
	}

	source := strings.TrimSpace(r.FormValue("source"))
	if source == "" {
		source = strings.TrimSuffix(safeName, ext)
	}

	userID := auth.UserFromCtx(r.Context()).ID
	draftID, err := h.repo.CreateDraft(userID, source, draftQs)
	if err != nil {
		renderErr("Server error: could not create import draft.")
		return
	}

	http.Redirect(w, r, "/import/"+draftID, http.StatusSeeOther)
}

func (h *UploadHandler) ImportGet(w http.ResponseWriter, r *http.Request) {
	draftID := chi.URLParam(r, "id")
	userID := auth.UserFromCtx(r.Context()).ID

	draft, err := h.repo.GetDraftForUser(draftID, userID)
	if err != nil {
		HTTPError(w, http.StatusNotFound)
		return
	}

	data := BaseData(h.auth, r)
	data["Draft"] = draft
	h.render.Render(w, http.StatusOK, "import", data)
}

func (h *UploadHandler) ImportPost(w http.ResponseWriter, r *http.Request) {
	draftID := chi.URLParam(r, "id")
	userID := auth.UserFromCtx(r.Context()).ID

	if !h.auth.CheckCSRF(r) {
		HTTPError(w, http.StatusForbidden)
		return
	}

	draft, err := h.repo.GetDraftForUser(draftID, userID)
	if err != nil {
		HTTPError(w, http.StatusNotFound)
		return
	}

	source := strings.TrimSpace(r.FormValue("source"))

	var imported, skipped int
	for i, dq := range draft.Questions {
		if r.FormValue(fmt.Sprintf("delete_%d", i)) == "on" {
			skipped++
			continue
		}

		text := strings.TrimSpace(r.FormValue(fmt.Sprintf("q_text_%d", i)))
		if text == "" {
			skipped++
			continue
		}

		correctIdx := r.FormValue(fmt.Sprintf("correct_%d", i))
		var answers []*models.Answer
		for j := range dq.Answers {
			aText := strings.TrimSpace(r.FormValue(fmt.Sprintf("a_text_%d_%d", i, j)))
			if aText == "" {
				continue
			}
			answers = append(answers, &models.Answer{
				Text:      aText,
				IsCorrect: fmt.Sprintf("%d", j) == correctIdx,
			})
		}

		var nCorrect int
		for _, a := range answers {
			if a.IsCorrect {
				nCorrect++
			}
		}
		if len(answers) < 2 || nCorrect != 1 {
			skipped++
			continue
		}

		q := &models.Question{Text: text, Source: source}
		if err := h.repo.InsertQuestion(q, answers); err != nil {
			slog.Error("insert question", "err", err)
			skipped++
			continue
		}
		imported++
	}

	if err := h.repo.DeleteDraft(draftID); err != nil {
		slog.Error("delete draft", "id", draftID, "err", err)
	}

	h.auth.SetFlash(r, fmt.Sprintf("Imported %d question(s), %d skipped.", imported, skipped))
	http.Redirect(w, r, "/", http.StatusSeeOther)
}