Spanish/Conjuga/Conjuga/Services/LexemeSessionQueue.swift

import Foundation
import SharedModels
import SwiftData

/// Which pool a `LexemeSessionQueue` draws from. Mirrors `VocabSessionKind`.
enum LexemeSessionKind {
    /// Due-first + new lexemes from enabled CEFR levels, capped — the
    /// standard SRS session. Ratings update the long-term schedule.
    case standard
    /// Lexemes already studied at least once, most-recent first, uncapped
    /// and unfiltered — a consolidation cram. Ratings drive the in-session
    /// queue only and do NOT reschedule (long-term SM-2 due dates left
    /// untouched, parallel to `VocabSessionKind.reviewLearned`).
    case reviewLearned
}

/// In-session learning-step queue for `Lexeme`-based vocab practice — the
/// non-verb analog of `VocabSessionQueue`. Same Anki-style position-based
/// requeue: Again/Hard requeue close, Good advances state then graduates on
/// the second pass, Easy graduates immediately. `answer` returns a
/// `ReviewQuality` only when the card graduates — that's the rating fed to
/// the cross-session `LexemeReviewStore`.
struct LexemeSessionQueue {

    enum CardState: String {
        case new
        case learning
        case review
    }

    enum Rating {
        case again, hard, good, easy
    }

    struct Entry: Identifiable {
        let id = UUID()
        let lexeme: Lexeme
        var state: CardState
    }

    let drillMode: String
    private(set) var queue: [Entry]
    private(set) var learnedCount: Int = 0
    private let originalLexemes: [Lexeme]

    init(lexemes: [Lexeme], drillMode: String) {
        self.drillMode = drillMode
        self.originalLexemes = lexemes
        self.queue = lexemes.map { Entry(lexeme: $0, state: .new) }
    }

    init(entries: [(lexeme: Lexeme, state: CardState)], drillMode: String, learnedCount: Int) {
        self.drillMode = drillMode
        self.originalLexemes = entries.map(\.lexeme)
        self.queue = entries.map { Entry(lexeme: $0.lexeme, state: $0.state) }
        self.learnedCount = learnedCount
    }

    func snapshot() -> [(lexemeId: String, state: CardState)] {
        queue.map { ($0.lexeme.id, $0.state) }
    }

    var current: Entry? { queue.first }
    var isComplete: Bool { queue.isEmpty }
    var remainingCount: Int { queue.count }

    var progress: Double {
        let total = learnedCount + queue.count
        return total == 0 ? 1 : Double(learnedCount) / Double(total)
    }

    @discardableResult
    mutating func answer(_ rating: Rating) -> ReviewQuality? {
        guard !queue.isEmpty else { return nil }
        var entry = queue.removeFirst()

        switch rating {
        case .again:
            entry.state = .learning
            insert(entry, offset: Int.random(in: 5...8))
            return nil
        case .hard:
            entry.state = .learning
            insert(entry, offset: Int.random(in: 7...10))
            return nil
        case .good:
            if entry.state == .review {
                learnedCount += 1
                return .good
            }
            entry.state = .review
            insert(entry, offset: Int.random(in: 16...24))
            return nil
        case .easy:
            learnedCount += 1
            return .easy
        }
    }

    mutating func restart() {
        queue = originalLexemes.shuffled().map { Entry(lexeme: $0, state: .new) }
        learnedCount = 0
    }

    private mutating func insert(_ entry: Entry, offset: Int) {
        let idx = min(queue.count, offset)
        queue.insert(entry, at: idx)
    }
}

// MARK: - Session lexeme pool

/// Builds a session for a given POS + drill mode: due-first per
/// `LexemeReviewCard`, then fresh (never-studied) lexemes, capped.
enum LexemePool {

    /// Per-session cap. 0/unset → 20. Mirrors `VocabVerbPool.sessionCardLimit`.
    static var sessionCardLimit: Int {
        let stored = UserDefaults.standard.integer(forKey: "lexemeSessionCardLimit")
        return stored == 0 ? 20 : stored
    }

    static func sessionLexemes(
        partOfSpeech: String,
        drillMode: String,
        enabledLevels: Set<LexemeLevel>,
        localContext: ModelContext,
        cloudContext: ModelContext
    ) -> [Lexeme] {
        let pool = fetchStudyable(partOfSpeech: partOfSpeech, context: localContext)

        let descriptor = FetchDescriptor<LexemeReviewCard>(
            predicate: #Predicate<LexemeReviewCard> {
                $0.partOfSpeech == partOfSpeech && $0.drillMode == drillMode
            }
        )
        let reviewCards = (try? cloudContext.fetch(descriptor)) ?? []
        let cardById = Dictionary(
            reviewCards.map { ($0.lexemeId, $0) },
            uniquingKeysWith: { existing, _ in existing }
        )

        let now = Date()
        var due: [(lexeme: Lexeme, dueDate: Date)] = []
        var fresh: [Lexeme] = []
        for lexeme in pool {
            if let card = cardById[lexeme.id] {
                if card.dueDate <= now {
                    // Due cards surface regardless of current level toggles —
                    // SRS isn't level-gated. Already-studied cards keep
                    // coming back on their schedule.
                    due.append((lexeme, card.dueDate))
                }
            } else if enabledLevels.contains(LexemeLevel.level(forRank: lexeme.frequencyRank)) {
                // Fresh (never-studied) cards only enter the pool from
                // levels the user has on. Disabling a level is the lever
                // for "don't introduce me to harder/easier words yet."
                fresh.append(lexeme)
            }
        }

        due.sort { $0.dueDate < $1.dueDate }
        // Fresh cards surface in frequency order — most-useful words first.
        // Lexemes without a rank (frequencyRank == 0) sort last.
        fresh.sort { lhs, rhs in
            let l = lhs.frequencyRank == 0 ? Int.max : lhs.frequencyRank
            let r = rhs.frequencyRank == 0 ? Int.max : rhs.frequencyRank
            if l != r { return l < r }
            return lhs.baseForm < rhs.baseForm
        }

        let ordered = due.map(\.lexeme) + fresh
        return Array(ordered.prefix(sessionCardLimit))
    }

    /// Lexemes the user has already studied at least once for `(POS, drill)`,
    /// most-recently-studied first. Mirrors `VocabVerbPool.reviewLearnedVerbs`.
    static func reviewLearnedLexemes(
        partOfSpeech: String,
        drillMode: String,
        localContext: ModelContext,
        cloudContext: ModelContext
    ) -> [Lexeme] {
        let descriptor = FetchDescriptor<LexemeReviewCard>(
            predicate: #Predicate<LexemeReviewCard> {
                $0.partOfSpeech == partOfSpeech && $0.drillMode == drillMode
            }
        )
        let reviewCards = (try? cloudContext.fetch(descriptor)) ?? []
        let sorted = reviewCards.sorted {
            ($0.lastReviewDate ?? .distantPast) > ($1.lastReviewDate ?? .distantPast)
        }

        let pool = fetchStudyable(partOfSpeech: partOfSpeech, context: localContext)
        let byId = Dictionary(pool.map { ($0.id, $0) }, uniquingKeysWith: { existing, _ in existing })
        return sorted.compactMap { byId[$0.lexemeId] }
    }

    /// Lexemes for a POS. The catalog (`vocab_lexemes.json`) only emits
    /// nouns that have a known gender, so no extra filter is needed here.
    private static func fetchStudyable(partOfSpeech: String, context: ModelContext) -> [Lexeme] {
        let descriptor = FetchDescriptor<Lexeme>(
            predicate: #Predicate<Lexeme> { $0.partOfSpeech == partOfSpeech }
        )
        return (try? context.fetch(descriptor)) ?? []
    }
}