package assignmentgen

import (
	"context"

	"boostai-backend/internal/database"
	"boostai-backend/internal/questiongen"
	"boostai-backend/internal/sqlc"
)

type SourceBucket string

const (
	SourceBucketCoreTopic    SourceBucket = "core_topic"
	SourceBucketPersonalized SourceBucket = "personalized"
	defaultQuestionSource                 = "assignment_student_generated"
)

type PlanItem struct {
	Topic        sqlc.QuestionTopic
	Difficulty   sqlc.QuestionDifficulty
	Count        int
	SourceBucket SourceBucket
	Seed         int64
}

type GenerateStudentQuestionSetParams struct {
	AssignmentID   int64
	StudentID      int64
	TeacherID      int64
	Subject        string
	QuestionStatus sqlc.QuestionStatus
	QuestionSource string
	Plan           []PlanItem
}

type StoredStudentQuestion struct {
	Mapping      sqlc.AssignmentStudentQuestion
	Question     sqlc.Question
	Tags         []string
	UsedSeed     int64
	SourceBucket string
}

type Service struct {
	db        *database.DB
	generator *questiongen.Service
}

func NewService(db *database.DB, generator *questiongen.Service) *Service {
	return &Service{db: db, generator: generator}
}

func (s *Service) GenerateAndStoreStudentQuestions(ctx context.Context, params GenerateStudentQuestionSetParams) ([]StoredStudentQuestion, error) {
	if err := s.validateGenerateRequest(params); err != nil {
		return nil, err
	}

	tx, err := s.db.Pool.Begin(ctx)
	if err != nil {
		return nil, err
	}
	defer func() {
		_ = tx.Rollback(ctx)
	}()

	queries := sqlc.New(tx)

	if err := validateAssignmentOwnership(ctx, queries, params.AssignmentID, params.TeacherID); err != nil {
		return nil, err
	}
	if err := validateStudentAssignment(ctx, queries, params.AssignmentID, params.StudentID); err != nil {
		return nil, err
	}
	if err := clearStudentQuestionMappings(ctx, queries, params.AssignmentID, params.StudentID); err != nil {
		return nil, err
	}

	questionStatus, questionSource := normalizeQuestionDefaults(params)

	stored, err := s.generateAndStorePlan(ctx, queries, params, questionStatus, questionSource)
	if err != nil {
		return nil, err
	}

	if err := tx.Commit(ctx); err != nil {
		return nil, err
	}

	return stored, nil
}