BoostAI/Backend/internal/aireview/service.go

package aireview

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

type requestStyle string

const (
	requestStyleResponses       requestStyle = "responses"
	requestStyleChatCompletions requestStyle = "chat_completions"
)

type Service struct {
	endpoint string
	apiKey   string
	model    string
	client   *http.Client
}

type AssignmentReviewInput struct {
	AssignmentID   int64                   `json:"assignmentId"`
	StudentID      int64                   `json:"studentId"`
	AssignmentTitle string                 `json:"assignmentTitle"`
	Instructions   string                  `json:"instructions,omitempty"`
	PassThreshold  float64                 `json:"passThreshold"`
	Questions      []AssignmentQuestionInput `json:"questions"`
}

type AssignmentQuestionInput struct {
	QuestionID     int64    `json:"questionId"`
	Position       int32    `json:"position"`
	Title          string   `json:"title"`
	Prompt         string   `json:"prompt"`
	Subject        string   `json:"subject,omitempty"`
	Source         string   `json:"source,omitempty"`
	CorrectAnswer  string   `json:"correctAnswer,omitempty"`
	QuestionTags   []string `json:"questionTags,omitempty"`
	SolveMode      string   `json:"solveMode,omitempty"`
	AnswerText     string   `json:"answerText,omitempty"`
	WorkingSteps   string   `json:"workingSteps,omitempty"`
	IsCorrect      *bool    `json:"isCorrect,omitempty"`
	AnswerStatus   string   `json:"answerStatus,omitempty"`
}

type AssignmentReviewResult struct {
	Questions            []QuestionReviewResult `json:"questions"`
	AssignmentSummary    string                 `json:"assignmentSummary"`
	RecommendedNextStep  string                 `json:"recommendedNextStep"`
}

type RedoPlanInput struct {
	AssignmentID         int64               `json:"assignmentId"`
	StudentID            int64               `json:"studentId"`
	AssignmentTitle      string              `json:"assignmentTitle"`
	Instructions         string              `json:"instructions,omitempty"`
	TeacherFeedback      string              `json:"teacherFeedback,omitempty"`
	PassThreshold        float64             `json:"passThreshold"`
	TopicScores          map[string]float64  `json:"topicScores"`
	WeakTags             []string            `json:"weakTags,omitempty"`
	RecentIssues         []string            `json:"recentIssues,omitempty"`
	AllowedTopics        []string            `json:"allowedTopics"`
	AllowedDifficulties  []string            `json:"allowedDifficulties"`
}

type RedoPlanResult struct {
	Rationale   string             `json:"rationale"`
	QuestionSet []RedoPlanQuestion `json:"questionSet"`
}

type RedoPlanQuestion struct {
	Topic      string   `json:"topic"`
	Difficulty string   `json:"difficulty"`
	Tags       []string `json:"tags,omitempty"`
	Reason     string   `json:"reason"`
}

type QuestionReviewResult struct {
	QuestionID           int64   `json:"questionId"`
	AiFeedback           string  `json:"aiFeedback"`
	UnderstandingScore   float64 `json:"understandingScore"`
	Confidence           float64 `json:"confidence"`
	NeedsAttention       bool    `json:"needsAttention"`
	IssueReason          string  `json:"issueReason"`
}

func NewService(endpoint, apiKey, model string) *Service {
	return &Service{
		endpoint: strings.TrimSpace(endpoint),
		apiKey:   strings.TrimSpace(apiKey),
		model:    strings.TrimSpace(model),
		client: &http.Client{
			Timeout: 45 * time.Second,
		},
	}
}

func (s *Service) Enabled() bool {
	return s != nil && s.endpoint != "" && s.apiKey != "" && s.model != ""
}

func (s *Service) ReviewSubmission(ctx context.Context, input AssignmentReviewInput) (*AssignmentReviewResult, error) {
	if !s.Enabled() {
		return nil, fmt.Errorf("AI review is not configured")
	}

	payloadJSON, err := json.Marshal(input)
	if err != nil {
		return nil, fmt.Errorf("marshal AI review input: %w", err)
	}

	outputText, err := s.runStructuredRequest(
		ctx,
		strings.TrimSpace(`You are reviewing student homework submissions for a teacher workflow.

You must assess the student's understanding by looking at the student's final answer and working against the saved correct answer when one is available. Do not re-grade weighting.

Rules:
- correctness score is fixed at 1.0 externally and must not vary.
- question weighting is fixed at 1.0 externally and must not vary.
- understandingScore is the only variable score and must be between 0.0 and 1.0.
- confidence must be between 0.0 and 1.0.
- every question in the assignment must be included in the output, even if the student left it blank.
- if answerText and workingSteps are both empty, treat the question as unanswered and set understandingScore to 0.0.
- unanswered questions should normally set needsAttention to true and explain that no answer was submitted.
- when correctAnswer is present, explicitly compare the student's answerText and workingSteps against that correctAnswer before judging understanding.
- when the student's answer is wrong, issueReason should say what is mismatched, missing, or misunderstood relative to the correctAnswer, not just give a generic comment.
- when the student's answer is correct but the explanation is weak, issueReason should make clear that correctness was reached but understanding evidence is still limited.
- when correctAnswer is missing, fall back to judging from the student's explanation, steps, and internal consistency only.
- needsAttention should be true when the student likely needs follow-up help based on their understanding, explanation quality, or uncertainty.
- issueReason should be concise and directly tied to understanding gaps, explicitly grounded in the comparison to the correct answer whenever available.
- aiFeedback should be concise, teacher-facing, and about the student's understanding for that question, referencing the answer-vs-correct-answer comparison when it materially explains the judgment.
- recommendedNextStep must be exactly one of: redo, accept, support.

Interpretation guidance:
- accept = understanding is generally sufficient for the assignment so the student can continue.
- support = the student shows meaningful gaps and likely needs targeted help.
- redo = the student should redo the assignment because understanding is broadly too weak or incomplete.

		Review the full assignment in one pass and produce a short assignment-level summary.`),
		string(payloadJSON),
		"assignment_review",
		reviewSchema(),
	)
	if err != nil {
		return nil, err
	}

	var result AssignmentReviewResult
	if err := json.Unmarshal([]byte(outputText), &result); err != nil {
		return nil, fmt.Errorf("decode AI review structured output: %w", err)
	}

	sanitizeResult(&result)
	return &result, nil
}

func (s *Service) PlanRedoAssignment(ctx context.Context, input RedoPlanInput) (*RedoPlanResult, error) {
	if !s.Enabled() {
		return nil, fmt.Errorf("AI review is not configured")
	}

	payloadJSON, err := json.Marshal(input)
	if err != nil {
		return nil, fmt.Errorf("marshal redo plan input: %w", err)
	}

	outputText, err := s.runStructuredRequest(
		ctx,
		strings.TrimSpace(`You are planning the next redo assignment for a student.

You are NOT writing final math questions. You are only producing a structured topic+difficulty blueprint for a later generator layer.

Rules:
- Use only the allowedTopics values exactly as provided.
- Use only the allowedDifficulties values exactly as provided.
- Return between 5 and 10 planned items.
- Focus most heavily on the weakest topics and weak tags, while still keeping the redo assignment coherent with the current assignment title/instructions and any teacher feedback.
- Prefer a sensible progression of difficulty rather than making everything hard.
- If teacherFeedback contains a specific reteach direction, incorporate it.
- Tags should be short machine-friendly labels and may be empty.
- rationale should briefly explain why these topics/difficulties were chosen from the weakness summary.
- reason on each item should briefly explain why that topic/difficulty belongs in the redo set.
- Do not invent topics outside the allowed topic vocabulary.
- Do not output prose outside the JSON schema.`),
		string(payloadJSON),
		"redo_assignment_plan",
		redoPlanSchema(),
	)
	if err != nil {
		return nil, err
	}

	var result RedoPlanResult
	if err := json.Unmarshal([]byte(outputText), &result); err != nil {
		return nil, fmt.Errorf("decode redo plan structured output: %w", err)
	}

	sanitizeRedoPlan(&result, input.AllowedTopics, input.AllowedDifficulties)
	return &result, nil
}

func (s *Service) runStructuredRequest(ctx context.Context, systemPrompt, userPrompt, schemaName string, schema map[string]any) (string, error) {
	respBytes, err := s.sendStructuredRequest(ctx, systemPrompt, userPrompt, schemaName, schema)
	if err != nil {
		return "", err
	}

	switch s.requestStyle() {
	case requestStyleChatCompletions:
		return extractChatCompletionText(respBytes)
	default:
		return extractResponsesOutputText(respBytes)
	}
}

func (s *Service) sendStructuredRequest(ctx context.Context, systemPrompt, userPrompt, schemaName string, schema map[string]any) ([]byte, error) {
	body, err := s.buildStructuredRequestBody(systemPrompt, userPrompt, schemaName, schema)
	if err != nil {
		return nil, err
	}

	bodyBytes, err := json.Marshal(body)
	if err != nil {
		return nil, fmt.Errorf("marshal AI review request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, s.endpoint, bytes.NewReader(bodyBytes))
	if err != nil {
		return nil, fmt.Errorf("build AI review request: %w", err)
	}

	req.Header.Set("Content-Type", "application/json")
	s.applyAuthHeader(req)

	resp, err := s.client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("send AI review request: %w", err)
	}
	defer resp.Body.Close()

	respBytes, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("read AI review response: %w", err)
	}

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("AI review request failed: status %d body %s", resp.StatusCode, strings.TrimSpace(string(respBytes)))
	}

	return respBytes, nil
}

func (s *Service) buildStructuredRequestBody(systemPrompt, userPrompt, schemaName string, schema map[string]any) (map[string]any, error) {
	switch s.requestStyle() {
	case requestStyleChatCompletions:
		body := map[string]any{
			"model": s.model,
			"messages": []map[string]any{
				{
					"role":    "system",
					"content": systemPrompt,
				},
				{
					"role":    "user",
					"content": userPrompt,
				},
			},
			"temperature": 0,
			"response_format": map[string]any{
				"type": "json_schema",
				"json_schema": map[string]any{
					"name":   schemaName,
					"strict": true,
					"schema": schema,
				},
			},
		}

		if s.shouldDisableThinking() {
			body["chat_template_kwargs"] = map[string]any{
				"enable_thinking": false,
			}
		}

		return body, nil
	default:
		return map[string]any{
			"model": s.model,
			"input": []map[string]any{
				{
					"role": "system",
					"content": []map[string]any{{
						"type": "input_text",
						"text": systemPrompt,
					}},
				},
				{
					"role": "user",
					"content": []map[string]any{{
						"type": "input_text",
						"text": userPrompt,
					}},
				},
			},
			"text": map[string]any{
				"format": map[string]any{
					"type":   "json_schema",
					"name":   schemaName,
					"strict": true,
					"schema": schema,
				},
			},
		}, nil
	}
}

func (s *Service) requestStyle() requestStyle {
	endpoint := strings.ToLower(strings.TrimSpace(s.endpoint))
	if strings.Contains(endpoint, "/chat/completions") {
		return requestStyleChatCompletions
	}
	return requestStyleResponses
}

func (s *Service) applyAuthHeader(req *http.Request) {
	if s.isAzureEndpoint() {
		req.Header.Set("api-key", s.apiKey)
		return
	}
	req.Header.Set("Authorization", "Bearer "+s.apiKey)
}

func (s *Service) shouldDisableThinking() bool {
	return s.requestStyle() == requestStyleChatCompletions && !s.isAzureEndpoint()
}

func (s *Service) isAzureEndpoint() bool {
	endpoint := strings.ToLower(strings.TrimSpace(s.endpoint))
	return strings.Contains(endpoint, "cognitiveservices.azure.com") || strings.Contains(endpoint, ".openai.azure.com")
}

func reviewSchema() map[string]any {
	return map[string]any{
		"type": "object",
		"additionalProperties": false,
		"properties": map[string]any{
			"questions": map[string]any{
				"type": "array",
				"items": map[string]any{
					"type":                 "object",
					"additionalProperties": false,
					"properties": map[string]any{
						"questionId":         map[string]any{"type": "integer"},
						"aiFeedback":         map[string]any{"type": "string"},
						"understandingScore": map[string]any{"type": "number", "minimum": 0, "maximum": 1},
						"confidence":         map[string]any{"type": "number", "minimum": 0, "maximum": 1},
						"needsAttention":     map[string]any{"type": "boolean"},
						"issueReason":        map[string]any{"type": "string"},
					},
					"required": []string{"questionId", "aiFeedback", "understandingScore", "confidence", "needsAttention", "issueReason"},
				},
			},
			"assignmentSummary":   map[string]any{"type": "string"},
			"recommendedNextStep": map[string]any{"type": "string", "enum": []string{"redo", "accept", "support"}},
		},
		"required": []string{"questions", "assignmentSummary", "recommendedNextStep"},
	}
}

func redoPlanSchema() map[string]any {
	return map[string]any{
		"type": "object",
		"additionalProperties": false,
		"properties": map[string]any{
			"rationale": map[string]any{"type": "string"},
			"questionSet": map[string]any{
				"type": "array",
				"minItems": 5,
				"maxItems": 10,
				"items": map[string]any{
					"type":                 "object",
					"additionalProperties": false,
					"properties": map[string]any{
						"topic":      map[string]any{"type": "string"},
						"difficulty": map[string]any{"type": "string"},
						"tags": map[string]any{
							"type":  "array",
							"items": map[string]any{"type": "string"},
						},
						"reason": map[string]any{"type": "string"},
					},
					"required": []string{"topic", "difficulty", "tags", "reason"},
				},
			},
		},
		"required": []string{"rationale", "questionSet"},
	}
}

func extractResponsesOutputText(respBytes []byte) (string, error) {
	var direct struct {
		OutputText string `json:"output_text"`
	}
	if err := json.Unmarshal(respBytes, &direct); err == nil {
		if text := strings.TrimSpace(direct.OutputText); text != "" {
			return text, nil
		}
	}

	var raw map[string]any
	if err := json.Unmarshal(respBytes, &raw); err != nil {
		return "", fmt.Errorf("decode AI review raw response: %w", err)
	}

	if text := strings.TrimSpace(findOutputText(raw)); text != "" {
		return text, nil
	}

	return "", fmt.Errorf("AI review response did not contain structured output text")
}

func extractChatCompletionText(respBytes []byte) (string, error) {
	var payload struct {
		Choices []struct {
			Message struct {
				Content any `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}

	if err := json.Unmarshal(respBytes, &payload); err != nil {
		return "", fmt.Errorf("decode AI review chat completion response: %w", err)
	}

	for _, choice := range payload.Choices {
		if text := strings.TrimSpace(extractMessageContent(choice.Message.Content)); text != "" {
			return text, nil
		}
	}

	return "", fmt.Errorf("AI review chat completion response did not contain message content")
}

func extractMessageContent(content any) string {
	switch typed := content.(type) {
	case string:
		return typed
	case []any:
		parts := make([]string, 0, len(typed))
		for _, item := range typed {
			text := strings.TrimSpace(findOutputText(item))
			if text == "" {
				continue
			}
			parts = append(parts, text)
		}
		return strings.Join(parts, "\n")
	default:
		return ""
	}
}

func findOutputText(value any) string {
	switch typed := value.(type) {
	case map[string]any:
		if text, ok := typed["output_text"].(string); ok && strings.TrimSpace(text) != "" {
			return text
		}
		if text, ok := typed["text"].(string); ok && strings.TrimSpace(text) != "" {
			return text
		}
		for _, nested := range typed {
			if result := findOutputText(nested); result != "" {
				return result
			}
		}
	case []any:
		for _, nested := range typed {
			if result := findOutputText(nested); result != "" {
				return result
			}
		}
	}

	return ""
}

func sanitizeResult(result *AssignmentReviewResult) {
	result.AssignmentSummary = strings.TrimSpace(result.AssignmentSummary)
	switch result.RecommendedNextStep {
	case "redo", "accept", "support":
	default:
		result.RecommendedNextStep = "support"
	}

	for index := range result.Questions {
		question := &result.Questions[index]
		question.AiFeedback = strings.TrimSpace(question.AiFeedback)
		question.IssueReason = strings.TrimSpace(question.IssueReason)
		question.UnderstandingScore = clamp01(question.UnderstandingScore)
		question.Confidence = clamp01(question.Confidence)
	}
}

func sanitizeRedoPlan(result *RedoPlanResult, allowedTopics []string, allowedDifficulties []string) {
	allowedTopicSet := make(map[string]struct{}, len(allowedTopics))
	for _, topic := range allowedTopics {
		allowedTopicSet[strings.TrimSpace(topic)] = struct{}{}
	}
	allowedDifficultySet := make(map[string]struct{}, len(allowedDifficulties))
	for _, difficulty := range allowedDifficulties {
		allowedDifficultySet[strings.TrimSpace(difficulty)] = struct{}{}
	}

	result.Rationale = strings.TrimSpace(result.Rationale)
	filtered := make([]RedoPlanQuestion, 0, len(result.QuestionSet))
	for _, item := range result.QuestionSet {
		item.Topic = strings.TrimSpace(item.Topic)
		item.Difficulty = strings.ToLower(strings.TrimSpace(item.Difficulty))
		if _, ok := allowedTopicSet[item.Topic]; !ok {
			continue
		}
		if _, ok := allowedDifficultySet[item.Difficulty]; !ok {
			continue
		}
		item.Reason = strings.TrimSpace(item.Reason)
		cleanTags := make([]string, 0, len(item.Tags))
		for _, tag := range item.Tags {
			tag = strings.TrimSpace(tag)
			if tag == "" {
				continue
			}
			cleanTags = append(cleanTags, tag)
		}
		item.Tags = cleanTags
		filtered = append(filtered, item)
	}
	result.QuestionSet = filtered
}

func clamp01(value float64) float64 {
	if value < 0 {
		return 0
	}
	if value > 1 {
		return 1
	}
	return value
}