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

import Foundation

@MainActor
@Observable
class DemoModel {
    // MARK: - Input Properties
    
    var query: String = "What are Batman's secret weapons?"
    
    // MARK: - State Properties
    
    enum Step {
        case initial
        case slicing
        case indexing
        case retrieval
        case reranking
        case generation
        case success
    }
    var currentStep: Step = .initial
    var isRunning: Bool {
        ![Step.initial, .success].contains(currentStep)
        && error == nil
    }

    // MARK: - RAG Pipeline Results

    var slices: [String]?
    var indexed: [(chunk: String, embedding: [Double])]?
    var retrieved: [String]?
    var reranked: [String]?
    var generated: (prompt: String, response: String)?
    var error: Error?

    // MARK: - Dependencies

    private let ragService = RAGService()

    // MARK: - Public Methods

    func runRAGPipeline() async {
        reset()

        currentStep = .slicing

        let stream = ragService.runRAGPipeline(query: query)

        for await stepResult in stream {
            if let error = stepResult.error {
                self.error = error
                return
            }

            if let slices = stepResult.slices {
                self.slices = slices
                currentStep = .indexing
            }
            if let indexed = stepResult.indexed {
                self.indexed = indexed
                currentStep = .retrieval
            }
            if let retrieved = stepResult.retrieved {
                self.retrieved = retrieved
                currentStep = .reranking
            }
            if let reranked = stepResult.reranked {
                self.reranked = reranked
                currentStep = .generation
            }
            if let generated = stepResult.generated {
                self.generated = generated
                currentStep = .success
            }
        }
    }

    // MARK: - Private Methods

    private func reset() {
        slices = nil
        indexed = nil
        retrieved = nil
        reranked = nil
        generated = nil
        error = nil
        currentStep = .initial
    }
}
