//
// Created by Banghua Zhao on 07/07/2025
// Copyright Apps Bay Limited. All rights reserved.
//

import Foundation
import SwiftUI

struct DemoView: View {
    @State private var model = DemoModel()

    // Sheet state variables
    @State private var showSlicesSheet = false
    @State private var showIndexingSheet = false
    @State private var showRetrievalSheet = false
    @State private var showRerankingSheet = false
    @State private var showGenerationSheet = false

    var body: some View {
        NavigationView {
            List {
                Section("Ask a question:") {
                    TextField("Enter your question...", text: $model.query)
                        .textFieldStyle(RoundedBorderTextFieldStyle())

                    Button {
                        Task {
                            await model.runRAGPipeline()
                        }
                    } label: {
                        if model.isRunning {
                            ProgressView()
                        } else {
                            Text("Run RAG Pipeline")
                        }
                    }
                    .disabled(model.isRunning)
                }

                Section("RAG Pipeline Result") {
                    Button {
                        showSlicesSheet = true
                    } label: {
                        HStack {
                            Text("🧩 Slicing")
                            Spacer()
                            if model.currentStep == .slicing {
                                if model.error != nil {
                                    Text("❌")
                                } else {
                                    ProgressView()
                                }
                            } else if let slices = model.slices {
                                Text("✅ \(slices.count) chunks")
                                    .foregroundColor(.secondary)
                            }
                        }
                    }
                    .sheet(isPresented: $showSlicesSheet) {
                        NavigationView {
                            List {
                                if let slices = model.slices {
                                    ForEach(slices.indices, id: \.self) { index in
                                        VStack(alignment: .leading, spacing: 4) {
                                            Text("Chunk \(index + 1)")
                                                .font(.headline)
                                            Text(slices[index])
                                                .font(.body)
                                        }
                                        .padding(.vertical, 4)
                                    }
                                } else {
                                    Text("No Result")
                                }
                            }
                            .navigationTitle("Slicing Results")
                            .navigationBarTitleDisplayMode(.inline)
                            .toolbar {
                                ToolbarItem(placement: .navigationBarTrailing) {
                                    Button("Done") {
                                        showSlicesSheet = false
                                    }
                                }
                            }
                        }
                    }

                    Button {
                        showIndexingSheet = true
                    } label: {
                        HStack {
                            Text("🗂️ Indexing")
                            Spacer()
                            if model.currentStep == .indexing {
                                if model.error != nil {
                                    Text("❌")
                                } else {
                                    ProgressView()
                                }
                            } else if let indexed = model.indexed {
                                Text("✅ \(indexed.count) embeddings")
                                    .foregroundColor(.secondary)
                            }
                        }
                    }
                    .sheet(isPresented: $showIndexingSheet) {
                        NavigationView {
                            List {
                                if let indexed = model.indexed {
                                    ForEach(indexed.indices, id: \.self) { index in
                                        VStack(alignment: .leading, spacing: 4) {
                                            Text("Chunk \(index + 1)")
                                                .font(.headline)
                                            Text(indexed[index].chunk)
                                                .lineLimit(2)
                                                .font(.body)
                                            Text("Embedding: [\(indexed[index].embedding.prefix(15).map { String(format: "%.3f", $0) }.joined(separator: ", "))...]")
                                                .font(.caption)
                                                .foregroundColor(.secondary)
                                        }
                                        .padding(.vertical, 4)
                                    }
                                } else {
                                    Text("No Result")
                                }
                            }
                            .navigationTitle("Indexing Results")
                            .navigationBarTitleDisplayMode(.inline)
                            .toolbar {
                                ToolbarItem(placement: .navigationBarTrailing) {
                                    Button("Done") {
                                        showIndexingSheet = false
                                    }
                                }
                            }
                        }
                    }

                    Button {
                        showRetrievalSheet = true
                    } label: {
                        HStack {
                            Text("🔎 Retrieval")
                            Spacer()
                            if model.currentStep == .retrieval {
                                if model.error != nil {
                                    Text("❌")
                                } else {
                                    ProgressView()
                                }
                            } else if let retrieved = model.retrieved {
                                Text("✅ \(retrieved.count) retrieved")
                                    .foregroundColor(.secondary)
                            }
                        }
                    }
                    .sheet(isPresented: $showRetrievalSheet) {
                        NavigationView {
                            List {
                                if let retrieved = model.retrieved {
                                    ForEach(retrieved.indices, id: \.self) { index in
                                        VStack(alignment: .leading, spacing: 4) {
                                            Text("Retrieved \(index + 1)")
                                                .font(.headline)
                                            Text(retrieved[index])
                                                .font(.body)
                                        }
                                        .padding(.vertical, 4)
                                    }
                                } else {
                                    Text("No Result")
                                }
                            }
                            .navigationTitle("Retrieval Results")
                            .navigationBarTitleDisplayMode(.inline)
                            .toolbar {
                                ToolbarItem(placement: .navigationBarTrailing) {
                                    Button("Done") {
                                        showRetrievalSheet = false
                                    }
                                }
                            }
                        }
                    }

                    Button {
                        showRerankingSheet = true
                    } label: {
                        HStack {
                            Text("🔄 Reranking")
                            Spacer()
                            if model.currentStep == .reranking {
                                if model.error != nil {
                                    Text("❌")
                                } else {
                                    ProgressView()
                                }
                            } else if let reranked = model.reranked {
                                Text("✅ \(reranked.count) reranked")
                                    .foregroundColor(.secondary)
                            }
                        }
                    }
                    .sheet(isPresented: $showRerankingSheet) {
                        NavigationView {
                            List {
                                if let reranked = model.reranked {
                                    ForEach(reranked.indices, id: \.self) { index in
                                        VStack(alignment: .leading, spacing: 4) {
                                            Text("Reranked \(index + 1)")
                                                .font(.headline)
                                            Text(reranked[index])
                                                .font(.body)
                                        }
                                        .padding(.vertical, 4)
                                    }
                                } else {
                                    Text("No Result")
                                }
                            }
                            .navigationTitle("Reranking Results")
                            .navigationBarTitleDisplayMode(.inline)
                            .toolbar {
                                ToolbarItem(placement: .navigationBarTrailing) {
                                    Button("Done") {
                                        showRerankingSheet = false
                                    }
                                }
                            }
                        }
                    }

                    Button {
                        showGenerationSheet = true
                    } label: {
                        HStack {
                            Text("📝 Generation")
                            Spacer()
                            if model.currentStep == .generation {
                                if model.error != nil {
                                    Text("❌")
                                } else {
                                    ProgressView()
                                }
                            } else if model.generated != nil {
                                Text("✅ Generated")
                                    .foregroundColor(.secondary)
                            }
                        }
                    }
                    .sheet(isPresented: $showGenerationSheet) {
                        NavigationView {
                            ScrollView {
                                if let generated = model.generated {
                                    VStack(alignment: .leading, spacing: 16) {
                                        Text("Prompt:")
                                            .font(.headline)
                                        Text(generated.0)
                                            .font(.body)
                                        Text("Generated Answer")
                                            .font(.headline)
                                        Text(generated.1)
                                            .font(.body)
                                    }
                                    .padding()
                                } else {
                                    Text("No Result")
                                }
                            }
                            .navigationTitle("Generation Result")
                            .navigationBarTitleDisplayMode(.inline)
                            .toolbar {
                                ToolbarItem(placement: .navigationBarTrailing) {
                                    Button("Done") {
                                        showGenerationSheet = false
                                    }
                                }
                            }
                        }
                    }

                    if model.slices == nil && model.indexed == nil && model.retrieved == nil && model.reranked == nil && model.generated == nil && model.error == nil {
                        Text("No results yet. Run the RAG pipeline to see results.")
                            .foregroundColor(.secondary)
                            .italic()
                    } else if let error = model.error {
                        Text("Error: \(error.localizedDescription)")
                            .foregroundColor(.red)
                            .italic()
                    }
                }
            }
            .navigationTitle("Swift RAG Demo")
        }
    }
}

#Preview {
    DemoView()
}
