import Foundation
import FoundationModels
import SharedModels

/// Generates a set of example sentences for a single verb, one per core tense
/// (Issue #27). Mirrors the StoryGenerator pattern: @Generable response types,
/// a static availability flag, and a single generate(...) entry point.
///
/// Issue #33: the model used to drift onto other verbs partway through the
/// 6-example batch (a "tener" set would contain sentences built on estar / ir
/// / deber). Two defenses now apply:
///   1. The prompt embeds the verb's *exact* conjugated forms per tense, so
///      the model echoes a real form instead of recalling one.
///   2. Every generated sentence is validated against those forms; failures
///      are regenerated once, and anything still wrong is dropped rather than
///      shown.
@MainActor
struct VerbExampleGenerator {

    // MARK: - Generable Types

    @Generable
    struct GeneratedExampleSet {
        @Guide(
            description: "Six example sentences, one per tense in the exact order requested. Each sentence must actually use the target verb conjugated in that tense.",
            .count(6)
        )
        var examples: [GeneratedExample]
    }

    @Generable
    struct GeneratedExample {
        @Guide(description: "The tense id this sentence demonstrates. Must match one of the ids provided in the prompt exactly (e.g. ind_presente).")
        var tenseId: String

        @Guide(description: "A natural Spanish sentence, 6-14 words, that uses the target verb in the specified tense. For imperative tenses use tú or nosotros — never yo.")
        var spanish: String

        @Guide(description: "An accurate, idiomatic English translation of the Spanish sentence.")
        var english: String
    }

    // MARK: - Generation

    /// Generate one validated example per tense in `tenseIds`.
    ///
    /// - Parameter formsByTense: the verb's conjugated forms keyed by tenseId
    ///   (from `ReferenceStore.conjugatedForms`). Used to ground the prompt
    ///   and validate the output. A tense with no forms here is accepted
    ///   without validation.
    static func generate(
        verbInfinitive: String,
        verbEnglish: String,
        tenseIds: [String],
        formsByTense: [String: [String]]
    ) async throws -> [VerbExample] {
        // A thrown error here (model busy, context overflow) shouldn't abort the
        // whole generation — treat it as a fully-failed batch so the retry below
        // still runs, same as the retry path's own `try?`.
        let firstPass = (try? await generateBatch(
            verbInfinitive: verbInfinitive,
            verbEnglish: verbEnglish,
            tenseIds: tenseIds,
            formsByTense: formsByTense
        )) ?? [:]

        var valid: [String: VerbExample] = [:]
        var failedTenses: [String] = []
        for id in tenseIds {
            if let ex = firstPass[id], exampleUsesVerb(ex.spanish, forms: formsByTense[id] ?? []) {
                valid[id] = ex
            } else {
                failedTenses.append(id)
            }
        }

        // One focused retry — regenerate the whole batch, but only adopt the
        // results for tenses that failed the first pass.
        if !failedTenses.isEmpty {
            let retry = try? await generateBatch(
                verbInfinitive: verbInfinitive,
                verbEnglish: verbEnglish,
                tenseIds: tenseIds,
                formsByTense: formsByTense
            )
            if let retry {
                for id in failedTenses {
                    if let ex = retry[id], exampleUsesVerb(ex.spanish, forms: formsByTense[id] ?? []) {
                        valid[id] = ex
                    }
                }
            }
        }

        // Requested order, dropping any tense that never produced a valid
        // sentence — better to show fewer examples than wrong ones.
        return tenseIds.compactMap { valid[$0] }
    }

    // MARK: - Single batch call

    private static func generateBatch(
        verbInfinitive: String,
        verbEnglish: String,
        tenseIds: [String],
        formsByTense: [String: [String]]
    ) async throws -> [String: VerbExample] {
        let tenseBlock = tenseIds.compactMap { id -> String? in
            guard let info = TenseInfo.find(id) else { return nil }
            let forms = formsByTense[id] ?? []
            if forms.isEmpty {
                return "- \(id) (\(info.english))"
            }
            return "- \(id) (\(info.english)): use one of these exact conjugated forms — \(forms.joined(separator: ", "))"
        }.joined(separator: "\n")

        let session = LanguageModelSession(instructions: """
            You are a Spanish language teacher writing short example sentences for a learner.
            The learner is studying the verb "\(verbInfinitive)" (to \(verbEnglish)).
            EVERY sentence must use "\(verbInfinitive)" as its main verb, conjugated in the
            requested tense — never substitute a different verb (no estar, ir, deber, etc.
            unless the target verb itself is that verb). Each sentence must:
            - Contain one of the exact conjugated forms listed for its tense.
            - Be 6-14 words, natural and everyday.
            - Use vocabulary appropriate for intermediate learners.
            - Vary subjects and contexts across the set; do not reuse the same subject twice.
            For imperative tenses, address "tú" or "nosotros" — never "yo".
            """)

        let prompt = """
            Write one example sentence for "\(verbInfinitive)" per tense below, in this order:
            \(tenseBlock)

            Return one GeneratedExample per tense with the matching tenseId, spanish, and english.
            The Spanish sentence MUST contain one of the conjugated forms shown for that tense.
            """

        let response = try await session.respond(to: prompt, generating: GeneratedExampleSet.self)

        // `uniquingKeysWith` defensively — the schema forces 6 examples even
        // when fewer tenses are requested, so the model may repeat a tenseId.
        return Dictionary(
            response.content.examples.map {
                ($0.tenseId, VerbExample(tenseId: $0.tenseId, spanish: $0.spanish, english: $0.english))
            },
            uniquingKeysWith: { first, _ in first }
        )
    }

    // MARK: - Validation

    /// True when `sentence` contains at least one of `forms` as a whole word
    /// (accent- and case-insensitive). Empty `forms` → accept (can't validate).
    static func exampleUsesVerb(_ sentence: String, forms: [String]) -> Bool {
        guard !forms.isEmpty else { return true }
        let sentenceWords = foldedWords(sentence)
        let formWords = Set(forms.flatMap { foldedWords($0) })
        return !sentenceWords.isDisjoint(with: formWords)
    }

    private static func foldedWords(_ text: String) -> Set<String> {
        let folded = text.folding(
            options: [.diacriticInsensitive, .caseInsensitive],
            locale: nil
        )
        return Set(folded.split { !$0.isLetter }.map(String.init))
    }

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