import SwiftUI
import SharedModels
import FoundationModels

struct StoryReaderView: View {
    let story: Story

    @Environment(DictionaryService.self) private var dictionary
    @State private var selectedWord: WordAnnotation?
    @State private var showTranslation = false
    @State private var lookupCache: [String: WordAnnotation] = [:]

    private var annotations: [WordAnnotation] { story.decodedAnnotations }
    private var annotationMap: [String: WordAnnotation] {
        Dictionary(annotations.map { (cleanWord($0.word), $0) }, uniquingKeysWith: { first, _ in first })
    }

    var body: some View {
        ScrollView {
            VStack(alignment: .leading, spacing: 20) {
                // Title
                Text(story.title)
                    .font(.title2.bold())

                // Level badge
                Text(story.level.capitalized)
                    .font(.caption2.weight(.medium))
                    .foregroundStyle(.teal)
                    .padding(.horizontal, 8)
                    .padding(.vertical, 3)
                    .background(.teal.opacity(0.12), in: Capsule())

                Divider()

                // Tappable Spanish text
                tappableText
                    .padding()
                    .background(.fill.quinary, in: RoundedRectangle(cornerRadius: 12))

                // Translation toggle
                VStack(alignment: .leading, spacing: 8) {
                    Button {
                        withAnimation { showTranslation.toggle() }
                    } label: {
                        Label(
                            showTranslation ? "Hide Translation" : "Show Translation",
                            systemImage: showTranslation ? "eye.slash" : "eye"
                        )
                        .font(.subheadline.weight(.medium))
                    }
                    .tint(.secondary)

                    if showTranslation {
                        Text(story.bodyEN)
                            .font(.body)
                            .foregroundStyle(.secondary)
                            .padding()
                            .frame(maxWidth: .infinity, alignment: .leading)
                            .background(.fill.quinary, in: RoundedRectangle(cornerRadius: 12))
                            .transition(.opacity.combined(with: .move(edge: .top)))
                    }
                }

                // Quiz button
                if !story.decodedQuestions.isEmpty {
                    NavigationLink {
                        StoryQuizView(story: story)
                    } label: {
                        Label("Take Comprehension Quiz", systemImage: "questionmark.circle")
                            .font(.subheadline.weight(.semibold))
                            .frame(maxWidth: .infinity)
                            .padding(.vertical, 12)
                    }
                    .buttonStyle(.borderedProminent)
                    .tint(.teal)
                }
            }
            .padding()
            .adaptiveContainer(maxWidth: 800)
        }
        .navigationTitle("Story")
        .navigationBarTitleDisplayMode(.inline)
        .sheet(item: $selectedWord) { word in
            WordDetailSheet(word: word)
                .presentationDetents([.height(200)])
        }
    }

    // MARK: - Tappable Text

    private var tappableText: some View {
        let words = story.bodyES.components(separatedBy: " ")
        let map = annotationMap
        let cache = lookupCache
        let context = story.bodyES

        return WrappingHStack(words: words) { word in
            WordButton(word: word, map: map, cache: cache) { ann in
                if ann.english.isEmpty {
                    lookupWord(ann.word, inContext: context)
                } else {
                    selectedWord = ann
                }
            }
        }
    }

    private func lookupWord(_ word: String, inContext sentence: String) {
        // Try offline dictionary first
        if let entry = dictionary.lookup(word) {
            let annotation = WordAnnotation(
                word: word,
                baseForm: entry.baseForm,
                english: entry.english,
                partOfSpeech: entry.partOfSpeech
            )
            lookupCache[word] = annotation
            selectedWord = annotation
            return
        }

        // Fall back to on-device AI lookup
        selectedWord = WordAnnotation(word: word, baseForm: word, english: "Looking up...", partOfSpeech: "")

        Task {
            do {
                let annotation = try await WordLookup.lookup(word: word, inContext: sentence)
                lookupCache[word] = annotation
                selectedWord = annotation
            } catch {
                selectedWord = WordAnnotation(word: word, baseForm: word, english: "Lookup unavailable", partOfSpeech: "")
            }
        }
    }

    private func cleanWord(_ word: String) -> String {
        word.lowercased()
            .trimmingCharacters(in: .punctuationCharacters)
            .trimmingCharacters(in: .whitespaces)
    }
}

// MARK: - Word Button

private struct WordButton: View {
    let word: String
    let map: [String: WordAnnotation]
    let cache: [String: WordAnnotation]
    let onTap: (WordAnnotation) -> Void

    private var cleaned: String {
        word.lowercased()
            .trimmingCharacters(in: .punctuationCharacters)
            .trimmingCharacters(in: .whitespaces)
    }

    private var resolvedAnnotation: WordAnnotation {
        map[cleaned] ?? cache[cleaned] ?? WordAnnotation(word: cleaned, baseForm: cleaned, english: "", partOfSpeech: "")
    }

    var body: some View {
        Button {
            onTap(resolvedAnnotation)
        } label: {
            Text(word + " ")
                .font(.body)
                .foregroundStyle(.primary)
                .underline(true, color: .teal.opacity(0.3))
        }
        .buttonStyle(.plain)
    }
}

// MARK: - Wrapping HStack

private struct WrappingHStack<Content: View>: View {
    let words: [String]
    let content: (String) -> Content

    var body: some View {
        FlowLayout(spacing: 0) {
            ForEach(Array(words.enumerated()), id: \.offset) { _, word in
                content(word)
            }
        }
        .accessibilityElement(children: .combine)
    }
}

private struct FlowLayout: Layout {
    var spacing: CGFloat = 0

    func sizeThatFits(proposal: ProposedViewSize, subviews: Subviews, cache: inout ()) -> CGSize {
        let rows = computeRows(proposal: proposal, subviews: subviews)
        var height: CGFloat = 0
        for row in rows {
            height += row.map { $0.height }.max() ?? 0
        }
        height += CGFloat(max(0, rows.count - 1)) * spacing
        return CGSize(width: proposal.width ?? 0, height: height)
    }

    func placeSubviews(in bounds: CGRect, proposal: ProposedViewSize, subviews: Subviews, cache: inout ()) {
        let rows = computeRows(proposal: proposal, subviews: subviews)
        var y = bounds.minY
        var subviewIndex = 0
        for row in rows {
            var x = bounds.minX
            let rowHeight = row.map { $0.height }.max() ?? 0
            for size in row {
                subviews[subviewIndex].place(at: CGPoint(x: x, y: y), proposal: ProposedViewSize(size))
                x += size.width
                subviewIndex += 1
            }
            y += rowHeight + spacing
        }
    }

    private func computeRows(proposal: ProposedViewSize, subviews: Subviews) -> [[CGSize]] {
        let maxWidth = proposal.width ?? .infinity
        var rows: [[CGSize]] = [[]]
        var currentWidth: CGFloat = 0

        for subview in subviews {
            let size = subview.sizeThatFits(.unspecified)
            if currentWidth + size.width > maxWidth && !rows[rows.count - 1].isEmpty {
                rows.append([])
                currentWidth = 0
            }
            rows[rows.count - 1].append(size)
            currentWidth += size.width
        }
        return rows
    }
}

// MARK: - Word Detail Sheet

private struct WordDetailSheet: View {
    let word: WordAnnotation

    var body: some View {
        VStack(spacing: 16) {
            HStack {
                Text(word.word)
                    .font(.title2.bold())
                Spacer()
                if !word.partOfSpeech.isEmpty {
                    Text(word.partOfSpeech)
                        .font(.caption.weight(.medium))
                        .foregroundStyle(.secondary)
                        .padding(.horizontal, 8)
                        .padding(.vertical, 4)
                        .background(.fill.tertiary, in: Capsule())
                }
            }

            Divider()

            if word.english == "Looking up..." {
                HStack(spacing: 8) {
                    ProgressView()
                    Text("Looking up word...")
                        .font(.subheadline)
                        .foregroundStyle(.secondary)
                }
                .frame(maxWidth: .infinity)
            } else {
                VStack(alignment: .leading, spacing: 8) {
                    if !word.baseForm.isEmpty && word.baseForm != word.word {
                        HStack {
                            Text("Base form:")
                                .font(.subheadline)
                                .foregroundStyle(.secondary)
                            Text(word.baseForm)
                                .font(.subheadline.weight(.semibold))
                                .italic()
                        }
                    }

                    if !word.english.isEmpty {
                        HStack {
                            Text("English:")
                                .font(.subheadline)
                                .foregroundStyle(.secondary)
                            Text(word.english)
                                .font(.subheadline.weight(.semibold))
                        }
                    }
                }
                .frame(maxWidth: .infinity, alignment: .leading)
            }

            Spacer()
        }
        .padding()
    }
}

// MARK: - On-Demand Word Lookup

@MainActor
private enum WordLookup {
    @Generable
    struct WordInfo {
        @Guide(description: "The dictionary base form (infinitive for verbs, singular for nouns)")
        var baseForm: String
        @Guide(description: "English translation")
        var english: String
        @Guide(description: "Part of speech: verb, noun, adjective, adverb, preposition, conjunction, article, pronoun, or other")
        var partOfSpeech: String
    }

    static func lookup(word: String, inContext sentence: String) async throws -> WordAnnotation {
        let session = LanguageModelSession(instructions: """
            You are a Spanish dictionary. Given a word and the sentence it appears in, \
            provide its base form, English translation, and part of speech.
            """)

        let response = try await session.respond(
            to: "Word: \"\(word)\" in sentence: \"\(sentence)\"",
            generating: WordInfo.self
        )

        let info = response.content
        return WordAnnotation(
            word: word,
            baseForm: info.baseForm,
            english: info.english,
            partOfSpeech: info.partOfSpeech
        )
    }
}
