flutter_leap_sdk 0.2.4

flutter_leap_sdk: ^0.2.4 copied to clipboard

example/lib/main.dart

import 'package:flutter/material.dart';
import 'package:flutter_leap_sdk/flutter_leap_sdk.dart';
import 'package:image_picker/image_picker.dart';
import 'dart:io';

void main() {
  runApp(const MyApp());
}

class MyApp extends StatelessWidget {
  const MyApp({super.key});

  @override
  Widget build(BuildContext context) {
    return MaterialApp(
      title: 'LEAP SDK Demo',
      theme: ThemeData(
        colorScheme: ColorScheme.fromSeed(seedColor: Colors.blue),
        useMaterial3: true,
      ),
      home: const MainTabScreen(),
    );
  }
}

class MainTabScreen extends StatefulWidget {
  const MainTabScreen({super.key});

  @override
  State<MainTabScreen> createState() => _MainTabScreenState();
}

class _MainTabScreenState extends State<MainTabScreen> {
  int _currentIndex = 0;

  final List<Widget> _screens = [
    const RegularChatScreen(),
    const TextChatScreen(),
    const VisionChatScreen(),
    const CustomDownloadScreen(),
  ];

  @override
  Widget build(BuildContext context) {
    return Scaffold(
      appBar: AppBar(
        title: const Text('LEAP SDK Demo'),
        backgroundColor: Theme.of(context).colorScheme.inversePrimary,
      ),
      body: _screens[_currentIndex],
      bottomNavigationBar: BottomNavigationBar(
        type: BottomNavigationBarType.fixed,
        currentIndex: _currentIndex,
        onTap: (index) {
          setState(() {
            _currentIndex = index;
          });
        },
        items: const [
          BottomNavigationBarItem(
            icon: Icon(Icons.chat),
            label: 'Regular Chat',
          ),
          BottomNavigationBarItem(
            icon: Icon(Icons.functions),
            label: 'Function Calling',
          ),
          BottomNavigationBarItem(
            icon: Icon(Icons.image),
            label: 'Vision Chat',
          ),
          BottomNavigationBarItem(
            icon: Icon(Icons.download),
            label: 'Custom Download',
          ),
        ],
      ),
    );
  }
}

class TextChatScreen extends StatefulWidget {
  const TextChatScreen({super.key});

  @override
  State<TextChatScreen> createState() => _TextChatScreenState();
}

class _TextChatScreenState extends State<TextChatScreen> {
  bool _isDownloading = false;
  bool _isLoading = false;
  double _downloadProgress = 0.0;
  
  Conversation? _conversation;
  final List<ChatMessage> _messages = [];
  final TextEditingController _messageController = TextEditingController();
  bool _isGenerating = false;

  @override
  void initState() {
    super.initState();
    FlutterLeapSdkService.initialize();
    _checkStatus();
  }

  Future<void> _checkStatus() async {
    try {
      final isLoaded = FlutterLeapSdkService.modelLoaded;
      setState(() {
        if (isLoaded && _conversation != null) {
          // Status: '🚀 Ready to chat! Try: "What\'s the weather in Paris?"';
        } else if (isLoaded) {
          // Status: '✅ Model ready - click Load to start chat';
        } else {
          // Status: '⬇️ Need to download model';
        }
      });
    } catch (e) {
      setState(() {
        // Status: '❌ Error: $e';
      });
    }
  }

  Future<void> _downloadModel() async {
    setState(() {
      _isDownloading = true;
      _downloadProgress = 0.0;
    });

    try {
      await FlutterLeapSdkService.downloadModel(
        modelName: 'LFM2-350M',
        onProgress: (progress) {
          setState(() {
            _downloadProgress = progress.percentage / 100.0;
            if (progress.isComplete) {
              _isDownloading = false;
              // Status: '✅ Downloaded! Loading model...';
              _loadModel();
            }
          });
        },
      );
    } catch (e) {
      setState(() {
        _isDownloading = false;
        // Status: '❌ Download failed: $e';
      });
    }
  }

  Future<void> _loadModel() async {
    setState(() {
      _isLoading = true;
    });
    
    try {
      await FlutterLeapSdkService.loadModel(modelPath: 'LFM2-350M');
      
      // Create conversation through service (creates both Dart and native conversation)
      _conversation = await FlutterLeapSdkService.createConversation(
        systemPrompt: 'You are a helpful AI assistant.',
      );
      
      // Register function
      await _conversation!.registerFunction(
        LeapFunction(
          name: 'get_weather',
          description: 'Get weather information for a location',
          parameters: [
            LeapFunctionParameter(
              name: 'location',
              type: 'string',
              description: 'The city name',
              required: true,
            ),
          ],
          implementation: (args) async {
            final argsMap = Map<String, dynamic>.from(args);
            final location = argsMap['location']?.toString() ?? 'Unknown';
            return {
              'location': location,
              'temperature': 22,
              'description': 'Sunny',
            };
          },
        ),
      );
      
      setState(() {
        _isLoading = false;
        // Status: '🚀 Ready to chat! Try: "What\'s the weather in Paris?"';
      });
    } catch (e) {
      setState(() {
        _isLoading = false;
        // Status: '❌ Load failed: $e';
      });
    }
  }

  Future<void> _sendMessage() async {
    if (_messageController.text.trim().isEmpty || _conversation == null) return;

    final userMessage = _messageController.text.trim();
    _messageController.clear();
    
    setState(() {
      _messages.add(ChatMessage.user(userMessage));
      _isGenerating = true;
    });

    try {
      String response = '';
      
      await for (final chunk in _conversation!.generateResponseStream(userMessage)) {
        response += chunk;
        setState(() {
          if (_messages.isEmpty || _messages.last.role != MessageRole.assistant) {
            _messages.add(ChatMessage.assistant(''));
          }
          _messages[_messages.length - 1] = ChatMessage.assistant(response);
        });
      }
      
    } catch (e) {
      setState(() {
        _messages.add(ChatMessage.assistant('Error: $e'));
      });
    } finally {
      setState(() {
        _isGenerating = false;
      });
    }
  }

  @override
  Widget build(BuildContext context) {
    return Column(
      children: [
        if (_conversation == null) 
          Container(
            width: double.infinity,
            padding: const EdgeInsets.all(16),
            color: Colors.grey.shade100,
            child: Column(
              children: [
                const Text(
                  'Function Calling',
                  style: TextStyle(
                    fontSize: 18,
                    fontWeight: FontWeight.bold,
                  ),
                ),
                const SizedBox(height: 12),
                Row(
                  children: [
                    if (_isDownloading) ...[
                      const CircularProgressIndicator(),
                      const SizedBox(width: 16),
                      Text('Downloading... ${(_downloadProgress * 100).toStringAsFixed(1)}%'),
                    ] else ...[
                      Expanded(
                        child: ElevatedButton(
                          onPressed: _downloadModel,
                          child: const Text('Download Model'),
                        ),
                      ),
                      const SizedBox(width: 8),
                      Expanded(
                        child: ElevatedButton(
                          onPressed: _isLoading ? null : _loadModel,
                          child: _isLoading 
                              ? const SizedBox(
                                  width: 16,
                                  height: 16,
                                  child: CircularProgressIndicator(strokeWidth: 2),
                                )
                              : const Text('Load Model'),
                        ),
                      ),
                    ],
                  ],
                ),
              ],
            ),
          ),
        
        // Chat section
        if (_conversation != null) ...[
          // Messages
            Expanded(
              child: ListView.builder(
                padding: const EdgeInsets.all(16),
                itemCount: _messages.length,
                itemBuilder: (context, index) {
                  final message = _messages[index];
                  final isUser = message.role == MessageRole.user;
                  
                  return Container(
                    margin: const EdgeInsets.symmetric(vertical: 4),
                    padding: const EdgeInsets.all(12),
                    decoration: BoxDecoration(
                      color: isUser ? Colors.blue.shade100 : Colors.grey.shade100,
                      borderRadius: BorderRadius.circular(8),
                    ),
                    child: Column(
                      crossAxisAlignment: CrossAxisAlignment.start,
                      children: [
                        Text(
                          isUser ? 'You' : 'Assistant',
                          style: TextStyle(
                            fontWeight: FontWeight.bold,
                            color: isUser ? Colors.blue.shade800 : Colors.grey.shade800,
                          ),
                        ),
                        const SizedBox(height: 4),
                        Text(message.content),
                      ],
                    ),
                  );
                },
              ),
            ),
            
            // Input section
            Container(
              padding: const EdgeInsets.all(16),
              decoration: BoxDecoration(
                color: Colors.white,
                border: Border(top: BorderSide(color: Colors.grey.shade300)),
              ),
              child: Row(
                children: [
                  Expanded(
                    child: TextField(
                      controller: _messageController,
                      decoration: const InputDecoration(
                        hintText: 'Type a message... (try: "What\'s the weather in Paris?")',
                        border: OutlineInputBorder(),
                      ),
                      onSubmitted: (_) => _sendMessage(),
                      enabled: !_isGenerating,
                    ),
                  ),
                  const SizedBox(width: 8),
                  ElevatedButton(
                    onPressed: _isGenerating ? null : _sendMessage,
                    child: _isGenerating
                        ? const SizedBox(
                            width: 20,
                            height: 20,
                            child: CircularProgressIndicator(strokeWidth: 2),
                          )
                        : const Icon(Icons.send),
                  ),
                ],
              ),
            ),
        ],
      ],
    );
  }
}

class RegularChatScreen extends StatefulWidget {
  const RegularChatScreen({super.key});

  @override
  State<RegularChatScreen> createState() => _RegularChatScreenState();
}

class _RegularChatScreenState extends State<RegularChatScreen> {
  bool _isDownloading = false;
  bool _isLoading = false;
  double _downloadProgress = 0.0;
  
  Conversation? _conversation;
  final List<ChatMessage> _messages = [];
  final TextEditingController _messageController = TextEditingController();
  bool _isGenerating = false;

  @override
  void initState() {
    super.initState();
    FlutterLeapSdkService.initialize();
    _checkStatus();
  }

  Future<void> _checkStatus() async {
    try {
      final isLoaded = FlutterLeapSdkService.modelLoaded;
      setState(() {
        if (isLoaded && _conversation != null) {
          // Status: '💬 Ready to chat!';
        } else if (isLoaded) {
          // Status: '✅ Model ready - click Load to start chat';
        } else {
          // Status: '⬇️ Need to download model';
        }
      });
    } catch (e) {
      setState(() {
        // Status: '❌ Error: $e';
      });
    }
  }

  Future<void> _downloadModel() async {
    setState(() {
      _isDownloading = true;
      _downloadProgress = 0.0;
    });

    try {
      await FlutterLeapSdkService.downloadModel(
        modelName: 'LFM2-350M',
        onProgress: (progress) {
          setState(() {
            _downloadProgress = progress.percentage / 100.0;
            if (progress.isComplete) {
              _isDownloading = false;
              // Status: '✅ Downloaded! Loading model...';
              _loadModel();
            }
          });
        },
      );
    } catch (e) {
      setState(() {
        _isDownloading = false;
        // Status: '❌ Download failed: $e';
      });
    }
  }

  Future<void> _loadModel() async {
    setState(() {
      _isLoading = true;
    });
    
    try {
      await FlutterLeapSdkService.loadModel(modelPath: 'LFM2-350M');
      
      _conversation = await FlutterLeapSdkService.createConversation(
        systemPrompt: 'You are a helpful AI assistant.',
      );
      
      setState(() {
        _isLoading = false;
        // Status: '💬 Ready to chat!';
      });
    } catch (e) {
      setState(() {
        _isLoading = false;
        // Status: '❌ Load failed: $e';
      });
    }
  }

  Future<void> _sendMessage() async {
    if (_messageController.text.trim().isEmpty || _conversation == null) return;

    final userMessage = _messageController.text.trim();
    _messageController.clear();

    setState(() {
      _messages.add(ChatMessage.user(userMessage));
      _isGenerating = true;
    });

    try {
      String response = '';
      await for (final chunk in _conversation!.generateResponseStream(userMessage)) {
        response += chunk;
        setState(() {
          if (_messages.isEmpty || _messages.last.role != MessageRole.assistant) {
            _messages.add(ChatMessage.assistant(''));
          }
          _messages[_messages.length - 1] = ChatMessage.assistant(response);
        });
      }
    } catch (e) {
      setState(() {
        _messages.add(ChatMessage.assistant('Error: $e'));
      });
    } finally {
      setState(() {
        _isGenerating = false;
      });
    }
  }

  @override
  Widget build(BuildContext context) {
    return Column(
      children: [
        if (_conversation == null) 
          Container(
            width: double.infinity,
            padding: const EdgeInsets.all(16),
            color: Colors.grey.shade100,
            child: Column(
              children: [
                const Text(
                  'Regular Chat',
                  style: TextStyle(
                    fontSize: 18,
                    fontWeight: FontWeight.bold,
                  ),
                ),
                const SizedBox(height: 12),
                Row(
                  children: [
                    if (_isDownloading) ...[
                      const CircularProgressIndicator(),
                      const SizedBox(width: 16),
                      Text('Downloading... ${(_downloadProgress * 100).toStringAsFixed(1)}%'),
                    ] else ...[
                      Expanded(
                        child: ElevatedButton(
                          onPressed: _downloadModel,
                          child: const Text('Download Model'),
                        ),
                      ),
                      const SizedBox(width: 8),
                      Expanded(
                        child: ElevatedButton(
                          onPressed: _isLoading ? null : _loadModel,
                          child: _isLoading 
                              ? const SizedBox(
                                  width: 16,
                                  height: 16,
                                  child: CircularProgressIndicator(strokeWidth: 2),
                                )
                              : const Text('Load Model'),
                        ),
                      ),
                    ],
                  ],
                ),
              ],
            ),
          ),
        
        // Chat section
        if (_conversation != null) ...[
          // Messages
          Expanded(
            child: ListView.builder(
              padding: const EdgeInsets.all(16),
              itemCount: _messages.length,
              itemBuilder: (context, index) {
                final message = _messages[index];
                final isUser = message.role == MessageRole.user;
                
                return Container(
                  margin: const EdgeInsets.symmetric(vertical: 4),
                  padding: const EdgeInsets.all(12),
                  decoration: BoxDecoration(
                    color: isUser ? Colors.blue.shade100 : Colors.grey.shade100,
                    borderRadius: BorderRadius.circular(8),
                  ),
                  child: Column(
                    crossAxisAlignment: CrossAxisAlignment.start,
                    children: [
                      Text(
                        isUser ? 'You' : 'Assistant',
                        style: TextStyle(
                          fontWeight: FontWeight.bold,
                          color: isUser ? Colors.blue.shade800 : Colors.grey.shade800,
                        ),
                      ),
                      const SizedBox(height: 4),
                      Text(message.content),
                    ],
                  ),
                );
              },
            ),
          ),
          
          // Input section
          Container(
            padding: const EdgeInsets.all(16),
            decoration: BoxDecoration(
              color: Colors.white,
              border: Border(top: BorderSide(color: Colors.grey.shade300)),
            ),
            child: Row(
              children: [
                Expanded(
                  child: TextField(
                    controller: _messageController,
                    decoration: const InputDecoration(
                      hintText: 'Type a message...',
                      border: OutlineInputBorder(),
                    ),
                    onSubmitted: (_) => _sendMessage(),
                    enabled: !_isGenerating,
                  ),
                ),
                const SizedBox(width: 8),
                ElevatedButton(
                  onPressed: _isGenerating ? null : _sendMessage,
                  child: _isGenerating
                      ? const SizedBox(
                          width: 20,
                          height: 20,
                          child: CircularProgressIndicator(strokeWidth: 2),
                        )
                      : const Icon(Icons.send),
                ),
              ],
            ),
          ),
        ],
      ],
    );
  }
}

class VisionChatScreen extends StatefulWidget {
  const VisionChatScreen({super.key});

  @override
  State<VisionChatScreen> createState() => _VisionChatScreenState();
}

class _VisionChatScreenState extends State<VisionChatScreen> {
  bool _isDownloading = false;
  bool _isLoading = false;
  double _downloadProgress = 0.0;
  
  Conversation? _conversation;
  final TextEditingController _messageController = TextEditingController();
  bool _isGenerating = false;
  String _currentResponse = '';
  
  File? _selectedImage;
  final ImagePicker _picker = ImagePicker();

  @override
  void initState() {
    super.initState();
    FlutterLeapSdkService.initialize();
    _checkStatus();
  }

  Future<void> _checkStatus() async {
    try {
      final isLoaded = FlutterLeapSdkService.modelLoaded;
      final currentModel = FlutterLeapSdkService.currentModel;
      final isVisionModel = currentModel.contains('VL') || currentModel.contains('Vision');
      
      setState(() {
        if (isLoaded && isVisionModel && _conversation != null) {
          // Status: '🖼️ Vision model ready! Select an image and ask about it.';
        } else if (isLoaded && !isVisionModel) {
          // Status: '⚠️ Please load a vision model (LFM2-VL-450M) for image processing';
        } else if (isLoaded) {
          // Status: '✅ Model ready - click Load to start vision chat';
        } else {
          // Status: '⬇️ Need to download vision model';
        }
      });
    } catch (e) {
      setState(() {
        // Status: '❌ Error: $e';
      });
    }
  }

  Future<void> _downloadModel() async {
    setState(() {
      _isDownloading = true;
      _downloadProgress = 0.0;
    });

    try {
      await FlutterLeapSdkService.downloadModel(
        modelName: 'LFM2-VL-450M (Vision)',
        onProgress: (progress) {
          setState(() {
            _downloadProgress = progress.percentage / 100.0;
            if (progress.isComplete) {
              _isDownloading = false;
              // Status: '✅ Downloaded! Loading vision model...';
              _loadModel();
            }
          });
        },
      );
    } catch (e) {
      setState(() {
        _isDownloading = false;
        // Status: '❌ Download failed: $e';
      });
    }
  }

  Future<void> _loadModel() async {
    setState(() {
      _isLoading = true;
    });
    
    try {
      await FlutterLeapSdkService.loadModel(modelPath: 'LFM2-VL-450M (Vision)');
      
      _conversation = await FlutterLeapSdkService.createConversation(
        systemPrompt: 'You are a helpful AI assistant that can see and analyze images. Describe what you see in detail and answer questions about the images.',
      );
      
      setState(() {
        _isLoading = false;
        // Status: '🖼️ Vision model ready! Select an image and ask about it.';
      });
    } catch (e) {
      setState(() {
        _isLoading = false;
        // Status: '❌ Load failed: $e';
      });
    }
  }

  Future<void> _pickImage() async {
    try {
      final XFile? image = await _picker.pickImage(source: ImageSource.gallery);
      if (image != null) {
        setState(() {
          _selectedImage = File(image.path);
        });
      }
    } catch (e) {
      if (mounted) {
        ScaffoldMessenger.of(context).showSnackBar(
          SnackBar(content: Text('Failed to pick image: $e')),
        );
      }
    }
  }

  Future<void> _sendMessage() async {
    if (_messageController.text.trim().isEmpty || _conversation == null) return;
    if (_selectedImage == null) {
      ScaffoldMessenger.of(context).showSnackBar(
        const SnackBar(content: Text('Please select an image first')),
      );
      return;
    }

    final userMessage = _messageController.text.trim();
    _messageController.clear();

    setState(() {
      _currentResponse = '';
      _isGenerating = true;
    });

    try {
      final imageBytes = await _selectedImage!.readAsBytes();
      
      // Clear previous response and start streaming
      setState(() {
        _currentResponse = '';
      });
      
      String response = '';
      
      try {
        await for (final chunk in _conversation!.generateResponseWithImageStream(userMessage, imageBytes)) {
          response += chunk;
          
          if (mounted) {
            setState(() {
              _currentResponse = response;
            });
          }
        }
      } catch (streamError) {
        final nonStreamResponse = await _conversation!.generateResponseWithImage(userMessage, imageBytes);
        setState(() {
          _currentResponse = nonStreamResponse;
        });
      }
      
    } catch (e) {
      setState(() {
        _currentResponse = 'Error: $e';
      });
    } finally {
      setState(() {
        _isGenerating = false;
      });
    }
  }

  @override
  Widget build(BuildContext context) {
    return Column(
      children: [
        // Status section
        Container(
          width: double.infinity,
          padding: const EdgeInsets.all(16),
          color: Colors.grey.shade100,
          child: Column(
            children: [
              const Text(
                'Vision Chat',
                style: TextStyle(
                  fontSize: 18,
                  fontWeight: FontWeight.bold,
                ),
              ),
              const SizedBox(height: 12),
              Row(
                children: [
                  if (_isDownloading) ...[
                    const CircularProgressIndicator(),
                    const SizedBox(width: 16),
                    Text('Downloading... ${(_downloadProgress * 100).toStringAsFixed(1)}%'),
                  ] else if (_conversation == null) ...[
                    Expanded(
                      child: ElevatedButton(
                        onPressed: _downloadModel,
                        child: const Text('Download Vision'),
                      ),
                    ),
                    const SizedBox(width: 8),
                    Expanded(
                      child: ElevatedButton(
                        onPressed: _isLoading ? null : _loadModel,
                        child: _isLoading 
                            ? const SizedBox(
                                width: 16,
                                height: 16,
                                child: CircularProgressIndicator(strokeWidth: 2),
                              )
                            : const Text('Load Vision'),
                      ),
                    ),
                  ],
                ],
              ),
            ],
          ),
        ),
        
        // Chat section
        if (_conversation != null) ...[
          // Image selection
          Container(
            width: double.infinity,
            padding: const EdgeInsets.all(16),
            color: Colors.blue.shade50,
            child: Column(
              children: [
                if (_selectedImage != null) ...[
                  Row(
                    children: [
                      ClipRRect(
                        borderRadius: BorderRadius.circular(4),
                        child: Image.file(
                          _selectedImage!,
                          height: 32,
                          width: 32,
                          fit: BoxFit.cover,
                        ),
                      ),
                      const SizedBox(width: 8),
                      Expanded(
                        child: Text(
                          _selectedImage!.path.split('/').last,
                          style: const TextStyle(fontSize: 14),
                          overflow: TextOverflow.ellipsis,
                        ),
                      ),
                    ],
                  ),
                  const SizedBox(height: 8),
                ],
                Row(
                  children: [
                    ElevatedButton.icon(
                      onPressed: _pickImage,
                      icon: const Icon(Icons.photo_library),
                      label: Text(_selectedImage == null ? 'Select Image' : 'Change Image'),
                    ),
                    if (_selectedImage != null) ...[
                      const SizedBox(width: 8),
                      TextButton(
                        onPressed: () => setState(() => _selectedImage = null),
                        child: const Text('Clear'),
                      ),
                    ],
                  ],
                ),
              ],
            ),
          ),
          
          // Response Output
          Expanded(
            child: Container(
              margin: const EdgeInsets.all(16),
              padding: const EdgeInsets.all(16),
              decoration: BoxDecoration(
                color: Colors.grey.shade50,
                borderRadius: BorderRadius.circular(8),
                border: Border.all(color: Colors.grey.shade300),
              ),
              child: Column(
                crossAxisAlignment: CrossAxisAlignment.start,
                children: [
                  Row(
                    children: [
                      Icon(Icons.smart_toy, color: Colors.grey.shade600, size: 20),
                      const SizedBox(width: 8),
                      Text(
                        'Vision Assistant Response:',
                        style: TextStyle(
                          fontWeight: FontWeight.bold,
                          color: Colors.grey.shade700,
                          fontSize: 16,
                        ),
                      ),
                    ],
                  ),
                  const SizedBox(height: 12),
                  Expanded(
                    child: SingleChildScrollView(
                      child: Column(
                        crossAxisAlignment: CrossAxisAlignment.start,
                        children: [
                          if (_isGenerating && _currentResponse.isEmpty) ...[
                            const Center(child: CircularProgressIndicator()),
                            const SizedBox(height: 16),
                            Center(
                              child: Text(
                                'Analyzing image...',
                                style: TextStyle(
                                  color: Colors.grey.shade600,
                                  fontStyle: FontStyle.italic,
                                ),
                              ),
                            ),
                          ],
                          if (_currentResponse.isNotEmpty) ...[
                            if (_isGenerating) ...[
                              Row(
                                children: [
                                  SizedBox(
                                    width: 12,
                                    height: 12,
                                    child: CircularProgressIndicator(strokeWidth: 2),
                                  ),
                                  const SizedBox(width: 8),
                                  Text(
                                    'Generating...',
                                    style: TextStyle(
                                      color: Colors.blue.shade600,
                                      fontSize: 12,
                                      fontStyle: FontStyle.italic,
                                    ),
                                  ),
                                ],
                              ),
                              const SizedBox(height: 8),
                            ],
                            SelectableText(
                              _currentResponse,
                              style: const TextStyle(
                                fontSize: 16,
                                height: 1.5,
                              ),
                            ),
                          ],
                          if (!_isGenerating && _currentResponse.isEmpty)
                            Text(
                              'Select an image and ask a question to see the response here.',
                              style: TextStyle(
                                color: Colors.grey.shade500,
                                fontStyle: FontStyle.italic,
                              ),
                            ),
                        ],
                      ),
                    ),
                  ),
                ],
              ),
            ),
          ),
          
          // Input section
          Container(
            padding: const EdgeInsets.all(16),
            decoration: BoxDecoration(
              color: Colors.white,
              border: Border(top: BorderSide(color: Colors.grey.shade300)),
            ),
            child: Row(
              children: [
                Expanded(
                  child: TextField(
                    controller: _messageController,
                    decoration: const InputDecoration(
                      hintText: 'Ask about the image... (try: "What do you see in this image?")',
                      border: OutlineInputBorder(),
                    ),
                    onSubmitted: (_) => _sendMessage(),
                    enabled: !_isGenerating,
                  ),
                ),
                const SizedBox(width: 8),
                ElevatedButton(
                  onPressed: _isGenerating ? null : _sendMessage,
                  child: _isGenerating
                      ? const SizedBox(
                          width: 20,
                          height: 20,
                          child: CircularProgressIndicator(strokeWidth: 2),
                        )
                      : const Icon(Icons.send),
                ),
              ],
            ),
          ),
        ],
      ],
    );
  }
}

class CustomDownloadScreen extends StatefulWidget {
  const CustomDownloadScreen({super.key});

  @override
  State<CustomDownloadScreen> createState() => _CustomDownloadScreenState();
}

class _CustomDownloadScreenState extends State<CustomDownloadScreen> {
  final TextEditingController _urlController = TextEditingController();
  final TextEditingController _nameController = TextEditingController();
  bool _isDownloading = false;
  bool _isLoading = false;
  double _downloadProgress = 0.0;
  String _status = '📥 Ready to download custom models';
  String _downloadSpeed = '';
  
  Conversation? _conversation;
  final List<ChatMessage> _messages = [];
  final TextEditingController _messageController = TextEditingController();
  bool _isGenerating = false;

  @override
  void initState() {
    super.initState();
    FlutterLeapSdkService.initialize();
    _checkStatus();
  }

  Future<void> _checkStatus() async {
    try {
      final isLoaded = FlutterLeapSdkService.modelLoaded;
      setState(() {
        if (isLoaded && _conversation != null) {
          _status = '🚀 Custom model ready to chat!';
        } else if (isLoaded) {
          _status = '✅ Model ready - click Load to start chat';
        } else {
          _status = '📥 Enter URL and name to download custom model';
        }
      });
    } catch (e) {
      setState(() {
        // Status: '❌ Error: $e';
      });
    }
  }

  Future<void> _downloadCustomModel() async {
    if (_urlController.text.trim().isEmpty || _nameController.text.trim().isEmpty) {
      ScaffoldMessenger.of(context).showSnackBar(
        const SnackBar(content: Text('Please enter both URL and model name')),
      );
      return;
    }

    setState(() {
      _isDownloading = true;
      _downloadProgress = 0.0;
      _downloadSpeed = '';
      // Status: 'Downloading custom model...';
    });

    try {
      await FlutterLeapSdkService.downloadModel(
        modelUrl: _urlController.text.trim(),
        modelName: _nameController.text.trim(),
        onProgress: (progress) {
          setState(() {
            _downloadProgress = progress.percentage / 100.0;
            _downloadSpeed = progress.speed;
            // Status: 'Downloading... ${progress.percentage.toStringAsFixed(1)}% (${progress.speed})';
            if (progress.isComplete) {
              _isDownloading = false;
              // Status: '✅ Downloaded! Click Load to use the model.';
            }
          });
        },
      );
    } catch (e) {
      setState(() {
        _isDownloading = false;
        // Status: '❌ Download failed: $e';
      });
    }
  }

  Future<void> _loadCustomModel() async {
    if (_nameController.text.trim().isEmpty) {
      ScaffoldMessenger.of(context).showSnackBar(
        const SnackBar(content: Text('Please enter model name')),
      );
      return;
    }

    setState(() {
      _isLoading = true;
      // Status: 'Loading custom model...';
    });
    
    try {
      await FlutterLeapSdkService.loadModel(modelPath: _nameController.text.trim());
      
      _conversation = await FlutterLeapSdkService.createConversation(
        systemPrompt: 'You are a helpful AI assistant.',
      );
      
      setState(() {
        _isLoading = false;
        // Status: '🚀 Custom model ready to chat!';
      });
    } catch (e) {
      setState(() {
        _isLoading = false;
        // Status: '❌ Load failed: $e';
      });
    }
  }

  Future<void> _sendMessage() async {
    if (_messageController.text.trim().isEmpty || _conversation == null) return;

    final userMessage = _messageController.text.trim();
    _messageController.clear();

    setState(() {
      _messages.add(ChatMessage.user(userMessage));
      _isGenerating = true;
    });

    try {
      String response = '';
      await for (final chunk in _conversation!.generateResponseStream(userMessage)) {
        response += chunk;
        setState(() {
          if (_messages.isEmpty || _messages.last.role != MessageRole.assistant) {
            _messages.add(ChatMessage.assistant(''));
          }
          _messages[_messages.length - 1] = ChatMessage.assistant(response);
        });
      }
    } catch (e) {
      setState(() {
        _messages.add(ChatMessage.assistant('Error: $e'));
      });
    } finally {
      setState(() {
        _isGenerating = false;
      });
    }
  }

  @override
  Widget build(BuildContext context) {
    return Column(
      children: [
        // Header section
        Container(
          width: double.infinity,
          padding: const EdgeInsets.all(16),
          color: Colors.grey.shade100,
          child: Column(
            crossAxisAlignment: CrossAxisAlignment.start,
            children: [
              Row(
                children: [
                  const Icon(Icons.download, size: 24),
                  const SizedBox(width: 8),
                  const Text(
                    'Custom Model Download',
                    style: TextStyle(
                      fontSize: 18,
                      fontWeight: FontWeight.bold,
                    ),
                  ),
                ],
              ),
              const SizedBox(height: 8),
              Text(
                _status,
                style: TextStyle(
                  fontSize: 14,
                  color: Colors.grey.shade700,
                ),
              ),
            ],
          ),
        ),
        
        // Input section
        if (_conversation == null) ...[
          Container(
            padding: const EdgeInsets.all(16),
            child: Column(
              children: [
                TextField(
                  controller: _urlController,
                  decoration: const InputDecoration(
                    labelText: 'Model URL',
                    hintText: 'https://example.com/model.bundle',
                    border: OutlineInputBorder(),
                    prefixIcon: Icon(Icons.link),
                  ),
                  enabled: !_isDownloading && !_isLoading,
                ),
                const SizedBox(height: 16),
                TextField(
                  controller: _nameController,
                  decoration: const InputDecoration(
                    labelText: 'Model Name',
                    hintText: 'my-custom-model',
                    border: OutlineInputBorder(),
                    prefixIcon: Icon(Icons.edit),
                  ),
                  enabled: !_isDownloading && !_isLoading,
                ),
                const SizedBox(height: 16),
                
                if (_isDownloading) ...[
                  Column(
                    children: [
                      LinearProgressIndicator(value: _downloadProgress),
                      const SizedBox(height: 8),
                      Text(
                        '${(_downloadProgress * 100).toStringAsFixed(1)}% - $_downloadSpeed',
                        style: TextStyle(color: Colors.grey.shade600),
                      ),
                    ],
                  ),
                ] else ...[
                  Row(
                    children: [
                      Expanded(
                        child: ElevatedButton.icon(
                          onPressed: _downloadCustomModel,
                          icon: const Icon(Icons.download),
                          label: const Text('Download'),
                        ),
                      ),
                      const SizedBox(width: 16),
                      Expanded(
                        child: ElevatedButton.icon(
                          onPressed: _isLoading ? null : _loadCustomModel,
                          icon: _isLoading 
                              ? const SizedBox(
                                  width: 16,
                                  height: 16,
                                  child: CircularProgressIndicator(strokeWidth: 2),
                                )
                              : const Icon(Icons.play_arrow),
                          label: const Text('Load'),
                        ),
                      ),
                    ],
                  ),
                ],
                
                const SizedBox(height: 16),
                const Divider(),
                const SizedBox(height: 8),
                
                // Example URLs
                Text(
                  'Example URLs:',
                  style: TextStyle(
                    fontWeight: FontWeight.bold,
                    color: Colors.grey.shade700,
                  ),
                ),
                const SizedBox(height: 8),
                _buildExampleUrlTile(
                  'LFM2-350M',
                  'https://huggingface.co/LiquidAI/LeapBundles/resolve/main/LFM2-350M-8da4w_output_8da8w-seq_4096.bundle?download=true',
                ),
                _buildExampleUrlTile(
                  'LFM2-VL-450M',
                  'https://huggingface.co/LiquidAI/LeapBundles/resolve/main/LFM2-VL-450M_8da4w.bundle?download=true',
                ),
              ],
            ),
          ),
        ],
        
        // Chat section
        if (_conversation != null) ...[
          // Messages
          Expanded(
            child: ListView.builder(
              padding: const EdgeInsets.all(16),
              itemCount: _messages.length,
              itemBuilder: (context, index) {
                final message = _messages[index];
                final isUser = message.role == MessageRole.user;
                
                return Container(
                  margin: const EdgeInsets.symmetric(vertical: 4),
                  padding: const EdgeInsets.all(12),
                  decoration: BoxDecoration(
                    color: isUser ? Colors.blue.shade100 : Colors.grey.shade100,
                    borderRadius: BorderRadius.circular(8),
                  ),
                  child: Column(
                    crossAxisAlignment: CrossAxisAlignment.start,
                    children: [
                      Text(
                        isUser ? 'You' : 'Custom Model',
                        style: TextStyle(
                          fontWeight: FontWeight.bold,
                          color: isUser ? Colors.blue.shade800 : Colors.grey.shade800,
                        ),
                      ),
                      const SizedBox(height: 4),
                      Text(message.content),
                    ],
                  ),
                );
              },
            ),
          ),
          
          // Input section
          Container(
            padding: const EdgeInsets.all(16),
            decoration: BoxDecoration(
              color: Colors.white,
              border: Border(top: BorderSide(color: Colors.grey.shade300)),
            ),
            child: Row(
              children: [
                Expanded(
                  child: TextField(
                    controller: _messageController,
                    decoration: const InputDecoration(
                      hintText: 'Chat with your custom model...',
                      border: OutlineInputBorder(),
                    ),
                    onSubmitted: (_) => _sendMessage(),
                    enabled: !_isGenerating,
                  ),
                ),
                const SizedBox(width: 8),
                ElevatedButton(
                  onPressed: _isGenerating ? null : _sendMessage,
                  child: _isGenerating
                      ? const SizedBox(
                          width: 20,
                          height: 20,
                          child: CircularProgressIndicator(strokeWidth: 2),
                        )
                      : const Icon(Icons.send),
                ),
              ],
            ),
          ),
        ],
      ],
    );
  }

  Widget _buildExampleUrlTile(String name, String url) {
    return Card(
      child: ListTile(
        title: Text(name),
        subtitle: Text(
          url,
          maxLines: 2,
          overflow: TextOverflow.ellipsis,
          style: TextStyle(fontSize: 12, color: Colors.grey.shade600),
        ),
        trailing: IconButton(
          icon: const Icon(Icons.copy),
          onPressed: () {
            _urlController.text = url;
            _nameController.text = name;
            ScaffoldMessenger.of(context).showSnackBar(
              SnackBar(content: Text('Copied $name URL and name')),
            );
          },
        ),
      ),
    );
  }

  @override
  void dispose() {
    _urlController.dispose();
    _nameController.dispose();
    _messageController.dispose();
    super.dispose();
  }
}