// lib/screens/river/investigative/manual_sampling/river_inves_in_situ_step_3_data_capture.dart

import 'dart:async';
import 'package:flutter/material.dart';
import 'package:provider/provider.dart';
import 'package:flutter_bluetooth_serial/flutter_bluetooth_serial.dart';
import 'package:usb_serial/usb_serial.dart';
import 'package:intl/intl.dart';

import '../../../../auth_provider.dart';
import '../../../../models/river_inves_manual_sampling_data.dart'; // Updated model
//import '../../../../services/api_service.dart';
import 'package:environment_monitoring_app/services/database_helper.dart';
import '../../../../services/river_investigative_sampling_service.dart'; // Updated service
import '../../../../bluetooth/bluetooth_manager.dart';
import '../../../../serial/serial_manager.dart';
import '../../../../bluetooth/widgets/bluetooth_device_list_dialog.dart';
import '../../../../serial/widget/serial_port_list_dialog.dart';

class RiverInvesStep3DataCapture extends StatefulWidget {
  final RiverInvesManualSamplingData data;
  final VoidCallback onNext;

  const RiverInvesStep3DataCapture({
    super.key,
    required this.data,
    required this.onNext,
  });

  @override
  State<RiverInvesStep3DataCapture> createState() =>
      _RiverInvesStep3DataCaptureState();
}

class _RiverInvesStep3DataCaptureState extends State<RiverInvesStep3DataCapture>
    with WidgetsBindingObserver {
  final _formKey = GlobalKey<FormState>();
  bool _isLoading = false;
  bool _isAutoReading = false;
  StreamSubscription? _dataSubscription;
  Timer? _lockoutTimer;
  int _lockoutSecondsRemaining = 30;
  bool _isLockedOut = false;

  late final RiverInvestigativeSamplingService _samplingService; // Updated service type
  final DatabaseHelper _dbHelper = DatabaseHelper();

  Map<String, double>? _previousReadingsForComparison;
  Set<String> _outOfBoundsKeys = {};

  // Parameter mappings and definitions remain the same as River In-Situ
  final Map<String, String> _parameterKeyToLimitName = const {
    'oxygenConcentration': 'Oxygen Conc',
    'oxygenSaturation': 'Oxygen Sat',
    'ph': 'pH',
    'salinity': 'Salinity',
    'electricalConductivity': 'Conductivity',
    'temperature': 'Temperature',
    'tds': 'TDS',
    'turbidity': 'Turbidity',
    'ammonia': 'Ammonia',
    'batteryVoltage': 'Battery',
  };
  final List<Map<String, dynamic>> _parameters = [];

  // Controllers remain the same
  final _sondeIdController = TextEditingController();
  final _dateController = TextEditingController();
  final _timeController = TextEditingController();
  final _oxyConcController = TextEditingController();
  final _oxySatController = TextEditingController();
  final _phController = TextEditingController();
  final _salinityController = TextEditingController();
  final _ecController = TextEditingController();
  final _tempController = TextEditingController();
  final _tdsController = TextEditingController();
  final _turbidityController = TextEditingController();
  final _ammoniaController = TextEditingController();
  final _batteryController = TextEditingController();
  String? _selectedFlowrateMethod;
  final _flowrateValueController = TextEditingController();
  final _sdHeightController = TextEditingController();
  final _sdDistanceController = TextEditingController();
  final _sdTimeFirstController = TextEditingController();
  final _sdTimeLastController = TextEditingController();

  @override
  void initState() {
    super.initState();
    // Use the Investigative service
    _samplingService = Provider.of<RiverInvestigativeSamplingService>(context, listen: false);
    _initializeControllers();
    _initializeFlowrateControllers();
    WidgetsBinding.instance.addObserver(this);
  }

  @override
  void dispose() {
    _dataSubscription?.cancel();
    _lockoutTimer?.cancel();
    // Ensure disconnect calls use the correct service instance
    if (_samplingService.bluetoothConnectionState.value != BluetoothConnectionState.disconnected) {
      _samplingService.disconnectFromBluetooth();
    }
    if (_samplingService.serialConnectionState.value != SerialConnectionState.disconnected) {
      _samplingService.disconnectFromSerial();
    }
    _disposeControllers();
    _disposeFlowrateControllers();
    WidgetsBinding.instance.removeObserver(this);
    super.dispose();
  }

  @override
  void didChangeAppLifecycleState(AppLifecycleState state) {
    if (state == AppLifecycleState.resumed) {
      if (mounted) {
        // --- START MODIFICATION ---
        final service = _samplingService;
        final btConnecting = service.bluetoothConnectionState.value == BluetoothConnectionState.connecting;
        final serialConnecting = service.serialConnectionState.value == SerialConnectionState.connecting;

        // If the widget's local state is loading OR the service's state is stuck connecting
        if (_isLoading || btConnecting || serialConnecting) {
          // Force-call disconnect to reset both the service's state
          // and the local _isLoading flag (inside _disconnect).
          _disconnectFromAll();
        } else {
          // If not stuck, just a normal refresh
          setState(() {});
        }
        // --- END MODIFICATION ---
      }
    }
  }

  // --- All helper methods (_initializeControllers, _disposeControllers, _initializeFlowrateControllers,
  //      _disposeFlowrateControllers, _onFlowrateMethodChanged, _calculateFlowrate, _selectTime,
  //      _handleConnectionAttempt, _connectToDevice, _startLockoutTimer, _toggleAutoReading,
  //      _disconnect, _disconnectFromAll, _updateTextFields, _validateAndProceed, _captureReadingsToMap,
  //      _validateParameters, _saveDataAndMoveOn, _showSnackBar, _showStopReadingDialog,
  //      _getActiveConnectionDetails, _showParameterLimitDialog, _buildFlowrateSection, etc.)
  //      are copied directly from river_in_situ_step_3_data_capture.dart
  //      but ensure they use the _samplingService instance of type RiverInvestigativeSamplingService.
  // ---

  void _initializeControllers() {
    // Logic copied from RiverInSituStep3DataCaptureState._initializeControllers
    widget.data.dataCaptureDate = widget.data.samplingDate;
    widget.data.dataCaptureTime = widget.data.samplingTime;

    _sondeIdController.text = widget.data.sondeId ?? '';
    _dateController.text = widget.data.dataCaptureDate ?? '';
    _timeController.text = widget.data.dataCaptureTime ?? '';

    _oxyConcController.text = widget.data.oxygenConcentration?.toString() ?? '-999.0';
    _oxySatController.text = widget.data.oxygenSaturation?.toString() ?? '-999.0';
    _phController.text = widget.data.ph?.toString() ?? '-999.0';
    _salinityController.text = widget.data.salinity?.toString() ?? '-999.0';
    _ecController.text = widget.data.electricalConductivity?.toString() ?? '-999.0';
    _tempController.text = widget.data.temperature?.toString() ?? '-999.0';
    _tdsController.text = widget.data.tds?.toString() ?? '-999.0';
    _turbidityController.text = widget.data.turbidity?.toString() ?? '-999.0';
    _ammoniaController.text = widget.data.ammonia?.toString() ?? '-999.0';
    _batteryController.text = widget.data.batteryVoltage?.toString() ?? '-999.0';

    if (_parameters.isEmpty) {
      _parameters.addAll([
        {'key': 'oxygenConcentration', 'icon': Icons.air, 'label': 'Oxygen Conc.', 'unit': 'mg/L', 'controller': _oxyConcController},
        {'key': 'oxygenSaturation', 'icon': Icons.percent, 'label': 'Oxygen Sat.', 'unit': '%', 'controller': _oxySatController},
        {'key': 'ph', 'icon': Icons.science_outlined, 'label': 'pH', 'unit': '', 'controller': _phController},
        {'key': 'salinity', 'icon': Icons.waves, 'label': 'Salinity', 'unit': 'ppt', 'controller': _salinityController},
        {'key': 'electricalConductivity', 'icon': Icons.flash_on, 'label': 'Conductivity', 'unit': 'µS/cm', 'controller': _ecController},
        {'key': 'temperature', 'icon': Icons.thermostat, 'label': 'Temperature', 'unit': '°C', 'controller': _tempController},
        {'key': 'tds', 'icon': Icons.grain, 'label': 'TDS', 'unit': 'mg/L', 'controller': _tdsController},
        {'key': 'turbidity', 'icon': Icons.opacity, 'label': 'Turbidity', 'unit': 'NTU', 'controller': _turbidityController},
        {'key': 'ammonia', 'icon': Icons.science, 'label': 'Ammonia', 'unit': 'mg/L', 'controller': _ammoniaController},
        {'key': 'batteryVoltage', 'icon': Icons.battery_charging_full, 'label': 'Battery', 'unit': 'V', 'controller': _batteryController},
      ]);
    }
  }

  void _disposeControllers() {
    // Logic copied from RiverInSituStep3DataCaptureState._disposeControllers
    _sondeIdController.dispose();
    _dateController.dispose();
    _timeController.dispose();
    _oxyConcController.dispose();
    _oxySatController.dispose();
    _phController.dispose();
    _salinityController.dispose();
    _ecController.dispose();
    _tempController.dispose();
    _tdsController.dispose();
    _turbidityController.dispose();
    _ammoniaController.dispose();
    _batteryController.dispose();
  }

  void _initializeFlowrateControllers() {
    // Logic copied from RiverInSituStep3DataCaptureState._initializeFlowrateControllers
    _selectedFlowrateMethod = widget.data.flowrateMethod;
    _flowrateValueController.text = widget.data.flowrateValue?.toString() ?? '';
    _sdHeightController.text = widget.data.flowrateSurfaceDrifterHeight?.toString() ?? '';
    _sdDistanceController.text = widget.data.flowrateSurfaceDrifterDistance?.toString() ?? '';
    _sdTimeFirstController.text = widget.data.flowrateSurfaceDrifterTimeFirst ?? '';
    _sdTimeLastController.text = widget.data.flowrateSurfaceDrifterTimeLast ?? '';
  }

  void _disposeFlowrateControllers() {
    // Logic copied from RiverInSituStep3DataCaptureState._disposeFlowrateControllers
    _flowrateValueController.dispose();
    _sdHeightController.dispose();
    _sdDistanceController.dispose();
    _sdTimeFirstController.dispose();
    _sdTimeLastController.dispose();
  }

  void _onFlowrateMethodChanged(String? value) {
    // Logic copied from RiverInSituStep3DataCaptureState._onFlowrateMethodChanged
    setState(() {
      _selectedFlowrateMethod = value;
      widget.data.flowrateMethod = value; // Update model immediately
      if (value == 'NA') {
        _flowrateValueController.text = 'NA';
      } else if (value == 'Flowmeter') {
        // Keep existing value if user switches back, or clear if desired
        // _flowrateValueController.clear();
        _sdHeightController.clear();
        _sdDistanceController.clear();
        _sdTimeFirstController.clear();
        _sdTimeLastController.clear();

      } else { // Surface Drifter
        // _flowrateValueController.clear(); // Will be calculated
      }
    });
  }

  void _calculateFlowrate() {
    // Logic copied from RiverInSituStep3DataCaptureState._calculateFlowrate
    final distance = double.tryParse(_sdDistanceController.text);
    final timeFirstStr = _sdTimeFirstController.text;
    final timeLastStr = _sdTimeLastController.text;

    if (distance == null || timeFirstStr.isEmpty || timeLastStr.isEmpty) {
      _showSnackBar("Please fill in Distance, Time First, and Time Last.", isError: true);
      return;
    }

    try {
      final timeFormat = DateFormat("HH:mm:ss");
      // Use a common date (like today) to allow time difference calculation across midnight
      final now = DateTime.now();
      final timeFirst = timeFormat.parse(timeFirstStr);
      final dateTimeFirst = DateTime(now.year, now.month, now.day, timeFirst.hour, timeFirst.minute, timeFirst.second);

      final timeLast = timeFormat.parse(timeLastStr);
      var dateTimeLast = DateTime(now.year, now.month, now.day, timeLast.hour, timeLast.minute, timeLast.second);

      // Handle crossing midnight
      if (dateTimeLast.isBefore(dateTimeFirst)) {
        dateTimeLast = dateTimeLast.add(const Duration(days: 1));
      }

      final differenceInSeconds = dateTimeLast.difference(dateTimeFirst).inSeconds;

      if (differenceInSeconds <= 0) {
        _showSnackBar("Time Last Deploy must be after Time First Deploy.", isError: true);
        return;
      }
      final flowrate = distance / differenceInSeconds;
      setState(() {
        _flowrateValueController.text = flowrate.toStringAsFixed(4);
      });
    } catch (e) {
      _showSnackBar("Invalid time format. Please use HH:mm:ss.", isError: true);
    }
  }

  Future<void> _selectTime(BuildContext context, TextEditingController controller) async {
    // Logic copied from RiverInSituStep3DataCaptureState._selectTime
    final TimeOfDay? picked = await showTimePicker(
      context: context,
      initialTime: TimeOfDay.now(),
    );
    if (picked != null) {
      final now = DateTime.now();
      final dt = DateTime(now.year, now.month, now.day, picked.hour, picked.minute);
      setState(() {
        controller.text = DateFormat('HH:mm:ss').format(dt);
      });
    }
  }

  Future<void> _handleConnectionAttempt(String type) async {
    // Logic copied from RiverInSituStep3DataCaptureState._handleConnectionAttempt
    // Uses the correct _samplingService instance
    final bool hasPermissions = await _samplingService.requestDevicePermissions();
    if (!hasPermissions && mounted) {
      _showSnackBar("Bluetooth & Location permissions are required to connect.", isError: true);
      return;
    }
    _disconnectFromAll();
    await Future.delayed(const Duration(milliseconds: 250)); // Short delay after disconnect
    final bool connectionSuccess = await _connectToDevice(type);

    if (connectionSuccess && mounted) {
      _dataSubscription?.cancel(); // Cancel previous subscription if any
      final stream = type == 'bluetooth' ? _samplingService.bluetoothDataStream : _samplingService.serialDataStream;
      _dataSubscription = stream.listen((readings) {
        if (mounted) {
          _updateTextFields(readings);
        }
      }, onError: (error) {
        debugPrint("Error on data stream: $error");
        if (mounted) _showSnackBar("Data stream error: $error", isError: true);
        _disconnect(type); // Disconnect on stream error
      }, onDone: () {
        debugPrint("Data stream done.");
        if (mounted) _disconnect(type); // Disconnect when stream closes
      });
    }
  }

  Future<bool> _connectToDevice(String type) async {
    // Logic copied from RiverInSituStep3DataCaptureState._connectToDevice
    // Uses the correct _samplingService instance
    setState(() => _isLoading = true);
    bool success = false;
    try {
      if (type == 'bluetooth') {
        final devices = await _samplingService.getPairedBluetoothDevices();
        if (!mounted) return false; // Check mounted after async gap
        if (devices.isEmpty) {
          _showSnackBar('No paired Bluetooth devices found.', isError: true);
          return false;
        }
        final selectedDevice = await showBluetoothDeviceListDialog(context: context, devices: devices);
        if (selectedDevice != null) {
          await _samplingService.connectToBluetoothDevice(selectedDevice);
          success = true;
        }
      } else if (type == 'serial') {
        final devices = await _samplingService.getAvailableSerialDevices();
        if (!mounted) return false;
        if (devices.isEmpty) {
          _showSnackBar('No USB Serial devices found. Ensure device is plugged in.', isError: true);
          return false;
        }
        final selectedDevice = await showSerialPortListDialog(context: context, devices: devices);
        if (selectedDevice != null) {
          await _samplingService.connectToSerialDevice(selectedDevice);
          success = true;
        }
      }
    } catch (e) {
      if (mounted) _showSnackBar('Connection failed: $e', isError: true);
      success = false;
    } finally {
      if (mounted) setState(() => _isLoading = false);
    }
    return success;
  }

  void _startLockoutTimer() {
    // Logic copied from RiverInSituStep3DataCaptureState._startLockoutTimer
    _lockoutTimer?.cancel(); // Cancel any existing timer
    setState(() {
      _isLockedOut = true;
      _lockoutSecondsRemaining = 30; // Reset countdown
    });

    _lockoutTimer = Timer.periodic(const Duration(seconds: 1), (timer) {
      if (_lockoutSecondsRemaining > 0) {
        if (mounted) { // Check if widget is still in the tree
          setState(() { _lockoutSecondsRemaining--; });
        } else {
          timer.cancel(); // Stop timer if widget is disposed
        }
      } else {
        timer.cancel();
        if (mounted) { // Check before final setState
          setState(() { _isLockedOut = false; });
        }
      }
    });
  }

  void _toggleAutoReading(String activeType) {
    // Logic copied from RiverInSituStep3DataCaptureState._toggleAutoReading
    // Uses the correct _samplingService instance
    setState(() {
      _isAutoReading = !_isAutoReading;
      if (_isAutoReading) {
        if (activeType == 'bluetooth') _samplingService.startBluetoothAutoReading();
        else _samplingService.startSerialAutoReading();
        _startLockoutTimer(); // Start countdown when reading starts
      } else {
        if (activeType == 'bluetooth') _samplingService.stopBluetoothAutoReading();
        else _samplingService.stopSerialAutoReading();
        _lockoutTimer?.cancel(); // Stop countdown if reading is stopped manually
        _isLockedOut = false; // Ensure unlocked if stopped manually
      }
    });
  }

  void _disconnect(String type) {
    // Logic copied from RiverInSituStep3DataCaptureState._disconnect
    // Uses the correct _samplingService instance
    // --- START MODIFICATION ---
    final service = _samplingService; // NEW: Use the member variable
    // --- END MODIFICATION ---
    if (type == 'bluetooth') {
      service.disconnectFromBluetooth();
    } else {
      service.disconnectFromSerial();
    }
    _dataSubscription?.cancel();
    _dataSubscription = null;
    _lockoutTimer?.cancel(); // Cancel timer on disconnect
    if (mounted) {
      setState(() {
        _isAutoReading = false;
        _isLockedOut = false; // Reset lockout state
        _isLoading = false;   // --- NEW: Also reset the loading flag ---
      });
    }
  }

  void _disconnectFromAll() {
    // Logic copied from RiverInSituStep3DataCaptureState._disconnectFromAll
    // --- START MODIFICATION ---
    final service = _samplingService; // NEW: Use the member variable
    // --- END MODIFICATION ---
    if (service.bluetoothConnectionState.value != BluetoothConnectionState.disconnected) {
      _disconnect('bluetooth');
    }
    if (service.serialConnectionState.value != SerialConnectionState.disconnected) {
      _disconnect('serial');
    }
  }

  void _updateTextFields(Map<String, double> readings) {
    // Logic copied from RiverInSituStep3DataCaptureState._updateTextFields
    const defaultValue = -999.0;
    setState(() {
      _oxyConcController.text = (readings['Optical Dissolved Oxygen: Compensated mg/L'] ?? defaultValue).toStringAsFixed(5);
      _oxySatController.text = (readings['Optical Dissolved Oxygen: Compensated % Saturation'] ?? defaultValue).toStringAsFixed(5);
      _phController.text = (readings['PH: PH units'] ?? defaultValue).toStringAsFixed(5);
      _tempController.text = (readings['External Temp: Degrees Celcius'] ?? defaultValue).toStringAsFixed(5);
      _ecController.text = (readings['Conductivity: us/cm'] ?? defaultValue).toStringAsFixed(5);
      _salinityController.text = (readings['Conductivity: Salinity'] ?? defaultValue).toStringAsFixed(5);
      _tdsController.text = (readings['Conductivity:TDS mg/L'] ?? defaultValue).toStringAsFixed(5);
      _turbidityController.text = (readings['Turbidity: FNU'] ?? defaultValue).toStringAsFixed(5);
      _batteryController.text = (readings['Sonde: Battery Voltage'] ?? defaultValue).toStringAsFixed(5);
      // Handle Ammonia if the key exists in the readings map
      _ammoniaController.text = (readings['Ammonium (NH4+) mg/L'] ?? defaultValue).toStringAsFixed(5);
    });
  }

  void _validateAndProceed() async {
    // Logic copied from RiverInSituStep3DataCaptureState._validateAndProceed
    if (_isLockedOut) {
      _showSnackBar("Please wait for the initial reading period to complete.", isError: true);
      return;
    }

    if (_isAutoReading) {
      _showStopReadingDialog();
      return;
    }

    if (!_formKey.currentState!.validate()) {
      return;
    }
    _formKey.currentState!.save(); // Save manual inputs like Sonde ID

    final currentReadings = _captureReadingsToMap();

    // Load the generic river parameter limits (same as In-Situ)
    final List<Map<String, dynamic>> riverLimits = await _dbHelper.loadRiverParameterLimits() ?? [];

    final outOfBoundsParams = _validateParameters(currentReadings, riverLimits);

    setState(() {
      _outOfBoundsKeys = outOfBoundsParams.map((p) =>
      _parameters.firstWhere((param) => param['label'] == p['label'])['key'] as String
      ).toSet();
    });

    if (outOfBoundsParams.isNotEmpty) {
      _showParameterLimitDialog(outOfBoundsParams, currentReadings);
    } else {
      _saveDataAndMoveOn(currentReadings);
    }
  }

  Map<String, double> _captureReadingsToMap() {
    // Logic copied from RiverInSituStep3DataCaptureState._captureReadingsToMap
    final Map<String, double> readings = {};
    for (var param in _parameters) {
      final key = param['key'] as String;
      final controller = param['controller'] as TextEditingController;
      // Use -999.0 as the default if parsing fails or text is empty/invalid
      readings[key] = double.tryParse(controller.text) ?? -999.0;
    }
    return readings;
  }

  List<Map<String, dynamic>> _validateParameters(Map<String, double> readings, List<Map<String, dynamic>> limits) {
    // Logic copied from RiverInSituStep3DataCaptureState._validateParameters
    final List<Map<String, dynamic>> invalidParams = [];

    double? parseLimitValue(dynamic value) {
      if (value == null) return null;
      if (value is num) return value.toDouble();
      if (value is String) return double.tryParse(value);
      return null;
    }

    readings.forEach((key, value) {
      if (value == -999.0) return; // Skip validation for missing/default values

      final limitName = _parameterKeyToLimitName[key];
      if (limitName == null) return; // Skip if no mapping exists

      // Find the limits for this parameter
      final limitData = limits.firstWhere(
            (l) => l['param_parameter_list'] == limitName,
        orElse: () => <String, dynamic>{}, // Return empty map if not found
      );

      if (limitData.isNotEmpty) {
        final lowerLimit = parseLimitValue(limitData['param_lower_limit']);
        final upperLimit = parseLimitValue(limitData['param_upper_limit']);

        // Check if the value is outside the defined range (inclusive check is usually fine)
        bool isOutOfBounds = false;
        if (lowerLimit != null && value < lowerLimit) {
          isOutOfBounds = true;
        }
        if (upperLimit != null && value > upperLimit) {
          isOutOfBounds = true;
        }

        if (isOutOfBounds) {
          final paramInfo = _parameters.firstWhere((p) => p['key'] == key, orElse: () => {});
          invalidParams.add({
            'label': paramInfo['label'] ?? key,
            'value': value,
            'lower_limit': lowerLimit,
            'upper_limit': upperLimit,
          });
        }
      }
    });
    return invalidParams;
  }

  void _saveDataAndMoveOn(Map<String, double> readings) {
    // Logic copied from RiverInSituStep3DataCaptureState._saveDataAndMoveOn
    // Saves data to the RiverInvesManualSamplingData model
    try {
      const defaultValue = -999.0;
      widget.data.sondeId = _sondeIdController.text; // Make sure sonde ID is saved
      widget.data.temperature = readings['temperature'] ?? defaultValue;
      widget.data.ph = readings['ph'] ?? defaultValue;
      widget.data.salinity = readings['salinity'] ?? defaultValue;
      widget.data.electricalConductivity = readings['electricalConductivity'] ?? defaultValue;
      widget.data.oxygenConcentration = readings['oxygenConcentration'] ?? defaultValue;
      widget.data.oxygenSaturation = readings['oxygenSaturation'] ?? defaultValue;
      widget.data.tds = readings['tds'] ?? defaultValue;
      widget.data.turbidity = readings['turbidity'] ?? defaultValue;
      widget.data.ammonia = readings['ammonia'] ?? defaultValue;
      widget.data.batteryVoltage = readings['batteryVoltage'] ?? defaultValue;

      // Save flowrate data
      widget.data.flowrateMethod = _selectedFlowrateMethod;
      if (_selectedFlowrateMethod == 'Surface Drifter') {
        widget.data.flowrateSurfaceDrifterHeight = double.tryParse(_sdHeightController.text);
        widget.data.flowrateSurfaceDrifterDistance = double.tryParse(_sdDistanceController.text);
        widget.data.flowrateSurfaceDrifterTimeFirst = _sdTimeFirstController.text;
        widget.data.flowrateSurfaceDrifterTimeLast = _sdTimeLastController.text;
        widget.data.flowrateValue = double.tryParse(_flowrateValueController.text);
      } else if (_selectedFlowrateMethod == 'Flowmeter') {
        widget.data.flowrateSurfaceDrifterHeight = null;
        widget.data.flowrateSurfaceDrifterDistance = null;
        widget.data.flowrateSurfaceDrifterTimeFirst = null;
        widget.data.flowrateSurfaceDrifterTimeLast = null;
        widget.data.flowrateValue = double.tryParse(_flowrateValueController.text);
      } else { // NA
        widget.data.flowrateSurfaceDrifterHeight = null;
        widget.data.flowrateSurfaceDrifterDistance = null;
        widget.data.flowrateSurfaceDrifterTimeFirst = null;
        widget.data.flowrateSurfaceDrifterTimeLast = null;
        widget.data.flowrateValue = null; // Store null for NA
        _flowrateValueController.text = 'NA'; // Display NA
      }

      // Set data capture date/time right before moving on
      final now = DateTime.now();
      widget.data.dataCaptureDate = DateFormat('yyyy-MM-dd').format(now);
      widget.data.dataCaptureTime = DateFormat('HH:mm:ss').format(now);
      _dateController.text = widget.data.dataCaptureDate!;
      _timeController.text = widget.data.dataCaptureTime!;


    } catch (e) {
      _showSnackBar("Could not save parameters due to a data format error: $e", isError: true);
      return;
    }

    // Clear comparison state if moving on
    setState(() {
      _outOfBoundsKeys.clear();
      if (_previousReadingsForComparison != null) {
        _previousReadingsForComparison = null;
      }
    });

    widget.onNext(); // Proceed to the next step
  }

  void _showSnackBar(String message, {bool isError = false}) {
    // Logic copied from RiverInSituStep3DataCaptureState._showSnackBar
    if (mounted) {
      ScaffoldMessenger.of(context).showSnackBar(SnackBar(
        content: Text(message),
        backgroundColor: isError ? Colors.red : null,
      ));
    }
  }

  void _showStopReadingDialog() {
    // Logic copied from RiverInSituStep3DataCaptureState._showStopReadingDialog
    showDialog(
        context: context,
        builder: (BuildContext context) {
          return AlertDialog(
              title: const Text('Data Collection Active'),
              content: const Text('Please stop the live data collection before proceeding.'),
              actions: <Widget>[
                TextButton(child: const Text('OK'), onPressed: () => Navigator.of(context).pop())
              ]
          );
        }
    );
  }

  Map<String, dynamic>? _getActiveConnectionDetails() {
    // Logic copied from RiverInSituStep3DataCaptureState._getActiveConnectionDetails
    // Uses the correct _samplingService instance via context.watch
    final service = context.watch<RiverInvestigativeSamplingService>(); // Watch Investigative service
    if (service.bluetoothConnectionState.value != BluetoothConnectionState.disconnected) {
      return {'type': 'bluetooth', 'state': service.bluetoothConnectionState.value, 'name': service.connectedBluetoothDeviceName};
    }
    if (service.serialConnectionState.value != SerialConnectionState.disconnected) {
      return {'type': 'serial', 'state': service.serialConnectionState.value, 'name': service.connectedSerialDeviceName};
    }
    return null;
  }

  // --- BUILD METHOD and child build methods ---
  // All build methods (_buildParameterListItem, _buildConnectionCard, _buildComparisonView,
  // _showParameterLimitDialog, _buildFlowrateSection, _buildFlowrateRadioButton,
  // _buildSurfaceDrifterFields, _buildFlowmeterField, _buildNAField)
  // are copied directly from river_in_situ_step_3_data_capture.dart.
  // Ensure context.watch<RiverInvestigativeSamplingService>() is used in build.

  @override
  Widget build(BuildContext context) {
    // Watch the Investigative service for state changes
    final service = context.watch<RiverInvestigativeSamplingService>();
    final activeConnection = _getActiveConnectionDetails();
    final String? activeType = activeConnection?['type'] as String?;

    return WillPopScope(
      onWillPop: () async {
        if (_isLockedOut) {
          _showSnackBar("Please wait for the initial reading period to complete.", isError: true);
          return false; // Prevent back navigation
        }
        // Disconnect if navigating back while connected
        _disconnectFromAll();
        return true; // Allow back navigation
      },
      child: Form(
        key: _formKey,
        child: ListView(
          padding: const EdgeInsets.all(24.0),
          children: [
            Text("Investigative Data Capture", style: Theme.of(context).textTheme.headlineSmall), // Updated Title
            const SizedBox(height: 16),
            // Connection Buttons (Bluetooth/Serial)
            Row(
              children: [
                Expanded(
                  child: activeType == 'bluetooth'
                      ? FilledButton.icon(icon: const Icon(Icons.bluetooth_connected), label: const Text("Bluetooth"), onPressed: _isLoading ? null : () => _handleConnectionAttempt('bluetooth'))
                      : OutlinedButton.icon(icon: const Icon(Icons.bluetooth), label: const Text("Bluetooth"), onPressed: _isLoading ? null : () => _handleConnectionAttempt('bluetooth')),
                ),
                const SizedBox(width: 16),
                Expanded(
                  child: activeType == 'serial'
                      ? FilledButton.icon(icon: const Icon(Icons.usb), label: const Text("USB Serial"), onPressed: _isLoading ? null : () => _handleConnectionAttempt('serial'))
                      : OutlinedButton.icon(icon: const Icon(Icons.usb), label: const Text("USB Serial"), onPressed: _isLoading ? null : () => _handleConnectionAttempt('serial')),
                ),
              ],
            ),
            const SizedBox(height: 16),

            // Connection Status Card
            if (activeConnection != null)
              _buildConnectionCard(type: activeConnection['type'], connectionState: activeConnection['state'], deviceName: activeConnection['name']),
            const SizedBox(height: 24),

            // Sonde ID - Updates based on service ValueNotifier
            ValueListenableBuilder<String?>(
              valueListenable: service.sondeId, // Listen to the correct service instance
              builder: (context, sondeId, child) {
                final newSondeId = sondeId ?? '';
                // Use addPostFrameCallback to avoid setting state during build
                WidgetsBinding.instance.addPostFrameCallback((_) {
                  if (mounted && _sondeIdController.text != newSondeId) {
                    _sondeIdController.text = newSondeId;
                    widget.data.sondeId = newSondeId; // Update model
                  }
                });
                return TextFormField(
                  controller: _sondeIdController,
                  decoration: const InputDecoration(labelText: 'Sonde ID *', hintText: 'Connect device or enter manually'),
                  validator: (v) => v == null || v.isEmpty ? 'Sonde ID is required' : null,
                  onChanged: (value) { widget.data.sondeId = value; }, // Update model on change
                  onSaved: (v) => widget.data.sondeId = v,
                );
              },
            ),
            const SizedBox(height: 16),

            // Date & Time (Read-only, set during save)
            Row(
              children: [
                Expanded(child: TextFormField(controller: _dateController, readOnly: true, decoration: const InputDecoration(labelText: 'Capture Date'))),
                const SizedBox(width: 16),
                Expanded(child: TextFormField(controller: _timeController, readOnly: true, decoration: const InputDecoration(labelText: 'Capture Time'))),
              ],
            ),

            // Resample Comparison View (Conditional)
            if (_previousReadingsForComparison != null)
              _buildComparisonView(),

            const Divider(height: 32),

            // Parameter List
            Column(
              children: _parameters.map((param) {
                return _buildParameterListItem(
                  icon: param['icon'] as IconData,
                  label: param['label'] as String,
                  unit: param['unit'] as String,
                  controller: param['controller'] as TextEditingController,
                  isOutOfBounds: _outOfBoundsKeys.contains(param['key']),
                );
              }).toList(),
            ),
            const Divider(height: 32),

            // Flowrate Section
            _buildFlowrateSection(),
            const SizedBox(height: 32),

            // Next Button with Lockout Timer
            ElevatedButton(
              onPressed: _isLockedOut ? null : _validateAndProceed,
              style: ElevatedButton.styleFrom(padding: const EdgeInsets.symmetric(vertical: 16)),
              child: Text(_isLockedOut ? 'Next ($_lockoutSecondsRemaining\s)' : 'Next'),
            ),
          ],
        ),
      ),
    );
  }

  // --- All helper build methods (_buildParameterListItem, _buildConnectionCard, etc.) ---
  // --- are copied directly from RiverInSituStep3DataCaptureState ---

  Widget _buildParameterListItem({ required IconData icon, required String label, required String unit, required TextEditingController controller, bool isOutOfBounds = false}) {
    // Copied from RiverInSituStep3DataCaptureState._buildParameterListItem
    final bool isMissing = controller.text.isEmpty || controller.text.contains('-999');
    // Display value formatted nicely, use '-.--' for missing/default
    final String displayValue = isMissing ? '-.--' : (double.tryParse(controller.text) ?? -999.0).toStringAsFixed(5);
    final String displayLabel = unit.isEmpty ? label : '$label ($unit)';

    // Determine color based on limits and missing status
    final Color valueColor = isOutOfBounds
        ? Colors.red // Highlight out of bounds in red
        : (isMissing ? Colors.grey : Theme.of(context).textTheme.bodyLarge?.color ?? Colors.black); // Grey for missing, default otherwise

    return Card(
      margin: const EdgeInsets.symmetric(vertical: 4.0),
      child: ListTile(
        leading: Icon(icon, color: Theme.of(context).primaryColor, size: 32),
        title: Text(displayLabel),
        trailing: Text(
          displayValue,
          style: Theme.of(context).textTheme.titleMedium?.copyWith(
              fontWeight: FontWeight.bold,
              color: valueColor
          ),
        ),
      ),
    );
  }

  Widget _buildConnectionCard({required String type, required dynamic connectionState, String? deviceName}) {
    // Copied from RiverInSituStep3DataCaptureState._buildConnectionCard
    final bool isConnected = connectionState == BluetoothConnectionState.connected || connectionState == SerialConnectionState.connected;
    final bool isConnecting = connectionState == BluetoothConnectionState.connecting || connectionState == SerialConnectionState.connecting;

    Color statusColor;
    String statusText;

    if (isConnected) {
      statusColor = Colors.green;
      statusText = 'Connected to ${deviceName ?? 'device'} ($type)';
    } else if (isConnecting) {
      statusColor = Colors.orange;
      statusText = 'Connecting via $type...';
    } else {
      statusColor = Colors.red;
      statusText = 'Disconnected';
    }

    return Card(
      elevation: 2,
      child: Padding(
        padding: const EdgeInsets.all(16.0),
        child: Column(
          children: [
            Text(statusText, style: TextStyle(color: statusColor, fontWeight: FontWeight.bold, fontSize: 16)),
            const SizedBox(height: 16),
            if (isConnecting || _isLoading)
              const CircularProgressIndicator()
            else if (isConnected)
              Row(
                mainAxisAlignment: MainAxisAlignment.spaceEvenly,
                children: [
                  ElevatedButton.icon(
                    icon: Icon(_isAutoReading ? Icons.stop_circle_outlined : Icons.play_circle_outlined),
                    label: Text(_isAutoReading
                        ? (_isLockedOut ? 'Stop Reading ($_lockoutSecondsRemaining\s)' : 'Stop Reading')
                        : 'Start Reading'),
                    onPressed: (_isAutoReading && _isLockedOut) ? null : () => _toggleAutoReading(type), // Pass active type
                    style: ElevatedButton.styleFrom(
                      backgroundColor: _isAutoReading
                          ? (_isLockedOut ? Colors.grey.shade600 : Colors.orange) // Grey out if locked
                          : Colors.green,
                      foregroundColor: Colors.white,
                    ),
                  ),
                  TextButton.icon(
                    icon: const Icon(Icons.link_off),
                    label: const Text('Disconnect'),
                    onPressed: () => _disconnect(type), // Pass active type
                    style: TextButton.styleFrom(foregroundColor: Colors.red),
                  )
                ],
              )
            // Optionally add a button to reconnect if disconnected
            else
              ElevatedButton.icon(
                icon: Icon(type == 'bluetooth' ? Icons.bluetooth_searching : Icons.usb),
                label: Text('Reconnect $type'),
                onPressed: _isLoading ? null : () => _handleConnectionAttempt(type),
              )
          ],
        ),
      ),
    );
  }

  Widget _buildComparisonView() {
    // Copied from RiverInSituStep3DataCaptureState._buildComparisonView
    final previousReadings = _previousReadingsForComparison!;
    final isDarkTheme = Theme.of(context).brightness == Brightness.dark;

    return Card(
      margin: const EdgeInsets.only(top: 24.0),
      color: Theme.of(context).cardColor,
      child: Padding(
        padding: const EdgeInsets.all(16.0),
        child: DefaultTextStyle(
          style: TextStyle(color: Theme.of(context).textTheme.bodyLarge?.color),
          child: Column(
            crossAxisAlignment: CrossAxisAlignment.start,
            children: [
              Text(
                "Resample Comparison",
                style: Theme.of(context).textTheme.titleLarge?.copyWith(color: Theme.of(context).primaryColor),
              ),
              const SizedBox(height: 8),
              Table(
                columnWidths: const {
                  0: FlexColumnWidth(2),
                  1: FlexColumnWidth(1.5),
                  2: FlexColumnWidth(1.5),
                },
                border: TableBorder(
                  horizontalInside: BorderSide(width: 1, color: isDarkTheme ? Colors.grey.shade700 : Colors.grey.shade300, style: BorderStyle.solid),
                  verticalInside: BorderSide(width: 1, color: isDarkTheme ? Colors.grey.shade700 : Colors.grey.shade300, style: BorderStyle.solid),
                  top: BorderSide(width: 1.5, color: isDarkTheme ? Colors.grey.shade600 : Colors.grey.shade400),
                  bottom: BorderSide(width: 1.5, color: isDarkTheme ? Colors.grey.shade600 : Colors.grey.shade400),
                ),
                children: [
                  TableRow(
                    decoration: BoxDecoration(color: isDarkTheme ? Colors.grey.shade800 : Colors.grey.shade200),
                    children: [
                      Padding(padding: const EdgeInsets.all(8.0), child: Text('Parameter', style: TextStyle(fontWeight: FontWeight.bold, color: Theme.of(context).textTheme.titleMedium?.color))),
                      Padding(padding: const EdgeInsets.all(8.0), child: Text('Previous', style: TextStyle(fontWeight: FontWeight.bold, color: Theme.of(context).textTheme.titleMedium?.color))),
                      Padding(padding: const EdgeInsets.all(8.0), child: Text('Current', style: TextStyle(fontWeight: FontWeight.bold, color: Theme.of(context).textTheme.titleMedium?.color))),
                    ],
                  ),
                  ..._parameters.map((param) {
                    final key = param['key'] as String;
                    final label = param['label'] as String;
                    final controller = param['controller'] as TextEditingController;
                    final previousValue = previousReadings[key];
                    final bool isCurrentValueOutOfBounds = _outOfBoundsKeys.contains(key);
                    final currentValue = double.tryParse(controller.text) ?? -999.0;


                    return TableRow(
                      children: [
                        Padding(padding: const EdgeInsets.all(8.0), child: Text(label)),
                        Padding(
                          padding: const EdgeInsets.all(8.0),
                          child: Text(
                            previousValue == -999.0 ? '-.--' : previousValue!.toStringAsFixed(5),
                            style: TextStyle(color: isDarkTheme ? Colors.orange.shade200 : Colors.orange.shade700), // Previous always orange-ish
                          ),
                        ),
                        Padding(
                          padding: const EdgeInsets.all(8.0),
                          child: Text(
                            currentValue == -999.0 ? '-.--' : currentValue.toStringAsFixed(5),
                            style: TextStyle(
                                color: isCurrentValueOutOfBounds
                                    ? Colors.red // Current out of bounds = red
                                    : (isDarkTheme ? Colors.green.shade200 : Colors.green.shade700), // Current in bounds = green
                                fontWeight: FontWeight.bold
                            ),
                          ),
                        ),
                      ],
                    );
                  }).toList(),
                ],
              ),
            ],
          ),
        ),
      ),
    );
  }

  Future<void> _showParameterLimitDialog(List<Map<String, dynamic>> invalidParams, Map<String, double> readings) async {
    // Copied from RiverInSituStep3DataCaptureState._showParameterLimitDialog
    return showDialog(
      context: context,
      barrierDismissible: false, // User must choose an action
      builder: (BuildContext context) {
        final isDarkTheme = Theme.of(context).brightness == Brightness.dark;
        return AlertDialog(
          title: const Text('Parameter Limit Warning'),
          content: SingleChildScrollView(
            child: DefaultTextStyle(
              style: TextStyle(color: Theme.of(context).textTheme.bodyMedium?.color),
              child: Column(
                mainAxisSize: MainAxisSize.min,
                crossAxisAlignment: CrossAxisAlignment.start,
                children: [
                  const Text('The following parameters are outside the standard limits:'),
                  const SizedBox(height: 16),
                  Table(
                    columnWidths: const {
                      0: FlexColumnWidth(2),    // Parameter name
                      1: FlexColumnWidth(2.5),  // Limit range
                      2: FlexColumnWidth(1.5),  // Current value
                    },
                    border: TableBorder(
                      horizontalInside: BorderSide(width: 0.5, color: isDarkTheme ? Colors.grey.shade700 : Colors.grey.shade300),
                      verticalInside: BorderSide(width: 0.5, color: isDarkTheme ? Colors.grey.shade700 : Colors.grey.shade300),
                      top: BorderSide(width: 1, color: isDarkTheme ? Colors.grey.shade600 : Colors.grey.shade400),
                      bottom: BorderSide(width: 1, color: isDarkTheme ? Colors.grey.shade600 : Colors.grey.shade400),
                    ),
                    children: [
                      TableRow(
                        decoration: BoxDecoration(color: isDarkTheme ? Colors.grey.shade800 : Colors.grey.shade200),
                        children: [
                          Padding(padding: const EdgeInsets.all(6.0), child: Text('Parameter', style: TextStyle(fontWeight: FontWeight.bold, color: Theme.of(context).textTheme.titleSmall?.color))),
                          Padding(padding: const EdgeInsets.all(6.0), child: Text('Limit Range', style: TextStyle(fontWeight: FontWeight.bold, color: Theme.of(context).textTheme.titleSmall?.color))),
                          Padding(padding: const EdgeInsets.all(6.0), child: Text('Current', style: TextStyle(fontWeight: FontWeight.bold, color: Theme.of(context).textTheme.titleSmall?.color))),
                        ],
                      ),
                      ...invalidParams.map((p) => TableRow(
                        children: [
                          Padding(padding: const EdgeInsets.all(6.0), child: Text(p['label'])),
                          // Display limits nicely, handling nulls
                          Padding(padding: const EdgeInsets.all(6.0), child: Text('${p['lower_limit']?.toStringAsFixed(5) ?? 'N/A'} - ${p['upper_limit']?.toStringAsFixed(5) ?? 'N/A'}')),
                          Padding(
                            padding: const EdgeInsets.all(6.0),
                            child: Text(
                              p['value'].toStringAsFixed(5),
                              style: const TextStyle(color: Colors.redAccent, fontWeight: FontWeight.bold),
                            ),
                          ),
                        ],
                      )).toList(),
                    ],
                  ),
                  const SizedBox(height: 16),
                  const Text('Do you want to resample or proceed with the current values? Please verify with standard solution.'),
                ],
              ),
            ),
          ),
          actions: <Widget>[
            TextButton(
              child: const Text('Resample'),
              onPressed: () {
                setState(() {
                  // Store the current (out of bounds) readings for comparison view
                  _previousReadingsForComparison = readings;
                });
                Navigator.of(context).pop(); // Close the dialog, user will retake readings
              },
            ),
            FilledButton(
              child: const Text('Proceed Anyway'),
              onPressed: () {
                Navigator.of(context).pop(); // Close the dialog
                _saveDataAndMoveOn(readings); // Save current readings and move to next step
              },
            ),
          ],
        );
      },
    );
  }

  Widget _buildFlowrateSection() {
    // Copied from RiverInSituStep3DataCaptureState._buildFlowrateSection
    return Card(
      margin: const EdgeInsets.symmetric(vertical: 4.0),
      child: Padding(
        padding: const EdgeInsets.all(16.0),
        child: Column(
          crossAxisAlignment: CrossAxisAlignment.start,
          children: [
            Text("Flowrate", style: Theme.of(context).textTheme.titleLarge),
            const SizedBox(height: 8),
            // Radio buttons for method selection
            Row(
              mainAxisAlignment: MainAxisAlignment.spaceAround,
              children: [
                _buildFlowrateRadioButton("Surface Drifter"),
                _buildFlowrateRadioButton("Flowmeter"),
                _buildFlowrateRadioButton("NA"), // Not Applicable
              ],
            ),
            // Conditional fields based on selected method
            if (_selectedFlowrateMethod == 'Surface Drifter')
              _buildSurfaceDrifterFields(),
            if (_selectedFlowrateMethod == 'Flowmeter')
              _buildFlowmeterField(),
            if (_selectedFlowrateMethod == 'NA')
              _buildNAField(),
          ],
        ),
      ),
    );
  }

  Widget _buildFlowrateRadioButton(String title) {
    // Copied from RiverInSituStep3DataCaptureState._buildFlowrateRadioButton
    return Column(
      children: [
        Radio<String>(
          value: title,
          groupValue: _selectedFlowrateMethod,
          onChanged: _onFlowrateMethodChanged,
        ),
        Text(title),
      ],
    );
  }

  Widget _buildSurfaceDrifterFields() {
    // Copied from RiverInSituStep3DataCaptureState._buildSurfaceDrifterFields
    return Padding(
      padding: const EdgeInsets.only(top: 16.0),
      child: Column(
        children: [
          TextFormField(
            controller: _sdHeightController,
            decoration: const InputDecoration(labelText: 'Height (m)'),
            keyboardType: const TextInputType.numberWithOptions(decimal: true),
            // Add validation if needed
          ),
          const SizedBox(height: 16),
          TextFormField(
            controller: _sdDistanceController,
            decoration: const InputDecoration(labelText: 'Distance (m) *'),
            keyboardType: const TextInputType.numberWithOptions(decimal: true),
            validator: (v) => v == null || v.isEmpty ? 'Distance is required' : null,
          ),
          const SizedBox(height: 16),
          TextFormField(
            controller: _sdTimeFirstController,
            decoration: const InputDecoration(labelText: 'Time First Deploy (HH:mm:ss) *', suffixIcon: Icon(Icons.timer)),
            readOnly: true,
            onTap: () => _selectTime(context, _sdTimeFirstController),
            validator: (v) => v == null || v.isEmpty ? 'Start time is required' : null,
          ),
          const SizedBox(height: 16),
          TextFormField(
            controller: _sdTimeLastController,
            decoration: const InputDecoration(labelText: 'Time Last Deploy (HH:mm:ss) *', suffixIcon: Icon(Icons.timer)),
            readOnly: true,
            onTap: () => _selectTime(context, _sdTimeLastController),
            validator: (v) => v == null || v.isEmpty ? 'End time is required' : null,
          ),
          const SizedBox(height: 16),
          ElevatedButton(
            onPressed: _calculateFlowrate,
            child: const Text('Calculate Flowrate'),
          ),
          const SizedBox(height: 16),
          TextFormField(
            controller: _flowrateValueController,
            decoration: const InputDecoration(labelText: 'Calculated Flowrate (m/s)'),
            readOnly: true,
            // Add validator if calculation must be done?
          ),
        ],
      ),
    );
  }

  Widget _buildFlowmeterField() {
    // Copied from RiverInSituStep3DataCaptureState._buildFlowmeterField
    return Padding(
      padding: const EdgeInsets.only(top: 16.0),
      child: TextFormField(
        controller: _flowrateValueController,
        decoration: const InputDecoration(labelText: 'Flowrate (m/s) *'),
        keyboardType: const TextInputType.numberWithOptions(decimal: true),
        validator: (v) => v == null || v.isEmpty ? 'Flowrate value is required' : null,
      ),
    );
  }

  Widget _buildNAField() {
    // Copied from RiverInSituStep3DataCaptureState._buildNAField
    return Padding(
      padding: const EdgeInsets.only(top: 16.0),
      child: TextFormField(
        controller: _flowrateValueController,
        decoration: const InputDecoration(labelText: 'Flowrate (m/s)'),
        initialValue: 'NA', // Set initial value to NA
        readOnly: true, // Make it read-only
      ),
    );
  }

} // End of State class