package main

import (
	"bytes"
	"context"
	"crypto/rand"
	"database/sql"
	_ "embed"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"html/template"
	"image"
	_ "image/gif"
	"image/jpeg"
	_ "image/png"
	"io"
	"log"
	"math"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"time"

	"github.com/fogleman/gg"
	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	_ "github.com/mattn/go-sqlite3"
	"golang.org/x/image/font/basicfont"
)

//go:embed admin.html
var adminHTML string

const defaultTitlePrefix = "作业"

var adminTemplate = template.Must(template.New("admin").Parse(adminHTML))

type Config struct {
	Port          string
	DBPath        string
	LLMBaseURL    string
	LLMAPIKey     string
	LLMModel      string
	LLMTimeout    time.Duration
	LLMMaxRetries int
	UploadDir     string
	AdminUser     string
	AdminPass     string
}

type Assignment struct {
	ID                  int64     `json:"id"`
	Username            string    `json:"username"`
	Title               string    `json:"title"`
	OCRText             string    `json:"ocrText"`
	Markdown            string    `json:"markdown"`
	Feedback            string    `json:"feedback"`
	Score               int       `json:"score"`
	ImagePath           string    `json:"imagePath"`
	ImagePaths          []string  `json:"imagePaths"`
	CorrectedImagePath  string    `json:"correctedImagePath"`
	CorrectedImagePaths []string  `json:"correctedImagePaths"`
	CreatedAt           time.Time `json:"createdAt"`
}

type createAssignmentRequest struct {
	Username string `json:"username"`
	Title    string `json:"title"`
	OCRText  string `json:"ocrText"`
}

type LLMResult struct {
	Markdown    string   `json:"markdown"`
	Feedback    string   `json:"feedback"`
	Score       int      `json:"score"`
	Annotations []string `json:"annotations"`
}

type LLMImage struct {
	MimeType string
	Base64   string
}

type LLMClient struct {
	baseURL    string
	apiKey     string
	model      string
	httpClient *http.Client
	maxRetries int
}

func main() {
	cfg := loadConfig()

	db, err := initDB(cfg.DBPath)
	if err != nil {
		log.Fatalf("db init failed: %v", err)
	}
	defer db.Close()

	if err := os.MkdirAll(cfg.UploadDir, 0o755); err != nil {
		log.Fatalf("upload dir init failed: %v", err)
	}

	llmClient := &LLMClient{
		baseURL: cfg.LLMBaseURL,
		apiKey:  cfg.LLMAPIKey,
		model:   cfg.LLMModel,
		httpClient: &http.Client{
			Timeout: cfg.LLMTimeout,
		},
		maxRetries: cfg.LLMMaxRetries,
	}

	router := chi.NewRouter()
	router.Use(middleware.RequestID)
	router.Use(middleware.RealIP)
	router.Use(middleware.Logger)
	router.Use(middleware.Recoverer)
	router.Use(corsMiddleware)

	router.Get("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	})

	router.Route("/api", func(r chi.Router) {
		r.Get("/assignments", handleListAssignments(db))
		r.Get("/assignments/{id}", handleGetAssignment(db))
		r.Get("/assignments/{id}/image", handleGetAssignmentImage(db, cfg.UploadDir))
		r.Get("/assignments/{id}/corrected", handleGetAssignmentCorrectedImage(db, cfg.UploadDir))
		r.Post("/assignments", handleCreateAssignment(db, llmClient, cfg.UploadDir))
		r.Put("/assignments/{id}", handleUpdateAssignment(db, llmClient, cfg.UploadDir))
		r.Delete("/assignments/{id}", handleDeleteAssignment(db, cfg.UploadDir))
	})

	router.Route("/backend", func(r chi.Router) {
		r.Use(basicAuth(cfg.AdminUser, cfg.AdminPass))
		r.Get("/", handleAdminDashboard(db))
	})

	router.With(basicAuth(cfg.AdminUser, cfg.AdminPass)).Get("/backend", handleAdminDashboard(db))

	addr := ":" + cfg.Port
	log.Printf("backend listening on %s", addr)
	if err := http.ListenAndServe(addr, router); err != nil {
		log.Fatalf("server error: %v", err)
	}
}

func loadConfig() Config {
	timeoutSeconds := getEnvAsInt("LLM_TIMEOUT_SECONDS", 90)
	maxRetries := getEnvAsInt("LLM_MAX_RETRIES", 3)

	return Config{
		Port:          getEnv("BACKEND_PORT", "8080"),
		DBPath:        getEnv("DB_PATH", "/data/homework.db"),
		LLMBaseURL:    getEnv("LLM_BASE_URL", ""),
		LLMAPIKey:     getEnv("LLM_API_KEY", ""),
		LLMModel:      getEnv("LLM_MODEL", "qwen3-max"),
		LLMTimeout:    time.Duration(timeoutSeconds) * time.Second,
		LLMMaxRetries: maxRetries,
		UploadDir:     getEnv("UPLOAD_DIR", "/data/uploads"),
		AdminUser:     getEnv("ADMIN_USER", "admin"),
		AdminPass:     getEnv("ADMIN_PASS", "whoami139"),
	}
}

func initDB(path string) (*sql.DB, error) {
	db, err := sql.Open("sqlite3", path)
	if err != nil {
		return nil, err
	}
	if _, err := db.Exec(`
    CREATE TABLE IF NOT EXISTS assignments (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      username TEXT NOT NULL,
      title TEXT NOT NULL,
      ocr_text TEXT,
      markdown TEXT,
      feedback TEXT,
      score INTEGER,
      image_path TEXT,
      corrected_image_path TEXT,
      created_at DATETIME DEFAULT CURRENT_TIMESTAMP
    );
  `); err != nil {
		return nil, err
	}
	if err := ensureColumn(db, "assignments", "image_path", "TEXT"); err != nil {
		return nil, err
	}
	if err := ensureColumn(db, "assignments", "corrected_image_path", "TEXT"); err != nil {
		return nil, err
	}
	return db, nil
}

func handleListAssignments(db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		username := strings.TrimSpace(r.URL.Query().Get("username"))
		if username == "" {
			writeError(w, http.StatusBadRequest, "username is required")
			return
		}

		rows, err := db.Query(`
      SELECT id, username, title, score, image_path, corrected_image_path, created_at
      FROM assignments
      WHERE username = ?
      ORDER BY created_at DESC
    `, username)
		if err != nil {
			writeError(w, http.StatusInternalServerError, "failed to query assignments")
			return
		}
		defer rows.Close()

		assignments := make([]Assignment, 0)
		for rows.Next() {
			var a Assignment
			var rawImagePath string
			var rawCorrectedPath string
			if err := rows.Scan(&a.ID, &a.Username, &a.Title, &a.Score, &rawImagePath, &rawCorrectedPath, &a.CreatedAt); err != nil {
				writeError(w, http.StatusInternalServerError, "failed to read assignments")
				return
			}
			setImageFields(&a, rawImagePath)
			setCorrectedImageFields(&a, rawCorrectedPath)
			assignments = append(assignments, a)
		}

		writeJSON(w, http.StatusOK, assignments)
	}
}

func handleGetAssignment(db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		username := strings.TrimSpace(r.URL.Query().Get("username"))
		if username == "" {
			writeError(w, http.StatusBadRequest, "username is required")
			return
		}

		id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
		if err != nil {
			writeError(w, http.StatusBadRequest, "invalid assignment id")
			return
		}

		var a Assignment
		var rawImagePath string
		var rawCorrectedPath string
		row := db.QueryRow(`
      SELECT id, username, title, ocr_text, markdown, feedback, score, image_path, corrected_image_path, created_at
      FROM assignments
      WHERE id = ? AND username = ?
      LIMIT 1
    `, id, username)
		if err := row.Scan(&a.ID, &a.Username, &a.Title, &a.OCRText, &a.Markdown, &a.Feedback, &a.Score, &rawImagePath, &rawCorrectedPath, &a.CreatedAt); err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				writeError(w, http.StatusNotFound, "assignment not found")
				return
			}
			writeError(w, http.StatusInternalServerError, "failed to query assignment")
			return
		}
		setImageFields(&a, rawImagePath)
		setCorrectedImageFields(&a, rawCorrectedPath)

		writeJSON(w, http.StatusOK, a)
	}
}

func handleGetAssignmentImage(db *sql.DB, uploadDir string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		username := strings.TrimSpace(r.URL.Query().Get("username"))
		if username == "" {
			writeError(w, http.StatusBadRequest, "username is required")
			return
		}

		id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
		if err != nil {
			writeError(w, http.StatusBadRequest, "invalid assignment id")
			return
		}

		var imagePath sql.NullString
		row := db.QueryRow(`
      SELECT image_path
      FROM assignments
      WHERE id = ? AND username = ?
      LIMIT 1
    `, id, username)
		if err := row.Scan(&imagePath); err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				writeError(w, http.StatusNotFound, "assignment not found")
				return
			}
			writeError(w, http.StatusInternalServerError, "failed to query image")
			return
		}

		if !imagePath.Valid || strings.TrimSpace(imagePath.String) == "" {
			writeError(w, http.StatusNotFound, "image not found")
			return
		}

		paths := parseImagePaths(imagePath.String)
		if len(paths) == 0 {
			writeError(w, http.StatusNotFound, "image not found")
			return
		}

		index := 0
		if indexStr := r.URL.Query().Get("index"); indexStr != "" {
			if parsed, err := strconv.Atoi(indexStr); err == nil {
				index = parsed
			}
		}
		if index < 0 || index >= len(paths) {
			writeError(w, http.StatusNotFound, "image not found")
			return
		}

		absImage, err := filepath.Abs(paths[index])
		if err != nil {
			writeError(w, http.StatusInternalServerError, "invalid image path")
			return
		}
		absUpload, err := filepath.Abs(uploadDir)
		if err != nil {
			writeError(w, http.StatusInternalServerError, "invalid upload dir")
			return
		}
		if !strings.HasPrefix(absImage, absUpload+string(os.PathSeparator)) && absImage != absUpload {
			writeError(w, http.StatusForbidden, "invalid image path")
			return
		}

		if _, err := os.Stat(absImage); err != nil {
			writeError(w, http.StatusNotFound, "image not found")
			return
		}

		w.Header().Set("Cache-Control", "private, max-age=3600")
		http.ServeFile(w, r, absImage)
	}
}

func handleGetAssignmentCorrectedImage(db *sql.DB, uploadDir string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		username := strings.TrimSpace(r.URL.Query().Get("username"))
		if username == "" {
			writeError(w, http.StatusBadRequest, "username is required")
			return
		}

		id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
		if err != nil {
			writeError(w, http.StatusBadRequest, "invalid assignment id")
			return
		}

		var imagePath sql.NullString
		row := db.QueryRow(`
      SELECT corrected_image_path
      FROM assignments
      WHERE id = ? AND username = ?
      LIMIT 1
    `, id, username)
		if err := row.Scan(&imagePath); err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				writeError(w, http.StatusNotFound, "assignment not found")
				return
			}
			writeError(w, http.StatusInternalServerError, "failed to query image")
			return
		}

		if !imagePath.Valid || strings.TrimSpace(imagePath.String) == "" {
			writeError(w, http.StatusNotFound, "image not found")
			return
		}

		paths := parseImagePaths(imagePath.String)
		if len(paths) == 0 {
			writeError(w, http.StatusNotFound, "image not found")
			return
		}

		index := 0
		if indexStr := r.URL.Query().Get("index"); indexStr != "" {
			if parsed, err := strconv.Atoi(indexStr); err == nil {
				index = parsed
			}
		}
		if index < 0 || index >= len(paths) {
			writeError(w, http.StatusNotFound, "image not found")
			return
		}

		absImage, err := filepath.Abs(paths[index])
		if err != nil {
			writeError(w, http.StatusInternalServerError, "invalid image path")
			return
		}
		absUpload, err := filepath.Abs(uploadDir)
		if err != nil {
			writeError(w, http.StatusInternalServerError, "invalid upload dir")
			return
		}
		if !strings.HasPrefix(absImage, absUpload+string(os.PathSeparator)) && absImage != absUpload {
			writeError(w, http.StatusForbidden, "invalid image path")
			return
		}

		if _, err := os.Stat(absImage); err != nil {
			writeError(w, http.StatusNotFound, "image not found")
			return
		}

		w.Header().Set("Cache-Control", "private, max-age=3600")
		http.ServeFile(w, r, absImage)
	}
}

func handleCreateAssignment(db *sql.DB, llm *LLMClient, uploadDir string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		contentType := r.Header.Get("Content-Type")
		if strings.HasPrefix(contentType, "multipart/form-data") {
			handleCreateAssignmentWithImage(db, llm, uploadDir, w, r)
			return
		}

		var req createAssignmentRequest
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			writeError(w, http.StatusBadRequest, "invalid json")
			return
		}

		req.Username = strings.TrimSpace(req.Username)
		req.Title = strings.TrimSpace(req.Title)
		req.OCRText = strings.TrimSpace(req.OCRText)

		if req.Username == "" {
			writeError(w, http.StatusBadRequest, "username is required")
			return
		}
		if req.OCRText == "" {
			writeError(w, http.StatusBadRequest, "ocrText is required")
			return
		}
		if req.Title == "" {
			req.Title = fmt.Sprintf("%s %s", defaultTitlePrefix, time.Now().Format("2006-01-02 15:04"))
		}

		markdown, feedback, score, err := llm.FormatAndGrade(r.Context(), req.Title, req.OCRText)
		if err != nil {
			writeError(w, http.StatusBadGateway, fmt.Sprintf("llm failed: %v", err))
			return
		}

		result, err := db.Exec(`
      INSERT INTO assignments (username, title, ocr_text, markdown, feedback, score, image_path, corrected_image_path)
      VALUES (?, ?, ?, ?, ?, ?, ?, ?)
    `, req.Username, req.Title, req.OCRText, markdown, feedback, score, "[]", "[]")
		if err != nil {
			writeError(w, http.StatusInternalServerError, "failed to save assignment")
			return
		}

		id, _ := result.LastInsertId()
		assignment := Assignment{
			ID:                  id,
			Username:            req.Username,
			Title:               req.Title,
			OCRText:             req.OCRText,
			Markdown:            markdown,
			Feedback:            feedback,
			Score:               score,
			ImagePath:           "",
			ImagePaths:          []string{},
			CorrectedImagePath:  "",
			CorrectedImagePaths: []string{},
			CreatedAt:           time.Now(),
		}

		writeJSON(w, http.StatusOK, assignment)
	}
}

func handleUpdateAssignment(db *sql.DB, llm *LLMClient, uploadDir string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		contentType := r.Header.Get("Content-Type")
		if !strings.HasPrefix(contentType, "multipart/form-data") {
			writeError(w, http.StatusBadRequest, "multipart form required")
			return
		}

		r.Body = http.MaxBytesReader(w, r.Body, 50<<20)
		if err := r.ParseMultipartForm(60 << 20); err != nil {
			writeError(w, http.StatusBadRequest, "invalid multipart form")
			return
		}

		username := strings.TrimSpace(r.FormValue("username"))
		if username == "" {
			username = strings.TrimSpace(r.URL.Query().Get("username"))
		}
		if username == "" {
			writeError(w, http.StatusBadRequest, "username is required")
			return
		}

		id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
		if err != nil {
			writeError(w, http.StatusBadRequest, "invalid assignment id")
			return
		}

		var existingTitle string
		var createdAt time.Time
		var existingImagePath sql.NullString
		var existingCorrectedPath sql.NullString
		row := db.QueryRow(`
      SELECT title, created_at, image_path, corrected_image_path
      FROM assignments
      WHERE id = ? AND username = ?
      LIMIT 1
    `, id, username)
		if err := row.Scan(&existingTitle, &createdAt, &existingImagePath, &existingCorrectedPath); err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				writeError(w, http.StatusNotFound, "assignment not found")
				return
			}
			writeError(w, http.StatusInternalServerError, "failed to query assignment")
			return
		}

		title := strings.TrimSpace(r.FormValue("title"))
		if title == "" {
			title = existingTitle
		}
		if title == "" {
			title = fmt.Sprintf("%s %s", defaultTitlePrefix, time.Now().Format("2006-01-02 15:04"))
		}

		files := r.MultipartForm.File["images"]
		if len(files) == 0 {
			files = r.MultipartForm.File["image"]
		}
		if len(files) == 0 {
			writeError(w, http.StatusBadRequest, "image is required")
			return
		}

		if err := os.MkdirAll(uploadDir, 0o755); err != nil {
			writeError(w, http.StatusInternalServerError, "failed to prepare upload dir")
			return
		}

		imagePaths := make([]string, 0, len(files))
		images := make([]LLMImage, 0, len(files))
		for _, header := range files {
			file, err := header.Open()
			if err != nil {
				writeError(w, http.StatusBadRequest, "failed to read image")
				return
			}

			data, err := io.ReadAll(file)
			file.Close()
			if err != nil || len(data) == 0 {
				writeError(w, http.StatusBadRequest, "failed to read image")
				return
			}

			mimeType := header.Header.Get("Content-Type")
			if mimeType == "" {
				mimeType = http.DetectContentType(data)
			}
			if !strings.HasPrefix(mimeType, "image/") {
				writeError(w, http.StatusBadRequest, "invalid image type")
				return
			}

			filename := fmt.Sprintf("%s_%s%s", time.Now().Format("20060102_150405"), randomToken(6), extFromMime(mimeType))
			imagePath := filepath.Join(uploadDir, filename)
			if err := os.WriteFile(imagePath, data, 0o644); err != nil {
				writeError(w, http.StatusInternalServerError, "failed to save image")
				return
			}

			imagePaths = append(imagePaths, imagePath)
			images = append(images, LLMImage{
				MimeType: mimeType,
				Base64:   base64.StdEncoding.EncodeToString(data),
			})
		}

		llmResult, err := llm.FormatAndGradeFromImagesResult(r.Context(), title, images)
		if err != nil {
			writeError(w, http.StatusBadGateway, fmt.Sprintf("llm failed: %v", err))
			return
		}

		markdown := llmResult.Markdown
		feedback := llmResult.Feedback
		score := llmResult.Score

		annotations := buildAnnotations(llmResult, markdown, feedback, score)
		correctedPaths, err := generateCorrectedImages(uploadDir, imagePaths, title, annotations)
		if err != nil {
			log.Printf("corrected image failed: %v", err)
			correctedPaths = []string{}
		}

		imagePathsJSON, _ := json.Marshal(imagePaths)
		correctedPathsJSON, _ := json.Marshal(correctedPaths)

		_, err = db.Exec(`
      UPDATE assignments
      SET title = ?, ocr_text = ?, markdown = ?, feedback = ?, score = ?, image_path = ?, corrected_image_path = ?
      WHERE id = ? AND username = ?
    `, title, "", markdown, feedback, score, string(imagePathsJSON), string(correctedPathsJSON), id, username)
		if err != nil {
			writeError(w, http.StatusInternalServerError, "failed to update assignment")
			return
		}

		if existingImagePath.Valid {
			cleanupImages(uploadDir, existingImagePath.String)
		}
		if existingCorrectedPath.Valid {
			cleanupImages(uploadDir, existingCorrectedPath.String)
		}

		assignment := Assignment{
			ID:                  id,
			Username:            username,
			Title:               title,
			OCRText:             "",
			Markdown:            markdown,
			Feedback:            feedback,
			Score:               score,
			ImagePath:           firstPath(imagePaths),
			ImagePaths:          imagePaths,
			CorrectedImagePath:  firstPath(correctedPaths),
			CorrectedImagePaths: correctedPaths,
			CreatedAt:           createdAt,
		}

		writeJSON(w, http.StatusOK, assignment)
	}
}

func handleCreateAssignmentWithImage(db *sql.DB, llm *LLMClient, uploadDir string, w http.ResponseWriter, r *http.Request) {
	r.Body = http.MaxBytesReader(w, r.Body, 50<<20)
	if err := r.ParseMultipartForm(60 << 20); err != nil {
		writeError(w, http.StatusBadRequest, "invalid multipart form")
		return
	}

	username := strings.TrimSpace(r.FormValue("username"))
	title := strings.TrimSpace(r.FormValue("title"))
	if username == "" {
		writeError(w, http.StatusBadRequest, "username is required")
		return
	}
	if title == "" {
		title = fmt.Sprintf("%s %s", defaultTitlePrefix, time.Now().Format("2006-01-02 15:04"))
	}

	files := r.MultipartForm.File["images"]
	if len(files) == 0 {
		files = r.MultipartForm.File["image"]
	}
	if len(files) == 0 {
		writeError(w, http.StatusBadRequest, "image is required")
		return
	}

	if err := os.MkdirAll(uploadDir, 0o755); err != nil {
		writeError(w, http.StatusInternalServerError, "failed to prepare upload dir")
		return
	}

	imagePaths := make([]string, 0, len(files))
	images := make([]LLMImage, 0, len(files))
	for _, header := range files {
		file, err := header.Open()
		if err != nil {
			writeError(w, http.StatusBadRequest, "failed to read image")
			return
		}

		data, err := io.ReadAll(file)
		file.Close()
		if err != nil || len(data) == 0 {
			writeError(w, http.StatusBadRequest, "failed to read image")
			return
		}

		mimeType := header.Header.Get("Content-Type")
		if mimeType == "" {
			mimeType = http.DetectContentType(data)
		}
		if !strings.HasPrefix(mimeType, "image/") {
			writeError(w, http.StatusBadRequest, "invalid image type")
			return
		}

		filename := fmt.Sprintf("%s_%s%s", time.Now().Format("20060102_150405"), randomToken(6), extFromMime(mimeType))
		imagePath := filepath.Join(uploadDir, filename)
		if err := os.WriteFile(imagePath, data, 0o644); err != nil {
			writeError(w, http.StatusInternalServerError, "failed to save image")
			return
		}

		imagePaths = append(imagePaths, imagePath)
		images = append(images, LLMImage{
			MimeType: mimeType,
			Base64:   base64.StdEncoding.EncodeToString(data),
		})
	}

	llmResult, err := llm.FormatAndGradeFromImagesResult(r.Context(), title, images)
	if err != nil {
		writeError(w, http.StatusBadGateway, fmt.Sprintf("llm failed: %v", err))
		return
	}

	markdown := llmResult.Markdown
	feedback := llmResult.Feedback
	score := llmResult.Score

	annotations := buildAnnotations(llmResult, markdown, feedback, score)
	correctedPaths, err := generateCorrectedImages(uploadDir, imagePaths, title, annotations)
	if err != nil {
		log.Printf("corrected image failed: %v", err)
		correctedPaths = []string{}
	}

	imagePathsJSON, _ := json.Marshal(imagePaths)
	correctedPathsJSON, _ := json.Marshal(correctedPaths)

	result, err := db.Exec(`
      INSERT INTO assignments (username, title, ocr_text, markdown, feedback, score, image_path, corrected_image_path)
      VALUES (?, ?, ?, ?, ?, ?, ?, ?)
    `, username, title, "", markdown, feedback, score, string(imagePathsJSON), string(correctedPathsJSON))
	if err != nil {
		writeError(w, http.StatusInternalServerError, "failed to save assignment")
		return
	}

	id, _ := result.LastInsertId()
	assignment := Assignment{
		ID:                  id,
		Username:            username,
		Title:               title,
		OCRText:             "",
		Markdown:            markdown,
		Feedback:            feedback,
		Score:               score,
		ImagePath:           firstPath(imagePaths),
		ImagePaths:          imagePaths,
		CorrectedImagePath:  firstPath(correctedPaths),
		CorrectedImagePaths: correctedPaths,
		CreatedAt:           time.Now(),
	}

	writeJSON(w, http.StatusOK, assignment)
}

func handleDeleteAssignment(db *sql.DB, uploadDir string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		username := strings.TrimSpace(r.URL.Query().Get("username"))
		if username == "" {
			writeError(w, http.StatusBadRequest, "username is required")
			return
		}

		id, err := strconv.ParseInt(chi.URLParam(r, "id"), 10, 64)
		if err != nil {
			writeError(w, http.StatusBadRequest, "invalid assignment id")
			return
		}

		var imagePath sql.NullString
		var correctedPath sql.NullString
		row := db.QueryRow(`
      SELECT image_path, corrected_image_path
      FROM assignments
      WHERE id = ? AND username = ?
      LIMIT 1
    `, id, username)
		if err := row.Scan(&imagePath, &correctedPath); err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				writeError(w, http.StatusNotFound, "assignment not found")
				return
			}
			writeError(w, http.StatusInternalServerError, "failed to load assignment")
			return
		}

		res, err := db.Exec(`DELETE FROM assignments WHERE id = ? AND username = ?`, id, username)
		if err != nil {
			writeError(w, http.StatusInternalServerError, "failed to delete assignment")
			return
		}
		affected, _ := res.RowsAffected()
		if affected == 0 {
			writeError(w, http.StatusNotFound, "assignment not found")
			return
		}

		if imagePath.Valid {
			cleanupImages(uploadDir, imagePath.String)
		}
		if correctedPath.Valid {
			cleanupImages(uploadDir, correctedPath.String)
		}

		writeJSON(w, http.StatusOK, map[string]any{"deleted": id})
	}
}

func handleAdminDashboard(db *sql.DB) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		rows, err := db.Query(`
      SELECT id, username, title, score, created_at
      FROM assignments
      ORDER BY created_at DESC
      LIMIT 200
    `)
		if err != nil {
			http.Error(w, "failed to load assignments", http.StatusInternalServerError)
			return
		}
		defer rows.Close()

		assignments := make([]Assignment, 0)
		for rows.Next() {
			var a Assignment
			if err := rows.Scan(&a.ID, &a.Username, &a.Title, &a.Score, &a.CreatedAt); err != nil {
				http.Error(w, "failed to read assignments", http.StatusInternalServerError)
				return
			}
			assignments = append(assignments, a)
		}

		w.Header().Set("Content-Type", "text/html; charset=utf-8")
		if err := adminTemplate.Execute(w, map[string]any{"Assignments": assignments}); err != nil {
			http.Error(w, "failed to render", http.StatusInternalServerError)
			return
		}
	}
}

func (c *LLMClient) FormatAndGradeResult(ctx context.Context, title, ocrText string) (LLMResult, error) {
	if c.baseURL == "" || c.apiKey == "" {
		return LLMResult{}, errors.New("LLM_BASE_URL or LLM_API_KEY missing")
	}

	systemPrompt := "你是一位严谨的阅卷老师。请把 OCR 文字整理成清晰的 Markdown，保持原有题号、分点、段落与格式。然后进行详细批改：对每一道题给出参考答案、解题过程拆解、学生答案摘录、正确性判定（正确/部分正确/错误）、正确答案对比说明、考察知识点说明。最后给出整体评价与改进建议。请额外输出 annotations 字段，数组，每条为简短批注（<=25字），用于叠加到作业图片。仅输出 JSON，字段为 markdown、feedback、score(0-100整数)、annotations。不要加代码块或多余文字。"
	userPrompt := fmt.Sprintf("标题: %s\n\nOCR文本:\n%s", title, ocrText)

	payload := map[string]any{
		"model": c.model,
		"messages": []map[string]string{
			{"role": "system", "content": systemPrompt},
			{"role": "user", "content": userPrompt},
		},
		"stream": false,
	}

	content, err := c.chatCompletionWithRetry(ctx, payload)
	if err != nil {
		return LLMResult{}, err
	}

	result := LLMResult{}
	if err := decodeLLMResult(content, &result); err != nil {
		markdown, feedback, score := fallbackFromContent(content, ocrText)
		return LLMResult{Markdown: markdown, Feedback: feedback, Score: score}, nil
	}

	result.Markdown = strings.TrimSpace(result.Markdown)
	result.Feedback = strings.TrimSpace(result.Feedback)
	if result.Markdown == "" {
		result.Markdown = ocrText
	}

	return result, nil
}

func (c *LLMClient) FormatAndGrade(ctx context.Context, title, ocrText string) (string, string, int, error) {
	result, err := c.FormatAndGradeResult(ctx, title, ocrText)
	if err != nil {
		return "", "", 0, err
	}
	return result.Markdown, result.Feedback, result.Score, nil
}

func (c *LLMClient) FormatAndGradeFromImage(ctx context.Context, title, mimeType, imageBase64 string) (string, string, int, error) {
	return c.FormatAndGradeFromImages(ctx, title, []LLMImage{{MimeType: mimeType, Base64: imageBase64}})
}

func (c *LLMClient) FormatAndGradeFromImagesResult(ctx context.Context, title string, images []LLMImage) (LLMResult, error) {
	if c.baseURL == "" || c.apiKey == "" {
		return LLMResult{}, errors.New("LLM_BASE_URL or LLM_API_KEY missing")
	}
	if len(images) == 0 {
		return LLMResult{}, errors.New("no images provided")
	}

	systemPrompt := "你是一位严谨的阅卷老师。请根据作业图片识别内容，整理为清晰的 Markdown，保持题号、分点、段落与格式。然后进行详细批改：对每一道题给出参考答案、解题过程拆解、学生答案摘录、正确性判定（正确/部分正确/错误）、正确答案对比说明、考察知识点说明。最后给出整体评价与改进建议。若题目或答案无法识别，请明确标注“无法辨识”。请额外输出 annotations 字段，数组，每条为简短批注（<=25字），用于叠加到作业图片。仅输出 JSON，字段为 markdown、feedback、score(0-100整数)、annotations。不要加代码块或多余文字。"
	userText := fmt.Sprintf("标题: %s\n以下图片按上传顺序依次为第1页、第2页……请综合批改。", title)

	contentParts := make([]map[string]any, 0, len(images)+1)
	contentParts = append(contentParts, map[string]any{"type": "text", "text": userText})
	for _, img := range images {
		contentParts = append(contentParts, map[string]any{
			"type": "image_url",
			"image_url": map[string]string{
				"url": fmt.Sprintf("data:%s;base64,%s", img.MimeType, img.Base64),
			},
		})
	}

	payload := map[string]any{
		"model": c.model,
		"messages": []map[string]any{
			{"role": "system", "content": systemPrompt},
			{"role": "user", "content": contentParts},
		},
		"stream": false,
	}

	content, err := c.chatCompletionWithRetry(ctx, payload)
	if err != nil {
		return LLMResult{}, err
	}

	result := LLMResult{}
	if err := decodeLLMResult(content, &result); err != nil {
		markdown, feedback, score := fallbackFromContent(content, "未识别到有效内容。")
		return LLMResult{Markdown: markdown, Feedback: feedback, Score: score}, nil
	}

	result.Markdown = strings.TrimSpace(result.Markdown)
	result.Feedback = strings.TrimSpace(result.Feedback)
	if result.Markdown == "" {
		result.Markdown = "未识别到有效内容。"
	}

	return result, nil
}

func (c *LLMClient) FormatAndGradeFromImages(ctx context.Context, title string, images []LLMImage) (string, string, int, error) {
	result, err := c.FormatAndGradeFromImagesResult(ctx, title, images)
	if err != nil {
		return "", "", 0, err
	}
	return result.Markdown, result.Feedback, result.Score, nil
}

func (c *LLMClient) chatCompletionWithRetry(ctx context.Context, payload map[string]any) (string, error) {
	var lastErr error

	for attempt := 0; attempt <= c.maxRetries; attempt++ {
		if attempt > 0 {
			sleep := time.Duration(1<<uint(attempt-1)) * time.Second
			time.Sleep(sleep)
		}

		content, status, err := c.chatCompletion(ctx, payload)
		if err == nil {
			if status >= 500 {
				lastErr = fmt.Errorf("llm server error: %d", status)
				continue
			}
			return content, nil
		}

		lastErr = err
		if status >= 500 || status == 0 {
			continue
		}
		break
	}

	if lastErr == nil {
		lastErr = errors.New("llm request failed")
	}
	return "", lastErr
}

func (c *LLMClient) chatCompletion(ctx context.Context, payload map[string]any) (string, int, error) {
	body, err := json.Marshal(payload)
	if err != nil {
		return "", 0, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL, bytes.NewReader(body))
	if err != nil {
		return "", 0, err
	}
	req.Header.Set("Authorization", "Bearer "+c.apiKey)
	req.Header.Set("Content-Type", "application/json")

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return "", 0, err
	}
	defer resp.Body.Close()

	var result struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return "", resp.StatusCode, err
	}
	if len(result.Choices) == 0 {
		return "", resp.StatusCode, errors.New("llm response missing choices")
	}

	return result.Choices[0].Message.Content, resp.StatusCode, nil
}

func decodeLLMResult(content string, out *LLMResult) error {
	content = strings.TrimSpace(content)
	if content == "" {
		return errors.New("empty content")
	}

	if json.Unmarshal([]byte(content), out) == nil {
		return nil
	}

	start := strings.Index(content, "{")
	end := strings.LastIndex(content, "}")
	if start == -1 || end == -1 || end <= start {
		return errors.New("json not found")
	}

	snippet := content[start : end+1]
	if err := json.Unmarshal([]byte(snippet), out); err != nil {
		return err
	}
	return nil
}

var scorePattern = regexp.MustCompile(`(?i)(?:评分|score)[:：]?\s*([0-9]{1,3})`)

func parseScoreFromText(text string) int {
	matches := scorePattern.FindStringSubmatch(text)
	if len(matches) < 2 {
		return 0
	}
	value, err := strconv.Atoi(matches[1])
	if err != nil {
		return 0
	}
	if value < 0 {
		return 0
	}
	if value > 100 {
		return 100
	}
	return value
}

func fallbackFromContent(content, defaultMarkdown string) (string, string, int) {
	markdown := strings.TrimSpace(content)
	if markdown == "" {
		markdown = defaultMarkdown
	}
	score := parseScoreFromText(markdown)
	return markdown, "", score
}

func buildAnnotations(result LLMResult, markdown, feedback string, score int) []string {
	annotations := normalizeAnnotations(result.Annotations)
	if len(annotations) == 0 {
		annotations = deriveAnnotations(feedback)
	}
	if len(annotations) == 0 {
		annotations = deriveAnnotations(markdown)
	}
	if len(annotations) == 0 {
		annotations = []string{"已完成批改，请查看详细结果。"}
	}
	if score > 0 && !hasScoreAnnotation(annotations) {
		annotations = append([]string{fmt.Sprintf("评分：%d", score)}, annotations...)
	}
	if len(annotations) > 8 {
		annotations = append(annotations[:8], "更多细节请查看批改结果")
	}
	return annotations
}

func normalizeAnnotations(raw []string) []string {
	if len(raw) == 0 {
		return nil
	}
	normalized := make([]string, 0, len(raw))
	for _, item := range raw {
		clean := strings.TrimSpace(item)
		if clean == "" {
			continue
		}
		normalized = append(normalized, clean)
	}
	return normalized
}

func hasScoreAnnotation(lines []string) bool {
	for _, line := range lines {
		if strings.Contains(line, "评分") || strings.Contains(strings.ToLower(line), "score") {
			return true
		}
	}
	return false
}

func deriveAnnotations(text string) []string {
	text = strings.TrimSpace(text)
	if text == "" {
		return nil
	}
	parts := splitSentences(text)
	lines := make([]string, 0, len(parts))
	for _, part := range parts {
		clean := strings.TrimSpace(part)
		if clean == "" {
			continue
		}
		lines = append(lines, clean)
	}
	return lines
}

func splitSentences(text string) []string {
	var parts []string
	var buf strings.Builder
	for _, r := range []rune(text) {
		if r == '\r' {
			continue
		}
		buf.WriteRune(r)
		switch r {
		case '\n', '。', '；', ';', '!', '！', '?', '？':
			parts = append(parts, buf.String())
			buf.Reset()
		}
	}
	if buf.Len() > 0 {
		parts = append(parts, buf.String())
	}
	return parts
}

func ensureColumn(db *sql.DB, table, column, columnType string) error {
	rows, err := db.Query(fmt.Sprintf("PRAGMA table_info(%s);", table))
	if err != nil {
		return err
	}
	defer rows.Close()

	for rows.Next() {
		var cid int
		var name string
		var ctype string
		var notnull int
		var dfltValue sql.NullString
		var pk int
		if err := rows.Scan(&cid, &name, &ctype, &notnull, &dfltValue, &pk); err != nil {
			return err
		}
		if name == column {
			return nil
		}
	}

	_, err = db.Exec(fmt.Sprintf("ALTER TABLE %s ADD COLUMN %s %s;", table, column, columnType))
	return err
}

func parseImagePaths(raw string) []string {
	raw = strings.TrimSpace(raw)
	if raw == "" {
		return nil
	}
	if raw == "null" {
		return nil
	}
	if strings.HasPrefix(raw, "[") {
		var paths []string
		if json.Unmarshal([]byte(raw), &paths) == nil {
			return paths
		}
	}
	return []string{raw}
}

func setImageFields(a *Assignment, raw string) {
	paths := parseImagePaths(raw)
	a.ImagePaths = paths
	a.ImagePath = firstPath(paths)
}

func setCorrectedImageFields(a *Assignment, raw string) {
	paths := parseImagePaths(raw)
	a.CorrectedImagePaths = paths
	a.CorrectedImagePath = firstPath(paths)
}

func firstPath(paths []string) string {
	if len(paths) == 0 {
		return ""
	}
	return paths[0]
}

func cleanupImages(uploadDir, raw string) {
	paths := parseImagePaths(raw)
	if len(paths) == 0 {
		return
	}

	absUpload, err := filepath.Abs(uploadDir)
	if err != nil {
		return
	}

	for _, path := range paths {
		absImage, err := filepath.Abs(path)
		if err != nil {
			continue
		}
		if !strings.HasPrefix(absImage, absUpload+string(os.PathSeparator)) && absImage != absUpload {
			continue
		}
		_ = os.Remove(absImage)
	}
}

func generateCorrectedImages(uploadDir string, imagePaths []string, title string, annotations []string) ([]string, error) {
	if len(imagePaths) == 0 {
		return nil, nil
	}
	annotations = normalizeAnnotations(annotations)
	if len(annotations) == 0 {
		annotations = []string{"已完成批改，请查看详细结果。"}
	}

	results := make([]string, 0, len(imagePaths))
	for index, srcPath := range imagePaths {
		filename := fmt.Sprintf("%s_%s_corrected_%d.jpg", time.Now().Format("20060102_150405"), randomToken(6), index+1)
		dstPath := filepath.Join(uploadDir, filename)
		if err := renderCorrectionImage(srcPath, dstPath, title, annotations); err != nil {
			return nil, err
		}
		results = append(results, dstPath)
	}
	return results, nil
}

func renderCorrectionImage(srcPath, dstPath, title string, annotations []string) error {
	file, err := os.Open(srcPath)
	if err != nil {
		return err
	}
	defer file.Close()

	img, _, err := image.Decode(file)
	if err != nil {
		return err
	}

	dc := gg.NewContextForImage(img)
	width := float64(dc.Width())
	height := float64(dc.Height())

	fontSize := clamp(width/36.0, 18, 30)
	if err := loadFontFace(dc, fontSize); err != nil {
		dc.SetFontFace(basicfont.Face7x13)
		fontSize = 14
	}

	dc.SetRGBA(1, 1, 1, 0.05)
	dc.DrawRectangle(0, 0, width, height)
	dc.Fill()

	label := "订正"
	if strings.TrimSpace(title) != "" {
		label = "订正"
	}
	labelX := width - 24
	labelY := 24.0
	drawShadowText(dc, label, labelX, labelY, 1, 0)

	startX := width * 0.06
	startY := height * 0.12
	maxWidth := width * 0.7
	lineHeight := fontSize * 1.45
	maxY := height * 0.9

	y := startY
	for _, note := range annotations {
		lines := wrapText(dc, note, maxWidth)
		for _, line := range lines {
			if y > maxY {
				return saveJPEG(dstPath, dc.Image())
			}
			drawShadowText(dc, line, startX, y, 0, 0)
			y += lineHeight
		}
		y += lineHeight * 0.25
	}

	return saveJPEG(dstPath, dc.Image())
}

func loadFontFace(dc *gg.Context, size float64) error {
	candidates := []string{
		"/usr/share/fonts/noto/NotoSansCJK-Regular.ttc",
		"/usr/share/fonts/noto/NotoSansCJK-Bold.ttc",
		"/usr/share/fonts/noto/NotoSerifCJK-Regular.ttc",
		"/usr/share/fonts/noto/NotoSerifCJK-Bold.ttc",
		"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
	}
	for _, path := range candidates {
		if _, err := os.Stat(path); err != nil {
			continue
		}
		if err := dc.LoadFontFace(path, size); err == nil {
			return nil
		}
	}
	return errors.New("font not found")
}

func wrapText(dc *gg.Context, text string, maxWidth float64) []string {
	text = strings.TrimSpace(text)
	if text == "" {
		return nil
	}

	var lines []string
	var current []rune
	flush := func() {
		if len(current) == 0 {
			return
		}
		lines = append(lines, strings.TrimSpace(string(current)))
		current = nil
	}

	for _, r := range []rune(text) {
		if r == '\n' {
			flush()
			continue
		}
		current = append(current, r)
		if len(current) <= 1 {
			continue
		}
		if width, _ := dc.MeasureString(string(current)); width > maxWidth {
			last := current[len(current)-1]
			current = current[:len(current)-1]
			flush()
			current = []rune{last}
		}
	}
	flush()
	return lines
}

func drawShadowText(dc *gg.Context, text string, x, y, ax, ay float64) {
	dc.SetRGBA(1, 1, 1, 0.85)
	dc.DrawStringAnchored(text, x+1, y+1, ax, ay)
	dc.SetRGB(0.82, 0, 0)
	dc.DrawStringAnchored(text, x, y, ax, ay)
}

func saveJPEG(path string, img image.Image) error {
	tmp, err := os.Create(path)
	if err != nil {
		return err
	}
	defer tmp.Close()
	return jpeg.Encode(tmp, img, &jpeg.Options{Quality: 92})
}

func clamp(value, min, max float64) float64 {
	return math.Max(min, math.Min(value, max))
}

func randomToken(length int) string {
	const charset = "abcdefghijklmnopqrstuvwxyz0123456789"
	buf := make([]byte, length)
	if _, err := rand.Read(buf); err != nil {
		return fmt.Sprintf("%d", time.Now().UnixNano())
	}
	for i, b := range buf {
		buf[i] = charset[int(b)%len(charset)]
	}
	return string(buf)
}

func extFromMime(mimeType string) string {
	switch mimeType {
	case "image/png":
		return ".png"
	case "image/webp":
		return ".webp"
	case "image/gif":
		return ".gif"
	case "image/jpeg":
		return ".jpg"
	default:
		return ".img"
	}
}

func basicAuth(user, pass string) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			username, password, ok := r.BasicAuth()
			if !ok || username != user || password != pass {
				w.Header().Set("WWW-Authenticate", `Basic realm="Homework Admin"`)
				writeError(w, http.StatusUnauthorized, "unauthorized")
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

func corsMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusNoContent)
			return
		}
		next.ServeHTTP(w, r)
	})
}

func writeJSON(w http.ResponseWriter, status int, payload any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	if err := json.NewEncoder(w).Encode(payload); err != nil {
		log.Printf("failed to write json: %v", err)
	}
}

func writeError(w http.ResponseWriter, status int, message string) {
	writeJSON(w, status, map[string]string{"error": message})
}

func getEnv(key, fallback string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return fallback
}

func getEnvAsInt(key string, fallback int) int {
	if value := os.Getenv(key); value != "" {
		if parsed, err := strconv.Atoi(value); err == nil {
			return parsed
		}
	}
	return fallback
}