d4t_formulas/lib/ai/formula_editor.dart

import 'package:flutter/material.dart';
import 'package:flutter/services.dart';
import 'package:flutter_markdown_plus_latex/flutter_markdown_plus_latex.dart';
import 'package:flutter_markdown_plus/flutter_markdown_plus.dart';
import 'package:markdown/markdown.dart' as markdown;
import '../formula_models.dart';
import '../corpus.dart';
import '../database/database_service.dart';
import '../service_locator.dart';
import 'formula_screen.dart';
import 'unit_dropdown.dart';
import 'package:flutter_code_editor/flutter_code_editor.dart';
import 'package:flutter_highlight/themes/monokai-sublime.dart';
import 'package:highlight/languages/dart.dart';

/// A screen for editing a Formula's properties including name, description,
/// input/output variables, and d4rt code.
class FormulaEditor extends StatefulWidget {
  final Formula formula;
  final Corpus corpus;
  final Function(Formula)? onSave;

  const FormulaEditor({super.key, required this.formula, required this.corpus, this.onSave});

  @override
  State<FormulaEditor> createState() => _FormulaEditorState();
}

class _FormulaEditorState extends State<FormulaEditor> {
  final _formKey = GlobalKey<FormState>();
  late TextEditingController _nameController;
  late TextEditingController _descriptionController;
  late CodeController _d4rtCodeController;

  // Track input variables
  final List<_InputVariableRowData> _inputVariables = [];

  // Output variable
  late _OutputVariableRowData _outputVariable;

  bool _isPreviewVisible = false;

  @override
  void initState() {
    super.initState();
    _nameController = TextEditingController(text: widget.formula.name);
    _descriptionController = TextEditingController(text: widget.formula.description ?? '');
    _d4rtCodeController = CodeController(language: dart, text: widget.formula.d4rtCode ?? '');

    // Initialize input variables
    for (final input in widget.formula.input) {
      _inputVariables.add(
        _InputVariableRowData(
          nameController: TextEditingController(text: input.name),
          unit: input.unit,
          values: input.values != null ? List.from(input.values!) : null,
        ),
      );
    }

    // Initialize output variable
    _outputVariable = _OutputVariableRowData(
      nameController: TextEditingController(text: widget.formula.output.name),
      unit: widget.formula.output.unit,
    );
  }

  @override
  void dispose() {
    _nameController.dispose();
    _descriptionController.dispose();
    _d4rtCodeController.dispose();
    for (final variable in _inputVariables) {
      variable.nameController.dispose();
    }
    _outputVariable.nameController.dispose();
    super.dispose();
  }

  void _addInputVariable() {
    setState(() {
      _inputVariables.add(
        _InputVariableRowData(
          nameController: TextEditingController(text: 'var${_inputVariables.length + 1}'),
          unit: null,
          values: null,
        ),
      );
    });
  }

  void _removeInputVariable(int index) {
    setState(() {
      _inputVariables.removeAt(index);
    });
  }

  void _showPreview() {
    setState(() {
      _isPreviewVisible = true;
    });
  }

  void _hidePreview() {
    setState(() {
      _isPreviewVisible = false;
    });
  }

  void _testFormula() {
    // Validate the formula before testing
    if (!_validateFormula()) {
      return;
    }

    final formula = _buildFormula();
    if (formula == null) return;

    Navigator.push(
      context,
      MaterialPageRoute(
        builder: (context) => FormulaScreen(formula: formula, corpus: widget.corpus),
      ),
    );
  }

  bool _validateFormula() {
    // Validate name
    if (_nameController.text.trim().isEmpty) {
      _showErrorDialog('Formula name cannot be empty');
      return false;
    }

    // Validate output name
    if (_outputVariable.nameController.text.trim().isEmpty) {
      _showErrorDialog('Output variable name cannot be empty');
      return false;
    }

    // Validate input variable names
    for (final variable in _inputVariables) {
      if (variable.nameController.text.trim().isEmpty) {
        _showErrorDialog('Input variable names cannot be empty');
        return false;
      }
    }

    // Validate d4rt code
    if (_d4rtCodeController.text.trim().isEmpty) {
      _showErrorDialog('D4RT code cannot be empty');
      return false;
    }

    return true;
  }

  Formula? _buildFormula() {
    try {
      final input = <VariableSpec>[];
      for (final variable in _inputVariables) {
        input.add(
          VariableSpec(name: variable.nameController.text.trim(), unit: variable.unit, values: variable.values),
        );
      }

      final output = VariableSpec(name: _outputVariable.nameController.text.trim(), unit: _outputVariable.unit);

      return Formula(
        uuid: widget.formula.uuid,
        name: _nameController.text.trim(),
        description: _descriptionController.text.isEmpty ? null : _descriptionController.text,
        input: input,
        output: output,
        d4rtCode: _d4rtCodeController.fullText,
        tags: widget.formula.tags, // Preserve existing tags
      );
    } catch (e) {
      _showErrorDialog('Error building formula: $e');
      return null;
    }
  }

  Future<void> _saveFormula() async {
    if (!_validateFormula()) {
      return;
    }

    final formula = _buildFormula();
    if (formula == null) return;

    try {
      final database = getDatabase();

      // Update corpus in memory
      widget.corpus.updateFormula(formula);

      // Update database
      final updated = await database.updateFormula(formula);

      if (!updated) {
        // If formula wasn't found (e.g., name changed), add it as new
        await database.addFormula(formula);
      }

      // Call the onSave callback if provided
      widget.onSave?.call(formula);

      // Show success message
      ScaffoldMessenger.of(context).showSnackBar(
        SnackBar(
          content: Text('Formula "${formula.name}" saved successfully!'),
          backgroundColor: Theme.of(context).colorScheme.primary,
        ),
      );
    } catch (e, stack) {
      print('Error saving formula: $e\n$stack');
      _showErrorDialog('Error saving formula: $e');
    }
  }

  Future<void> _saveFormulaAsCopy() async {
    if (!_validateFormula()) {
      return;
    }

    final formula = _buildFormula();
    if (formula == null) return;

    try {
      final database = getDatabase();

      // Create a copy with a new UUID
      final formulaCopy = Formula(
        name: '${formula.name} (Copy)',
        description: formula.description,
        input: formula.input,
        output: formula.output,
        d4rtCode: formula.d4rtCode,
        tags: formula.tags,
      );

      // Add to corpus
      widget.corpus.addFormula(formulaCopy);

      // Add to database
      await database.addFormula(formulaCopy);

      // Call the onSave callback if provided
      widget.onSave?.call(formulaCopy);

      // Show success message
      ScaffoldMessenger.of(context).showSnackBar(
        SnackBar(
          content: Text('Formula "${formulaCopy.name}" saved successfully!'),
          backgroundColor: Theme.of(context).colorScheme.primary,
        ),
      );

      // Navigate back to the formula list with the new formula
      Navigator.pop(context, formulaCopy);
    } catch (e, stack) {
      print('Error saving formula copy: $e\n$stack');
      _showErrorDialog('Error saving formula copy: $e');
    }
  }

  void _showErrorDialog(String message) {
    showDialog(
      context: context,
      builder: (BuildContext context) {
        return AlertDialog(
          title: const Text('Error'),
          content: Text(message),
          actions: <Widget>[
            TextButton(
              onPressed: () {
                Navigator.of(context).pop();
              },
              child: const Text('OK'),
            ),
          ],
        );
      },
    );
  }

  @override
  Widget build(BuildContext context) {
    return Scaffold(
      appBar: AppBar(
        title: const Text('Edit Formula'),
        actions: [
          IconButton(icon: const Icon(Icons.play_arrow), onPressed: _testFormula, tooltip: 'Test Formula'),
          IconButton(icon: const Icon(Icons.copy), onPressed: _saveFormulaAsCopy, tooltip: 'Save as copy'),
          IconButton(icon: const Icon(Icons.save), onPressed: _saveFormula, tooltip: 'Save'),
        ],
      ),
      body: Form(
        key: _formKey,
        child: Padding(
          padding: const EdgeInsets.all(16.0),
          child: ListView(
            children: [
              _buildNameSection(),
              const SizedBox(height: 16),
              _buildDescriptionSection(),
              const SizedBox(height: 16),
              _buildInputVariablesSection(),
              const SizedBox(height: 16),
              _buildOutputVariableSection(),
              const SizedBox(height: 16),
              _buildD4rtCodeSection(),
              const SizedBox(height: 32),
            ],
          ),
        ),
      ),
    );
  }

  Widget _buildNameSection() {
    return TextFormField(
      controller: _nameController,
      decoration: const InputDecoration(
        labelText: 'Formula Name',
        border: OutlineInputBorder(),
        prefixIcon: Icon(Icons.title),
      ),
      validator: (value) {
        if (value == null || value.trim().isEmpty) {
          return 'Name is required';
        }
        return null;
      },
    );
  }

  Widget _buildDescriptionSection() {
    return Card(
      child: Padding(
        padding: const EdgeInsets.all(16.0),
        child: Column(
          crossAxisAlignment: CrossAxisAlignment.start,
          children: [
            Row(
              mainAxisAlignment: MainAxisAlignment.spaceBetween,
              children: [
                const Text('Description (Markdown)', style: TextStyle(fontSize: 16, fontWeight: FontWeight.bold)),
                Row(
                  mainAxisSize: MainAxisSize.min,
                  children: [
                    if (_isPreviewVisible)
                      TextButton.icon(
                        icon: const Icon(Icons.visibility_off),
                        label: const Text('Hide Preview'),
                        onPressed: _hidePreview,
                      )
                    else
                      TextButton.icon(
                        icon: const Icon(Icons.visibility),
                        label: const Text('Preview'),
                        onPressed: _showPreview,
                      ),
                  ],
                ),
              ],
            ),
            const SizedBox(height: 8),
            if (_isPreviewVisible) ...[
              Container(
                padding: const EdgeInsets.all(8),
                decoration: BoxDecoration(
                  color: Theme.of(context).colorScheme.surfaceVariant,
                  borderRadius: BorderRadius.circular(8),
                ),
                child: Markdown(
                  data: _descriptionController.text,
                  shrinkWrap: true,
                  builders: {'latex': LatexElementBuilder()},
                  extensionSet: markdown.ExtensionSet([LatexBlockSyntax()], [LatexInlineSyntax()]),
                ),
              ),
              const SizedBox(height: 16),
            ],
            TextFormField(
              controller: _descriptionController,
              decoration: const InputDecoration(
                hintText: 'Enter formula description (supports Markdown and LaTeX)',
                border: OutlineInputBorder(),
              ),
              maxLines: 5,
            ),
          ],
        ),
      ),
    );
  }

  Widget _buildInputVariablesSection() {
    return Card(
      child: Padding(
        padding: const EdgeInsets.all(16.0),
        child: Column(
          crossAxisAlignment: CrossAxisAlignment.start,
          children: [
            const Text('Input Variables', style: TextStyle(fontSize: 16, fontWeight: FontWeight.bold)),
            const SizedBox(height: 8),
            ..._inputVariables.asMap().entries.map((entry) {
              final index = entry.key;
              final variable = entry.value;
              return _buildInputVariableRow(index, variable);
            }).toList(),
            const SizedBox(height: 8),
            ElevatedButton.icon(
              icon: const Icon(Icons.add),
              label: const Text('Add Input Variable'),
              onPressed: _addInputVariable,
            ),
          ],
        ),
      ),
    );
  }

  Widget _buildInputVariableRow(int index, _InputVariableRowData variable) {
    return Card(
      margin: const EdgeInsets.only(bottom: 8),
      child: Padding(
        padding: const EdgeInsets.all(8.0),
        child: Column(
          children: [
            Row(
              children: [
                Expanded(
                  flex: 2,
                  child: TextFormField(
                    controller: variable.nameController,
                    decoration: const InputDecoration(labelText: 'Name', border: OutlineInputBorder()),
                  ),
                ),
                const SizedBox(width: 8),
                Flexible(
                  flex: 1,
                  child: DropdownButtonFormField<String?>(
                    isDense: true,
                    isExpanded: true,
                    value: _getBaseUnit(variable.unit),
                    decoration: const InputDecoration(
                      border: OutlineInputBorder(),
                      labelText: "Base unit",
                      contentPadding: EdgeInsets.symmetric(horizontal: 8, vertical: 0),
                    ),
                    dropdownColor: Theme.of(context).colorScheme.surface,
                    menuMaxHeight: 300,
                    items: [
                      const DropdownMenuItem<String?>(
                        value: null,
                         child: Text('None', style: TextStyle(fontSize: 11)),
                      ),
                      ..._getAllBaseUnits().map((baseUnit) {
                        return DropdownMenuItem<String?>(
                          value: baseUnit,
                          child: Text(baseUnit, style: const TextStyle(fontSize: 11)),
                        );
                      }).toList(),
                    ],
                    onChanged: (baseUnit) {
                      setState(() {
                        variable.unit = baseUnit;
                      });
                    },
                  ),
                ),
                const SizedBox(width: 8),
                Flexible(
                  flex: 1,
                  child: DropdownButtonFormField<String?>(
                    isDense: true,
                    isExpanded: true,
                    value: variable.unit,
                    decoration: const InputDecoration(
                      border: OutlineInputBorder(),
                      labelText: "Unit",
                      contentPadding: EdgeInsets.symmetric(horizontal: 8, vertical: 0),
                    ),
                    dropdownColor: Theme.of(context).colorScheme.surface,
                    menuMaxHeight: 300,
                    items: [
                      const DropdownMenuItem<String?>(
                        value: null,
                        child: Text('None', style: TextStyle(fontSize: 10)),
                      ),
                      ..._getDerivedUnits(variable.unit).map((unit) {
                        final unitSpec = widget.corpus.getUnit(unit);
                        return DropdownMenuItem<String?>(
                          value: unit,
                          child: Text(
                            '${unitSpec.symbol}',
                            style: const TextStyle(fontSize: 10),
                            overflow: TextOverflow.ellipsis,
                          ),
                        );
                      }).toList(),
                    ],
                    onChanged: (unit) {
                      setState(() {
                        variable.unit = unit;
                      });
                    },
                  ),
                ),
                const SizedBox(width: 8),
                IconButton(
                  icon: const Icon(Icons.delete, color: Colors.red),
                  onPressed: () => _removeInputVariable(index),
                  tooltip: 'Delete variable',
                  padding: EdgeInsets.zero,
                  constraints: const BoxConstraints(),
                ),
              ],
            ),
          ],
        ),
      ),
    );
  }

  Widget _buildOutputVariableSection() {
    return Card(
      child: Padding(
        padding: const EdgeInsets.all(16.0),
        child: Column(
          crossAxisAlignment: CrossAxisAlignment.start,
          children: [
            const Text('Output Variable', style: TextStyle(fontSize: 16, fontWeight: FontWeight.bold)),
            const SizedBox(height: 8),
            Row(
              children: [
                Expanded(
                  flex: 2,
                  child: TextFormField(
                    controller: _outputVariable.nameController,
                    decoration: const InputDecoration(labelText: 'Name', border: OutlineInputBorder()),
                  ),
                ),
                const SizedBox(width: 8),
                Flexible(
                  flex: 1,
                  child: DropdownButtonFormField<String?>(
                    isDense: true,
                    isExpanded: true,
                    value: _getBaseUnit(_outputVariable.unit),
                    decoration: const InputDecoration(
                      border: OutlineInputBorder(),
                      labelText: "Base unit",
                      contentPadding: EdgeInsets.symmetric(horizontal: 8, vertical: 0),
                    ),
                    dropdownColor: Theme.of(context).colorScheme.surface,
                    menuMaxHeight: 300,
                    items: [
                      const DropdownMenuItem<String?>(
                        value: null,
                        child: Text('None', style: TextStyle(fontSize: 11)),
                      ),
                      ..._getAllBaseUnits().map((baseUnit) {
                        return DropdownMenuItem<String?>(
                          value: baseUnit,
                          child: Text(baseUnit, style: const TextStyle(fontSize: 11)),
                        );
                      }).toList(),
                    ],
                    onChanged: (baseUnit) {
                      setState(() {
                        _outputVariable.unit = baseUnit;
                      });
                    },
                  ),
                ),
                const SizedBox(width: 8),
                Flexible(
                  flex: 1,
                  child: DropdownButtonFormField<String?>(
                    isDense: true,
                    isExpanded: true,
                    value: _outputVariable.unit,
                    decoration: const InputDecoration(
                      border: OutlineInputBorder(),
                      labelText: "Unit",
                      contentPadding: EdgeInsets.symmetric(horizontal: 8, vertical: 0),
                    ),
                    dropdownColor: Theme.of(context).colorScheme.surface,
                    menuMaxHeight: 300,
                    items: [
                      const DropdownMenuItem<String?>(
                        value: null,
                        child: Text('None', style: TextStyle(fontSize: 10)),
                      ),
                      ..._getDerivedUnits(_outputVariable.unit).map((unit) {
                        final unitSpec = widget.corpus.getUnit(unit);
                        return DropdownMenuItem<String?>(
                          value: unit,
                          child: Text(
                            '${unitSpec.symbol}',
                            style: const TextStyle(fontSize: 10),
                            overflow: TextOverflow.ellipsis,
                          ),
                        );
                      }).toList(),
                    ],
                    onChanged: (unit) {
                      setState(() {
                        _outputVariable.unit = unit;
                      });
                    },
                  ),
                ),
              ],
            ),
          ],
        ),
      ),
    );
  }

  Widget _buildD4rtCodeSection() {
    return Card(
      child: Padding(
        padding: const EdgeInsets.all(16.0),
        child: Column(
          crossAxisAlignment: CrossAxisAlignment.start,
          children: [
            const Text('D4RT Code (Dart syntax)', style: TextStyle(fontSize: 16, fontWeight: FontWeight.bold)),
            const SizedBox(height: 8),
            ConstrainedBox(
              constraints: const BoxConstraints(minHeight: 200),
              child: CodeTheme(
                data: CodeThemeData(styles: monokaiSublimeTheme),
                child: SingleChildScrollView(
                  child: CodeField(controller: _d4rtCodeController),
                ),
              ),
            ),
          ],
        ),
      ),
    );
  }

  // Helper methods for unit management
  String? _getBaseUnit(String? unit) {
    if (unit == null) return null;
    try {
      return widget.corpus.getUnit(unit).baseUnit;
    } catch (e) {
      return null;
    }
  }

  List<String> _getAllBaseUnits() {
    final baseUnits = <String>{};
    for (final unit in widget.corpus.allUnits()) {
      baseUnits.add(unit.baseUnit);
    }
    return baseUnits.toList()..sort();
  }

  List<String> _getDerivedUnits(String? baseUnit) {
    if (baseUnit == null) return [];
    return widget.corpus.unitsOfSameMagnitude(baseUnit)..sort();
  }
}

// Data classes to track variable state
class _InputVariableRowData {
  final TextEditingController nameController;
  String? unit;
  List<dynamic>? values;

  _InputVariableRowData({required this.nameController, this.unit, this.values});
}

class _OutputVariableRowData {
  final TextEditingController nameController;
  String? unit;

  _OutputVariableRowData({required this.nameController, this.unit});
}