Spanish/Conjuga/Conjuga/Services/StoryGenerator.swift

import Foundation
import FoundationModels
import SharedModels
import SwiftData

@MainActor
struct StoryGenerator {

    // MARK: - Generable Types

    @Generable
    struct GeneratedStory {
        @Guide(description: "A short creative title for the story in Spanish, 3-6 words")
        var title: String

        @Guide(description: "A one-paragraph story in Spanish, 5-8 sentences long, using vocabulary and grammar appropriate for the student level")
        var bodyES: String

        @Guide(description: "An accurate English translation of bodyES")
        var bodyEN: String

        @Guide(description: "Every word from the story annotated with its base form, English meaning, and part of speech. Include articles, prepositions, and all other words.")
        var words: [GeneratedAnnotation]

        @Guide(description: "3 reading comprehension questions about the story, each with 4 answer options in Spanish", .count(3))
        var questions: [GeneratedQuestion]
    }

    @Generable
    struct GeneratedAnnotation {
        @Guide(description: "The exact word as it appears in the story")
        var word: String
        @Guide(description: "The dictionary base form (infinitive for verbs, singular for nouns)")
        var baseForm: String
        @Guide(description: "English translation of the word")
        var english: String
        @Guide(description: "Part of speech: verb, noun, adjective, adverb, preposition, or other")
        var partOfSpeech: String
    }

    @Generable
    struct GeneratedQuestion {
        @Guide(description: "A comprehension question about the story in Spanish")
        var question: String
        @Guide(description: "4 answer options in Spanish", .count(4))
        var options: [String]
        @Guide(description: "Index of the correct answer (0-3)", .range(0...3))
        var correctIndex: Int
    }

    // MARK: - Generation

    static func generate(level: String, tenses: [String]) async throws -> Story {
        let tenseNames = tenses.isEmpty
            ? "present, preterite, imperfect, and future"
            : tenses.joined(separator: ", ")

        let session = LanguageModelSession(instructions: """
            You are a Spanish language teacher creating a short reading exercise.
            The student's level is: \(level).
            Focus on these verb tenses: \(tenseNames).
            Write naturally but keep vocabulary appropriate for the level.
            Use common, everyday scenarios (shopping, travel, family, school, work, food).
            The story should be exactly one paragraph of 5-8 sentences.
            """)

        let response = try await session.respond(
            to: "Create a short Spanish story for reading practice.",
            generating: GeneratedStory.self
        )

        let story = response.content

        let annotations = story.words.map {
            WordAnnotation(word: $0.word, baseForm: $0.baseForm, english: $0.english, partOfSpeech: $0.partOfSpeech)
        }
        let questions = story.questions.map {
            QuizQuestion(question: $0.question, options: $0.options, correctIndex: $0.correctIndex)
        }

        let annotationsJSON = (try? String(data: JSONEncoder().encode(annotations), encoding: .utf8)) ?? "[]"
        let questionsJSON = (try? String(data: JSONEncoder().encode(questions), encoding: .utf8)) ?? "[]"

        return Story(
            title: story.title,
            bodyES: story.bodyES,
            bodyEN: story.bodyEN,
            level: level,
            wordAnnotations: annotationsJSON,
            quizQuestions: questionsJSON
        )
    }

    static var isAvailable: Bool {
        SystemLanguageModel.default.availability == .available
    }
}