LabWiseiOS/LabWise/ScanView.swift

import SwiftUI
import VisionKit
import LabWiseKit

// MARK: - Scan stages

enum ScanStage {
    case camera
    case analyzing
    case review
}

// MARK: - Extracted label data (Foundation Models @Generable output)

import FoundationModels

@Generable
struct ExtractedLabelData {
    @Guide(description: "CAS registry number, e.g. 7732-18-5")
    var casNumber: String

    @Guide(description: "Full chemical name")
    var chemicalName: String

    @Guide(description: "Chemical formula, e.g. H2O")
    var chemicalFormula: String

    @Guide(description: "Molecular weight with unit, e.g. 18.015 g/mol")
    var molecularWeight: String

    @Guide(description: "Vendor or manufacturer name")
    var vendor: String

    @Guide(description: "Catalog number from vendor")
    var catalogNumber: String

    @Guide(description: "Lot number")
    var lotNumber: String

    @Guide(description: "Concentration, e.g. 1M, 98%")
    var concentration: String

    @Guide(description: "Expiration date in ISO format YYYY-MM-DD if found")
    var expirationDate: String

    @Guide(description: "Physical state: liquid, solid, or gas")
    var physicalState: String

    @Guide(description: "Storage device inferred from label: flammable cabinet, fridge, freezer, ambient, etc.")
    var storageDevice: String

    @Guide(description: "Container type: bottle, cylinder, drum, vial, etc.")
    var containerType: String
}

// MARK: - ScanViewModel

@Observable
final class ScanViewModel {
    var stage: ScanStage = .camera
    var capturedImage: UIImage?
    var capturedTexts: [String] = []
    var scannedBarcode: String?

    // Extracted fields
    var casNumber = ""
    var chemicalName = ""
    var chemicalFormula = ""
    var molecularWeight = ""
    var vendor = ""
    var catalogNumber = ""
    var lotNumber = ""
    var concentration = ""
    var expirationDate = ""
    var physicalState = ""
    var storageDevice = ""
    var numberOfContainers = "1"
    var amountPerContainer = ""
    var unitOfMeasure = ""
    var percentageFull: Double = 50

    // Pre-populated from profile
    var piFirstName = ""
    var bldgCode = ""
    var lab = ""
    var storageLocation = ""
    var contact = ""

    var analysisError: String?
    var isSaving = false
    var saveError: String?
    var didSave = false

    private let chemicalsClient = ChemicalsClient()
    private let profileClient = ProfileClient()

    func loadProfile() async {
        if let profile = try? await profileClient.get() {
            piFirstName = profile.piFirstName
            bldgCode = profile.bldgCode
            lab = profile.lab
            contact = profile.contact ?? ""
        }
    }

    func analyzeTexts() async {
        stage = .analyzing
        analysisError = nil

        guard SystemLanguageModel.default.isAvailable else {
            // Fall through to review with empty fields
            stage = .review
            return
        }

        do {
            let combinedText = capturedTexts.joined(separator: "\n")
            let session = LanguageModelSession(instructions: """
                You are a chemical label OCR assistant. Extract structured fields from the raw OCR text of a chemical reagent label.
                Return only the fields you can identify. Leave fields empty string if not found.
                For storageDevice infer from hazard symbols: FLAMMABLE → flammable cabinet, KEEP REFRIGERATED → fridge, KEEP FROZEN → freezer.
                """)
            let response = try await session.respond(
                to: "Extract fields from this chemical label text:\n\n\(combinedText)",
                generating: ExtractedLabelData.self
            )
            let data = response.content
            await MainActor.run {
                casNumber = data.casNumber
                chemicalName = data.chemicalName
                chemicalFormula = data.chemicalFormula
                molecularWeight = data.molecularWeight
                vendor = data.vendor
                catalogNumber = data.catalogNumber
                lotNumber = data.lotNumber
                concentration = data.concentration
                expirationDate = data.expirationDate
                physicalState = data.physicalState
                storageDevice = data.storageDevice
            }
        } catch {
            await MainActor.run {
                analysisError = "Label analysis failed: \(error.localizedDescription)"
            }
        }

        await MainActor.run { stage = .review }
    }

    var missingRequiredFields: [String] {
        var missing: [String] = []
        if piFirstName.trimmingCharacters(in: .whitespaces).isEmpty { missing.append("PI First Name") }
        if bldgCode.trimmingCharacters(in: .whitespaces).isEmpty { missing.append("Building Code") }
        if lab.trimmingCharacters(in: .whitespaces).isEmpty { missing.append("Lab") }
        if storageLocation.trimmingCharacters(in: .whitespaces).isEmpty { missing.append("Storage Location") }
        if chemicalName.trimmingCharacters(in: .whitespaces).isEmpty { missing.append("Chemical Name") }
        if casNumber.trimmingCharacters(in: .whitespaces).isEmpty { missing.append("CAS Number") }
        return missing
    }

    func save() async {
        isSaving = true
        saveError = nil
        defer { isSaving = false }
        do {
            let body = ChemicalCreateBody(
                piFirstName: piFirstName,
                physicalState: physicalState.isEmpty ? "unknown" : physicalState,
                chemicalName: chemicalName,
                bldgCode: bldgCode,
                lab: lab,
                storageLocation: storageLocation,
                storageDevice: storageDevice.isEmpty ? "unknown" : storageDevice,
                numberOfContainers: numberOfContainers.isEmpty ? "1" : numberOfContainers,
                amountPerContainer: amountPerContainer.isEmpty ? "unknown" : amountPerContainer,
                unitOfMeasure: unitOfMeasure.isEmpty ? "unknown" : unitOfMeasure,
                casNumber: casNumber,
                chemicalFormula: chemicalFormula.isEmpty ? nil : chemicalFormula,
                molecularWeight: molecularWeight.isEmpty ? nil : molecularWeight,
                vendor: vendor.isEmpty ? nil : vendor,
                catalogNumber: catalogNumber.isEmpty ? nil : catalogNumber,
                lotNumber: lotNumber.isEmpty ? nil : lotNumber,
                expirationDate: expirationDate.isEmpty ? nil : expirationDate,
                concentration: concentration.isEmpty ? nil : concentration,
                percentageFull: percentageFull,
                barcode: scannedBarcode,
                contact: contact.isEmpty ? nil : contact
            )
            _ = try await chemicalsClient.create(body)
            await MainActor.run { didSave = true }
        } catch {
            await MainActor.run { saveError = "Failed to save: \(error.localizedDescription)" }
        }
    }

    func reset() {
        stage = .camera
        capturedImage = nil
        capturedTexts = []
        scannedBarcode = nil
        casNumber = ""; chemicalName = ""; chemicalFormula = ""
        molecularWeight = ""; vendor = ""; catalogNumber = ""
        lotNumber = ""; concentration = ""; expirationDate = ""
        physicalState = ""; storageDevice = ""; storageLocation = ""
        numberOfContainers = "1"; amountPerContainer = ""; unitOfMeasure = ""
        percentageFull = 50; analysisError = nil; didSave = false; saveError = nil
    }
}

// MARK: - Main ScanView

struct ScanView: View {
    @State private var viewModel = ScanViewModel()
    @State private var showChemicalsList = false

    var body: some View {
        NavigationStack {
            Group {
                switch viewModel.stage {
                case .camera:
                    DataScannerWrapperView(viewModel: viewModel)
                case .analyzing:
                    AnalyzingView()
                case .review:
                    ReviewFormView(viewModel: viewModel, onSaved: {
                        showChemicalsList = true
                    })
                }
            }
            .navigationTitle(stageName)
            .navigationBarTitleDisplayMode(.inline)
            .toolbar {
                if viewModel.stage != .camera {
                    ToolbarItem(placement: .topBarLeading) {
                        Button("Rescan") { viewModel.reset() }
                    }
                }
            }
        }
        .task { await viewModel.loadProfile() }
        .onChange(of: viewModel.didSave) { _, saved in
            if saved { showChemicalsList = true }
        }
    }

    private var stageName: String {
        switch viewModel.stage {
        case .camera: return "Scan Label"
        case .analyzing: return "Analysing..."
        case .review: return "Review"
        }
    }
}

// MARK: - Analysing placeholder

struct AnalyzingView: View {
    var body: some View {
        VStack(spacing: 20) {
            ProgressView()
                .scaleEffect(1.5)
            Text("Analysing label...")
                .font(.headline)
                .foregroundStyle(.secondary)
        }
        .frame(maxWidth: .infinity, maxHeight: .infinity)
    }
}

// MARK: - DataScanner wrapper

struct DataScannerWrapperView: View {
    let viewModel: ScanViewModel
    @State private var recognizedTexts: [RecognizedItem] = []
    @State private var showUnsupported = false

    var body: some View {
        ZStack(alignment: .bottom) {
            if DataScannerViewController.isSupported && DataScannerViewController.isAvailable {
                DataScannerRepresentable(
                    recognizedItems: $recognizedTexts,
                    onCapture: { image, texts, barcode in
                        viewModel.capturedImage = image
                        viewModel.capturedTexts = texts
                        viewModel.scannedBarcode = barcode
                        Task { await viewModel.analyzeTexts() }
                    }
                )
                .ignoresSafeArea()

                // Capture button
                VStack {
                    Spacer()
                    CaptureButton {
                        // Capture handled inside representable via button tap callback
                    }
                    .padding(.bottom, 40)
                }
            } else {
                ContentUnavailableView(
                    "Scanner Unavailable",
                    systemImage: "camera.slash",
                    description: Text("This device does not support data scanning.")
                )
            }
        }
    }
}

struct CaptureButton: View {
    let action: () -> Void

    var body: some View {
        Button(action: action) {
            ZStack {
                Circle()
                    .fill(.white)
                    .frame(width: 72, height: 72)
                Circle()
                    .strokeBorder(.white, lineWidth: 4)
                    .frame(width: 84, height: 84)
            }
        }
    }
}

// MARK: - UIViewControllerRepresentable for DataScanner

struct DataScannerRepresentable: UIViewControllerRepresentable {
    @Binding var recognizedItems: [RecognizedItem]
    let onCapture: (UIImage, [String], String?) -> Void

    func makeUIViewController(context: Context) -> DataScannerViewController {
        let scanner = DataScannerViewController(
            recognizedDataTypes: [.text(), .barcode()],
            qualityLevel: .balanced,
            recognizesMultipleItems: true,
            isHighFrameRateTrackingEnabled: true,
            isHighlightingEnabled: true
        )
        scanner.delegate = context.coordinator
        context.coordinator.scanner = scanner
        context.coordinator.onCapture = onCapture
        try? scanner.startScanning()
        return scanner
    }

    func updateUIViewController(_ uiViewController: DataScannerViewController, context: Context) {}

    func makeCoordinator() -> Coordinator {
        Coordinator(recognizedItems: $recognizedItems)
    }

    class Coordinator: NSObject, DataScannerViewControllerDelegate {
        @Binding var recognizedItems: [RecognizedItem]
        weak var scanner: DataScannerViewController?
        var onCapture: ((UIImage, [String], String?) -> Void)?

        init(recognizedItems: Binding<[RecognizedItem]>) {
            _recognizedItems = recognizedItems
        }

        func dataScanner(_ dataScanner: DataScannerViewController, didAdd addedItems: [RecognizedItem], allItems: [RecognizedItem]) {
            recognizedItems = allItems
        }

        func dataScanner(_ dataScanner: DataScannerViewController, didUpdate updatedItems: [RecognizedItem], allItems: [RecognizedItem]) {
            recognizedItems = allItems
        }

        func dataScanner(_ dataScanner: DataScannerViewController, didRemove removedItems: [RecognizedItem], allItems: [RecognizedItem]) {
            recognizedItems = allItems
        }

        func captureCurrentFrame() {
            guard let scanner else { return }

            // Collect all recognized text and barcodes
            var texts: [String] = []
            var barcode: String?

            for item in recognizedItems {
                switch item {
                case .text(let t):
                    texts.append(t.transcript)
                case .barcode(let b):
                    if let payload = b.payloadStringValue {
                        barcode = payload
                        texts.append(payload)
                    }
                default:
                    break
                }
            }

            // Capture a snapshot of the current frame
            let renderer = UIGraphicsImageRenderer(bounds: scanner.view.bounds)
            let image = renderer.image { ctx in
                scanner.view.layer.render(in: ctx.cgContext)
            }

            onCapture?(image, texts, barcode)
        }
    }
}

// MARK: - ReviewFormView

struct ReviewFormView: View {
    @Bindable var viewModel: ScanViewModel
    let onSaved: () -> Void

    var body: some View {
        ScrollView {
            VStack(spacing: 0) {
                // Captured frame + percentage slider
                if let image = viewModel.capturedImage {
                    CapturedImageWithSlider(image: image, percentageFull: $viewModel.percentageFull)
                        .frame(height: 260)
                        .clipped()
                }

                if let error = viewModel.analysisError {
                    Text(error)
                        .foregroundStyle(.orange)
                        .font(.footnote)
                        .padding(.horizontal)
                        .padding(.top, 8)
                }

                Form {
                    // Required fields section
                    Section {
                        ReviewField("Chemical Name", text: $viewModel.chemicalName, required: true, missing: viewModel.missingRequiredFields.contains("Chemical Name"))
                        ReviewField("CAS Number", text: $viewModel.casNumber, required: true, missing: viewModel.missingRequiredFields.contains("CAS Number"))
                        ReviewField("Physical State", text: $viewModel.physicalState)
                        ReviewField("Storage Device", text: $viewModel.storageDevice)
                    } header: {
                        Text("Identity")
                    }

                    Section {
                        ReviewField("PI First Name", text: $viewModel.piFirstName, required: true, missing: viewModel.missingRequiredFields.contains("PI First Name"))
                        ReviewField("Building Code", text: $viewModel.bldgCode, required: true, missing: viewModel.missingRequiredFields.contains("Building Code"))
                        ReviewField("Lab", text: $viewModel.lab, required: true, missing: viewModel.missingRequiredFields.contains("Lab"))
                        ReviewField("Storage Location", text: $viewModel.storageLocation, required: true, missing: viewModel.missingRequiredFields.contains("Storage Location"))
                    } header: {
                        Text("Location")
                    }

                    Section {
                        ReviewField("# Containers", text: $viewModel.numberOfContainers)
                        ReviewField("Amount / Container", text: $viewModel.amountPerContainer)
                        ReviewField("Unit of Measure", text: $viewModel.unitOfMeasure)
                    } header: {
                        Text("Quantity")
                    }

                    Section {
                        ReviewField("Formula", text: $viewModel.chemicalFormula)
                        ReviewField("Molecular Weight", text: $viewModel.molecularWeight)
                        ReviewField("Concentration", text: $viewModel.concentration)
                        ReviewField("Vendor", text: $viewModel.vendor)
                        ReviewField("Catalog #", text: $viewModel.catalogNumber)
                        ReviewField("Lot #", text: $viewModel.lotNumber)
                        ReviewField("Expiration Date", text: $viewModel.expirationDate)
                    } header: {
                        Text("Details")
                    }

                    Section {
                        if let saveError = viewModel.saveError {
                            Text(saveError)
                                .foregroundStyle(.red)
                                .font(.footnote)
                        }
                        Button {
                            Task { await viewModel.save() }
                        } label: {
                            if viewModel.isSaving {
                                ProgressView()
                                    .frame(maxWidth: .infinity)
                            } else {
                                Text("Save Chemical")
                                    .frame(maxWidth: .infinity)
                                    .fontWeight(.semibold)
                            }
                        }
                        .buttonStyle(.borderedProminent)
                        .disabled(!viewModel.missingRequiredFields.isEmpty || viewModel.isSaving)
                    }

                    if !viewModel.missingRequiredFields.isEmpty {
                        Section {
                            Text("Required: \(viewModel.missingRequiredFields.joined(separator: ", "))")
                                .foregroundStyle(.red)
                                .font(.footnote)
                        }
                    }
                }
                .frame(minHeight: 800)
                .scrollDisabled(true)
            }
        }
    }
}

// MARK: - Captured image with vertical % slider

struct CapturedImageWithSlider: View {
    let image: UIImage
    @Binding var percentageFull: Double

    var body: some View {
        ZStack(alignment: .trailing) {
            Image(uiImage: image)
                .resizable()
                .scaledToFill()

            // Dim overlay
            Color.black.opacity(0.2)

            // Vertical slider on right edge
            HStack {
                Spacer()
                VStack(spacing: 8) {
                    Text("\(Int(percentageFull))%")
                        .font(.caption.bold())
                        .foregroundStyle(.white)
                        .shadow(radius: 2)

                    VerticalSlider(value: $percentageFull, range: 0...100)
                        .frame(width: 32, height: 180)
                }
                .padding(.trailing, 12)
            }
        }
    }
}

// MARK: - Vertical slider

struct VerticalSlider: View {
    @Binding var value: Double
    let range: ClosedRange<Double>

    @GestureState private var isDragging = false

    var body: some View {
        GeometryReader { geo in
            ZStack(alignment: .bottom) {
                // Track
                Capsule()
                    .fill(Color.white.opacity(0.3))

                // Fill
                Capsule()
                    .fill(Color.white.opacity(0.8))
                    .frame(height: geo.size.height * fillFraction)

                // Thumb
                Circle()
                    .fill(.white)
                    .frame(width: 28, height: 28)
                    .shadow(radius: 3)
                    .offset(y: -geo.size.height * fillFraction + 14)
            }
            .gesture(
                DragGesture(minimumDistance: 0)
                    .onChanged { drag in
                        let fraction = 1 - (drag.location.y / geo.size.height)
                        let clamped = max(range.lowerBound, min(range.upperBound, fraction * (range.upperBound - range.lowerBound) + range.lowerBound))
                        value = clamped
                    }
            )
        }
    }

    private var fillFraction: Double {
        (value - range.lowerBound) / (range.upperBound - range.lowerBound)
    }
}

// MARK: - ReviewField helper

struct ReviewField: View {
    let label: String
    @Binding var text: String
    var required: Bool = false
    var missing: Bool = false

    init(_ label: String, text: Binding<String>, required: Bool = false, missing: Bool = false) {
        self.label = label
        self._text = text
        self.required = required
        self.missing = missing
    }

    var body: some View {
        HStack {
            Text(label + (required ? " *" : ""))
                .foregroundStyle(missing ? .red : .primary)
                .frame(minWidth: 120, alignment: .leading)
            TextField(label, text: $text)
                .multilineTextAlignment(.trailing)
        }
        .listRowBackground(missing ? Color.red.opacity(0.05) : Color(uiColor: .secondarySystemGroupedBackground))
    }
}