app/lib/screens/baseline_screen.dart

import 'package:flutter/material.dart';
import '../models/models.dart';
import '../risk_engine/mortality_tables.dart';
import '../storage/local_storage.dart';
import '../theme.dart';
import 'about_screen.dart';
import 'behavioral_screen.dart';

class BaselineScreen extends StatefulWidget {
  final bool readOnly;
  final UserProfile? initialProfile;

  const BaselineScreen({
    super.key,
    this.readOnly = false,
    this.initialProfile,
  });

  @override
  State<BaselineScreen> createState() => _BaselineScreenState();
}

class _BaselineScreenState extends State<BaselineScreen> {
  int _age = 35;
  Sex _sex = Sex.male;
  String _country = 'United States';
  double _heightCm = 170;
  double _weightKg = 70;
  final Set<Diagnosis> _diagnoses = {};
  bool _useMetric = false;

  late List<String> _countries;

  @override
  void initState() {
    super.initState();
    _countries = getSupportedCountries();
    _loadInitialData();
  }

  Future<void> _loadInitialData() async {
    // Load unit preference
    final useMetric = await LocalStorage.getUseMetricUnits();
    setState(() => _useMetric = useMetric);

    // If initial profile provided, use that
    if (widget.initialProfile != null) {
      _applyProfile(widget.initialProfile!);
      return;
    }

    // Otherwise load from storage
    final profile = await LocalStorage.getProfile();
    if (profile != null) {
      _applyProfile(profile);
    }
  }

  void _applyProfile(UserProfile profile) {
    setState(() {
      _age = profile.age;
      _sex = profile.sex;
      _country = profile.country;
      _heightCm = profile.heightCm;
      _weightKg = profile.weightKg;
      _diagnoses.clear();
      _diagnoses.addAll(profile.diagnoses);
    });
  }

  double get _bmi => _weightKg / ((_heightCm / 100) * (_heightCm / 100));

  String get _bmiCategory {
    if (_bmi < 18.5) return 'Underweight';
    if (_bmi < 25) return 'Normal';
    if (_bmi < 30) return 'Overweight';
    if (_bmi < 35) return 'Obese I';
    if (_bmi < 40) return 'Obese II';
    return 'Obese III';
  }

  @override
  Widget build(BuildContext context) {
    return Scaffold(
      appBar: AppBar(
        title: Text(widget.readOnly ? 'Baseline (View Only)' : 'Baseline'),
        leading: IconButton(
          icon: const Icon(Icons.arrow_back),
          onPressed: () => Navigator.pop(context),
        ),
        actions: [
          IconButton(
            icon: Icon(_useMetric ? Icons.straighten : Icons.square_foot),
            tooltip: _useMetric ? 'Using Metric' : 'Using Imperial',
            onPressed: widget.readOnly ? null : _toggleUnits,
          ),
          IconButton(
            icon: const Icon(Icons.info_outline),
            onPressed: () => Navigator.push(
              context,
              MaterialPageRoute(builder: (_) => const AboutScreen()),
            ),
          ),
        ],
      ),
      body: SingleChildScrollView(
        padding: const EdgeInsets.all(24),
        child: Column(
          crossAxisAlignment: CrossAxisAlignment.start,
          children: [
            Text(
              'Demographics',
              style: Theme.of(context).textTheme.headlineSmall,
            ),
            const SizedBox(height: 8),
            Text(
              'This information establishes your baseline life expectancy.',
              style: Theme.of(context).textTheme.bodyMedium,
            ),
            const SizedBox(height: 24),

            // Age
            _buildSectionLabel('Age'),
            const SizedBox(height: 8),
            _buildAgeSelector(),
            const SizedBox(height: 24),

            // Sex
            _buildSectionLabel('Biological Sex'),
            const SizedBox(height: 8),
            _buildSexSelector(),
            const SizedBox(height: 24),

            // Country
            _buildSectionLabel('Country'),
            const SizedBox(height: 8),
            _buildCountryDropdown(),
            const SizedBox(height: 24),

            // Height
            _buildSectionLabel('Height'),
            const SizedBox(height: 8),
            _buildHeightSlider(),
            const SizedBox(height: 24),

            // Weight
            _buildSectionLabel('Weight'),
            const SizedBox(height: 8),
            _buildWeightSlider(),
            const SizedBox(height: 16),

            // BMI display
            _buildBmiDisplay(),
            const SizedBox(height: 32),

            // Existing conditions
            Text(
              'Existing Conditions',
              style: Theme.of(context).textTheme.headlineSmall,
            ),
            const SizedBox(height: 8),
            Text(
              'Select any diagnosed conditions. These affect baseline calculations but are not modifiable factors.',
              style: Theme.of(context).textTheme.bodyMedium,
            ),
            const SizedBox(height: 16),
            _buildDiagnosisCheckboxes(),
            const SizedBox(height: 32),

            // Continue button (hidden in readOnly mode)
            if (!widget.readOnly)
              SizedBox(
                width: double.infinity,
                child: ElevatedButton(
                  onPressed: _continue,
                  child: const Text('Continue'),
                ),
              ),
            const SizedBox(height: 16),
          ],
        ),
      ),
    );
  }

  Widget _buildSectionLabel(String label) {
    return Text(
      label,
      style: Theme.of(context).textTheme.labelLarge,
    );
  }

  Widget _buildAgeSelector() {
    return Container(
      padding: const EdgeInsets.symmetric(horizontal: 16, vertical: 8),
      decoration: BoxDecoration(
        color: AppColors.surfaceVariant,
        borderRadius: BorderRadius.circular(12),
      ),
      child: Row(
        children: [
          IconButton(
            icon: const Icon(Icons.remove),
            onPressed: widget.readOnly || _age <= 18
                ? null
                : () => setState(() => _age--),
          ),
          Expanded(
            child: Text(
              '$_age years',
              textAlign: TextAlign.center,
              style: Theme.of(context).textTheme.headlineMedium,
            ),
          ),
          IconButton(
            icon: const Icon(Icons.add),
            onPressed: widget.readOnly || _age >= 100
                ? null
                : () => setState(() => _age++),
          ),
        ],
      ),
    );
  }

  Widget _buildSexSelector() {
    return Row(
      children: [
        Expanded(
          child: _buildToggleButton(
            'Male',
            _sex == Sex.male,
            widget.readOnly ? null : () => setState(() => _sex = Sex.male),
          ),
        ),
        const SizedBox(width: 12),
        Expanded(
          child: _buildToggleButton(
            'Female',
            _sex == Sex.female,
            widget.readOnly ? null : () => setState(() => _sex = Sex.female),
          ),
        ),
      ],
    );
  }

  Widget _buildToggleButton(String label, bool selected, VoidCallback? onTap) {
    return GestureDetector(
      onTap: onTap,
      child: Container(
        padding: const EdgeInsets.symmetric(vertical: 14),
        decoration: BoxDecoration(
          color: selected ? AppColors.primary : AppColors.surfaceVariant,
          borderRadius: BorderRadius.circular(12),
        ),
        child: Text(
          label,
          textAlign: TextAlign.center,
          style: TextStyle(
            fontSize: 16,
            fontWeight: FontWeight.w600,
            color: selected ? Colors.white : AppColors.textSecondary,
          ),
        ),
      ),
    );
  }

  Widget _buildCountryDropdown() {
    return Container(
      padding: const EdgeInsets.symmetric(horizontal: 16),
      decoration: BoxDecoration(
        color: AppColors.surfaceVariant,
        borderRadius: BorderRadius.circular(12),
      ),
      child: DropdownButtonHideUnderline(
        child: DropdownButton<String>(
          value: _country,
          isExpanded: true,
          icon: const Icon(Icons.keyboard_arrow_down),
          items: _countries.map((country) {
            return DropdownMenuItem(
              value: country,
              child: Text(country),
            );
          }).toList(),
          onChanged: widget.readOnly
              ? null
              : (value) {
                  if (value != null) setState(() => _country = value);
                },
        ),
      ),
    );
  }

  Widget _buildHeightSlider() {
    final primaryText = _useMetric
        ? '${_heightCm.round()} cm'
        : _cmToFeetInches(_heightCm);
    final secondaryText = _useMetric
        ? _cmToFeetInches(_heightCm)
        : '${_heightCm.round()} cm';

    return Column(
      children: [
        Row(
          mainAxisAlignment: MainAxisAlignment.spaceBetween,
          children: [
            Text(
              primaryText,
              style: Theme.of(context).textTheme.headlineMedium,
            ),
            Text(
              secondaryText,
              style: Theme.of(context).textTheme.bodyMedium,
            ),
          ],
        ),
        Slider(
          value: _heightCm,
          min: 120,
          max: 220,
          divisions: 100,
          onChanged: widget.readOnly
              ? null
              : (value) => setState(() => _heightCm = value),
        ),
      ],
    );
  }

  Widget _buildWeightSlider() {
    final lbs = (_weightKg * 2.205).round();
    final primaryText = _useMetric ? '${_weightKg.round()} kg' : '$lbs lbs';
    final secondaryText = _useMetric ? '$lbs lbs' : '${_weightKg.round()} kg';

    return Column(
      children: [
        Row(
          mainAxisAlignment: MainAxisAlignment.spaceBetween,
          children: [
            Text(
              primaryText,
              style: Theme.of(context).textTheme.headlineMedium,
            ),
            Text(
              secondaryText,
              style: Theme.of(context).textTheme.bodyMedium,
            ),
          ],
        ),
        Slider(
          value: _weightKg,
          min: 30,
          max: 200,
          divisions: 170,
          onChanged: widget.readOnly
              ? null
              : (value) => setState(() => _weightKg = value),
        ),
      ],
    );
  }

  Widget _buildBmiDisplay() {
    return Container(
      padding: const EdgeInsets.all(16),
      decoration: BoxDecoration(
        color: AppColors.surfaceVariant,
        borderRadius: BorderRadius.circular(12),
      ),
      child: Row(
        mainAxisAlignment: MainAxisAlignment.spaceBetween,
        children: [
          Column(
            crossAxisAlignment: CrossAxisAlignment.start,
            children: [
              Text(
                'BMI',
                style: Theme.of(context).textTheme.bodyMedium,
              ),
              Text(
                _bmi.toStringAsFixed(1),
                style: Theme.of(context).textTheme.headlineMedium,
              ),
            ],
          ),
          Container(
            padding: const EdgeInsets.symmetric(horizontal: 12, vertical: 6),
            decoration: BoxDecoration(
              color: _getBmiColor(),
              borderRadius: BorderRadius.circular(8),
            ),
            child: Text(
              _bmiCategory,
              style: const TextStyle(
                color: Colors.white,
                fontWeight: FontWeight.w600,
              ),
            ),
          ),
        ],
      ),
    );
  }

  Color _getBmiColor() {
    if (_bmi < 18.5 || _bmi >= 30) return AppColors.warning;
    if (_bmi >= 25) return AppColors.primary;
    return AppColors.success;
  }

  Widget _buildDiagnosisCheckboxes() {
    return Column(
      children: Diagnosis.values.map((diagnosis) {
        return CheckboxListTile(
          value: _diagnoses.contains(diagnosis),
          onChanged: widget.readOnly
              ? null
              : (checked) {
                  setState(() {
                    if (checked == true) {
                      _diagnoses.add(diagnosis);
                    } else {
                      _diagnoses.remove(diagnosis);
                    }
                  });
                },
          title: Text(_getDiagnosisLabel(diagnosis)),
          controlAffinity: ListTileControlAffinity.leading,
          contentPadding: EdgeInsets.zero,
        );
      }).toList(),
    );
  }

  String _getDiagnosisLabel(Diagnosis diagnosis) {
    switch (diagnosis) {
      case Diagnosis.cardiovascular:
        return 'Cardiovascular disease';
      case Diagnosis.diabetes:
        return 'Diabetes';
      case Diagnosis.cancer:
        return 'Cancer (active)';
      case Diagnosis.copd:
        return 'COPD';
      case Diagnosis.hypertension:
        return 'Hypertension';
    }
  }

  String _cmToFeetInches(double cm) {
    final totalInches = cm / 2.54;
    final feet = (totalInches / 12).floor();
    final inches = (totalInches % 12).round();
    return "$feet'$inches\"";
  }

  Future<void> _toggleUnits() async {
    final newValue = !_useMetric;
    await LocalStorage.setUseMetricUnits(newValue);
    setState(() => _useMetric = newValue);
  }

  void _continue() async {
    final profile = UserProfile(
      age: _age,
      sex: _sex,
      country: _country,
      heightCm: _heightCm,
      weightKg: _weightKg,
      diagnoses: _diagnoses,
    );

    await LocalStorage.saveProfile(profile);

    if (mounted) {
      Navigator.of(context).push(
        MaterialPageRoute(
          builder: (_) => BehavioralScreen(profile: profile),
        ),
      );
    }
  }
}