# app.py
from flask import Flask, request, jsonify, render_template_string, Response
import os
import requests
import json
import logging
from logging.handlers import RotatingFileHandler
from typing import Dict, Any, List
import time
import re

app = Flask(__name__)

# Configure logging with file output in a writable directory
log_dir = '/app/logs'
log_path = os.path.join(log_dir, 'openwebui.log')
try:
    os.makedirs(log_dir, exist_ok=True)
    log_handler = RotatingFileHandler(log_path, maxBytes=1000000, backupCount=5)
    log_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
    logging.getLogger().addHandler(log_handler)
except Exception as e:
    logging.getLogger().addHandler(logging.StreamHandler())  # Fallback to stderr
    logging.error(f"Failed to initialize file logging: {e}")
logging.getLogger().setLevel(logging.INFO)

# Configuration
OLLAMA_BASE_URL = os.getenv('OLLAMA_BASE_URL', 'https://huggingface.co/spaces/tommytracx/ollama-api')
ALLOWED_MODELS = os.getenv('ALLOWED_MODELS', 'llama2,llama2:13b,llama2:70b,codellama,neural-chat,gemma-3-270m').split(',')
MAX_TOKENS = int(os.getenv('MAX_TOKENS', '2048'))
TEMPERATURE = float(os.getenv('TEMPERATURE', '0.7'))

class OllamaManager:
    def __init__(self, base_url: str):
        self.base_url = base_url.rstrip('/')
        self.available_models = ALLOWED_MODELS
        self.refresh_models()
    
    def refresh_models(self) -> None:
        """Refresh the list of available models from Ollama API, falling back to allowed models."""
        try:
            response = requests.get(f"{self.base_url}/api/tags", timeout=10)
            response.raise_for_status()
            data = response.json()
            models = [model['name'] for model in data.get('models', [])]
            self.available_models = [model for model in models if model in ALLOWED_MODELS]
            if not self.available_models:
                self.available_models = ALLOWED_MODELS
                logging.warning("No allowed models found in API response, using ALLOWED_MODELS")
            logging.info(f"Available models: {self.available_models}")
        except requests.exceptions.ConnectionError as e:
            logging.error(f"Connection error refreshing models: {e}")
            self.available_models = ALLOWED_MODELS
        except requests.exceptions.HTTPError as e:
            logging.error(f"HTTP error refreshing models: {e}")
            self.available_models = ALLOWED_MODELS
        except Exception as e:
            logging.error(f"Unexpected error refreshing models: {e}")
            self.available_models = ALLOWED_MODELS
    
    def list_models(self) -> List[str]:
        """Return the list of available models."""
        return self.available_models
    
    def generate(self, model_name: str, prompt: str, stream: bool = False, **kwargs) -> Any:
        """Generate text using a model, with optional streaming."""
        if model_name not in self.available_models:
            logging.warning(f"Attempted to generate with unavailable model: {model_name}")
            return {"status": "error", "message": f"Model {model_name} not available"}
        
        try:
            payload = {
                "model": model_name,
                "prompt": prompt,
                "stream": stream,
                **kwargs
            }
            if stream:
                response = requests.post(f"{self.base_url}/api/generate", json=payload, stream=True, timeout=120)
                response.raise_for_status()
                return response
            else:
                response = requests.post(f"{self.base_url}/api/generate", json=payload, timeout=120)
                response.raise_for_status()
                data = response.json()
                logging.info(f"Generated response with model {model_name}")
                return {
                    "status": "success",
                    "response": data.get('response', ''),
                    "model": model_name,
                    "usage": data.get('usage', {})
                }
        except requests.exceptions.ConnectionError as e:
            logging.error(f"Connection error generating response: {e}")
            return {"status": "error", "message": f"Connection error: {str(e)}"}
        except requests.exceptions.HTTPError as e:
            logging.error(f"HTTP error generating response: {e}")
            return {"status": "error", "message": f"HTTP error: {str(e)}"}
        except Exception as e:
            logging.error(f"Unexpected error generating response: {e}")
            return {"status": "error", "message": str(e)}
    
    def health_check(self) -> Dict[str, Any]:
        """Check the health of the Ollama API."""
        try:
            response = requests.get(f"{self.base_url}/api/tags", timeout=10)
            response.raise_for_status()
            logging.info("Health check successful")
            return {"status": "healthy", "available_models": len(self.available_models)}
        except requests.exceptions.ConnectionError as e:
            logging.error(f"Health check connection error: {e}")
            return {"status": "unhealthy", "error": f"Connection error: {str(e)}"}
        except requests.exceptions.HTTPError as e:
            logging.error(f"Health check HTTP error: {e}")
            return {"status": "unhealthy", "error": f"HTTP error: {str(e)}"}
        except Exception as e:
            logging.error(f"Health check unexpected error: {e}")
            return {"status": "unhealthy", "error": str(e)}

# Initialize Ollama manager
ollama_manager = OllamaManager(OLLAMA_BASE_URL)

# HTML template for the chat interface with comprehensive, mobile-optimized UI
HTML_TEMPLATE = '''
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>OpenWebUI - Ollama Chat</title>
    <script type="importmap">
      {
        "imports": {
          "react": "https://esm.sh/react@18",
          "react-dom": "https://esm.sh/react-dom@18",
          "react-dom/": "https://esm.sh/react-dom@18/",
          "@codesandbox/sandpack-react": "https://esm.sh/@codesandbox/sandpack-react@latest"
        }
      }
    </script>
    <style>
        :root {
            --primary-color: #5a4bff;
            --secondary-color: #7b3fe4;
            --text-color: #333;
            --bg-color: #f8fafc;
            --message-bg-user: var(--primary-color);
            --message-bg-assistant: white;
            --avatar-user: var(--primary-color);
            --avatar-assistant: #2ea44f;
            --border-color: #e2e8f0;
            --input-bg: white;
            --sidebar-bg: #ffffff;
            --sidebar-border: #e2e8f0;
        }
        .dark-mode {
            --primary-color: #3b4a8c;
            --secondary-color: #4a2e6b;
            --text-color: #e2e8f0;
            --bg-color: #1a202c;
            --message-bg-user: var(--primary-color);
            --message-bg-assistant: #2d3748;
            --avatar-user: var(--primary-color);
            --avatar-assistant: #276749;
            --border-color: #4a5568;
            --input-bg: #2d3748;
            --sidebar-bg: #2d3748;
            --sidebar-border: #4a5568;
        }
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }
        body {
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
            background: linear-gradient(135deg, var(--primary-color) 0%, var(--secondary-color) 100%);
            color: var(--text-color);
            min-height: 100vh;
            overflow-x: hidden;
        }
        .container {
            display: flex;
            max-width: 100%;
            min-height: 100vh;
            background: var(--bg-color);
        }
        .sidebar {
            width: 250px;
            background: var(--sidebar-bg);
            border-right: 1px solid var(--sidebar-border);
            padding: 20px;
            position: fixed;
            height: 100%;
            transform: translateX(-100%);
            transition: transform 0.3s ease;
            z-index: 1000;
        }
        .sidebar.open {
            transform: translateX(0);
        }
        .sidebar-toggle {
            position: fixed;
            top: 10px;
            left: 10px;
            background: var(--primary-color);
            color: white;
            border: none;
            padding: 10px;
            border-radius: 8px;
            cursor: pointer;
            z-index: 1100;
        }
        .sidebar h2 {
            font-size: 1.5rem;
            margin-bottom: 20px;
        }
        .chat-history {
            list-style: none;
            overflow-y: auto;
            max-height: calc(100vh - 100px);
        }
        .chat-history-item {
            padding: 10px;
            border-radius: 8px;
            margin-bottom: 10px;
            cursor: pointer;
            transition: background 0.2s;
        }
        .chat-history-item:hover {
            background: var(--border-color);
        }
        .chat-history-item.active {
            background: var(--primary-color);
            color: white;
        }
        .main-content {
            flex: 1;
            display: flex;
            flex-direction: column;
            min-height: 100vh;
        }
        .header {
            background: linear-gradient(135deg, var(--primary-color) 0%, var(--secondary-color) 100%);
            color: white;
            padding: 20px;
            text-align: center;
            position: relative;
        }
        .theme-toggle {
            position: absolute;
            top: 20px;
            right: 20px;
            background: none;
            border: none;
            cursor: pointer;
            font-size: 1.5rem;
            color: white;
        }
        .header h1 {
            font-size: 2rem;
            margin-bottom: 10px;
            font-weight: 700;
        }
        .header p {
            font-size: 1rem;
            opacity: 0.9;
        }
        .controls {
            padding: 15px 20px;
            background: var(--bg-color);
            border-bottom: 1px solid var(--border-color);
            display: flex;
            gap: 10px;
            flex-wrap: wrap;
            justify-content: center;
        }
        .control-group {
            display: flex;
            align-items: center;
            gap: 8px;
            flex: 1;
            min-width: 200px;
        }
        .control-group label {
            font-weight: 600;
            color: var(--text-color);
            min-width: 80px;
        }
        .control-group select,
        .control-group input {
            flex: 1;
            padding: 10px;
            border: 2px solid var(--border-color);
            border-radius: 8px;
            font-size: 14px;
            transition: border-color 0.3s;
            background: var(--input-bg);
            color: var(--text-color);
        }
        .control-group select:focus,
        .control-group input:focus {
            outline: none;
            border-color: var(--primary-color);
        }
        .chat-container {
            flex: 1;
            overflow-y: auto;
            padding: 20px;
            background: var(--bg-color);
            transition: background 0.3s;
        }
        .message {
            margin-bottom: 20px;
            display: flex;
            gap: 15px;
            position: relative;
            animation: fadeIn 0.3s ease-in;
        }
        .message.user {
            flex-direction: row-reverse;
        }
        .message-avatar {
            width: 40px;
            height: 40px;
            border-radius: 50%;
            display: flex;
            align-items: center;
            justify-content: center;
            font-weight: bold;
            color: white;
            flex-shrink: 0;
            background: var(--avatar-assistant);
        }
        .message.user .message-avatar {
            background: var(--avatar-user);
        }
        .message-content {
            background: var(--message-bg-assistant);
            padding: 15px 20px;
            border-radius: 18px;
            max-width: 80%;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            line-height: 1.5;
            color: var(--text-color);
        }
        .message.user .message-content {
            background: var(--message-bg-user);
            color: white;
        }
        .code-actions {
            position: absolute;
            top: 5px;
            right: 5px;
            display: flex;
            gap: 5px;
        }
        .code-button {
            background: rgba(0,0,0,0.1);
            border: none;
            padding: 5px 10px;
            border-radius: 5px;
            cursor: pointer;
            font-size: 12px;
            color: var(--text-color);
            transition: background 0.2s;
        }
        .code-button:hover {
            background: rgba(0,0,0,0.2);
        }
        .sandbox-container {
            margin-top: 10px;
            border: 1px solid var(--border-color);
            border-radius: 8px;
            overflow: hidden;
        }
        .input-container {
            padding: 15px 20px;
            background: var(--bg-color);
            border-top: 1px solid var(--border-color);
        }
        .input-form {
            display: flex;
            gap: 10px;
            align-items: center;
        }
        .input-field {
            flex: 1;
            padding: 12px 15px;
            border: 2px solid var(--border-color);
            border-radius: 25px;
            font-size: 16px;
            transition: border-color 0.3s;
            resize: none;
            min-height: 50px;
            max-height: 150px;
            background: var(--input-bg);
            color: var(--text-color);
        }
        .input-field:focus {
            outline: none;
            border-color: var(--primary-color);
        }
        .send-button {
            padding: 12px 20px;
            background: linear-gradient(135deg, var(--primary-color) 0%, var(--secondary-color) 100%);
            color: white;
            border: none;
            border-radius: 25px;
            font-size: 16px;
            font-weight: 600;
            cursor: pointer;
            transition: transform 0.2s;
            min-width: 80px;
        }
        .send-button:hover {
            transform: translateY(-2px);
        }
        .send-button:disabled {
            opacity: 0.6;
            cursor: not-allowed;
            transform: none;
        }
        .status {
            text-align: center;
            padding: 10px;
            font-size: 14px;
            color: #6c757d;
        }
        .status.error {
            color: #dc3545;
        }
        .status.success {
            color: #28a745;
        }
        .typing-indicator {
            display: none;
            padding: 15px 20px;
            background: var(--message-bg-assistant);
            border-radius: 18px;
            color: #6c757d;
            font-style: italic;
            margin: 20px;
        }
        @keyframes fadeIn {
            from { opacity: 0; transform: translateY(10px); }
            to { opacity: 1; transform: translateY(0); }
        }
        @media (max-width: 768px) {
            .container {
                flex-direction: column;
            }
            .sidebar {
                width: 100%;
                height: auto;
                max-height: 80vh;
                position: fixed;
                top: 0;
                left: 0;
                transform: translateY(-100%);
                border-right: none;
                border-bottom: 1px solid var(--sidebar-border);
            }
            .sidebar.open {
                transform: translateY(0);
            }
            .sidebar-toggle {
                top: 10px;
                left: 10px;
                z-index: 1100;
            }
            .main-content {
                margin-top: 60px;
            }
            .controls {
                flex-direction: column;
                gap: 15px;
            }
            .control-group {
                flex-direction: column;
                align-items: stretch;
            }
            .control-group select,
            .control-group input {
                width: 100%;
            }
            .message-content {
                max-width: 90%;
            }
            .header {
                padding: 15px;
            }
            .header h1 {
                font-size: 1.8rem;
            }
            .header p {
                font-size: 0.9rem;
            }
            .input-container {
                padding: 10px 15px;
            }
            .send-button {
                padding: 10px 15px;
                min-width: 60px;
            }
        }
    </style>
</head>
<body>
    <button class="sidebar-toggle" id="sidebar-toggle">☰</button>
    <div class="container">
        <div class="sidebar" id="sidebar">
            <h2>Chat History</h2>
            <ul class="chat-history" id="chat-history">
                <!-- Chat history items will be populated here -->
            </ul>
        </div>
        <div class="main-content">
            <div class="header">
                <button class="theme-toggle" id="theme-toggle">🌙</button>
                <h1>🤖 OpenWebUI</h1>
                <p>Chat with AI models powered by Ollama on Hugging Face Spaces</p>
            </div>
            
            <div class="controls">
                <div class="control-group">
                    <label for="model-select">Model:</label>
                    <select id="model-select">
                        <option value="">Select a model...</option>
                    </select>
                </div>
                <div class="control-group">
                    <label for="temperature">Temperature:</label>
                    <div style="flex: 1; display: flex; align-items: center; gap: 8px;">
                        <input type="range" id="temperature" min="0" max="2" step="0.1" value="0.7">
                        <span id="temp-value">0.7</span>
                    </div>
                </div>
                <div class="control-group">
                    <label for="max-tokens">Max Tokens:</label>
                    <input type="number" id="max-tokens" min="1" max="4096" value="2048">
                </div>
            </div>
            
            <div class="chat-container" id="chat-container">
                <div class="message assistant">
                    <div class="message-avatar">AI</div>
                    <div class="message-content">
                        Hello! I'm your AI assistant powered by Ollama. How can I help you today?
                    </div>
                </div>
            </div>
            
            <div class="typing-indicator" id="typing-indicator">
                AI is thinking...
            </div>
            
            <div class="input-container">
                <form class="input-form" id="chat-form">
                    <textarea 
                        class="input-field" 
                        id="message-input" 
                        placeholder="Type your message here..."
                        rows="1"
                    ></textarea>
                    <button type="submit" class="send-button" id="send-button">
                        Send
                    </button>
                </form>
            </div>
            
            <div class="status" id="status"></div>
        </div>
    </div>

    <script type="module">
        import { Sandpack } from 'https://esm.sh/@codesandbox/sandpack-react@latest';

        let conversationHistory = JSON.parse(localStorage.getItem('chatHistory')) || [];
        let currentConversationId = null;
        let currentMessageDiv = null;
        let currentCodeBlocks = [];

        document.addEventListener('DOMContentLoaded', function() {
            loadModels();
            loadChatHistory();
            setupEventListeners();
            autoResizeTextarea();
        });

        function toggleTheme() {
            document.body.classList.toggle('dark-mode');
            const themeToggle = document.getElementById('theme-toggle');
            themeToggle.textContent = document.body.classList.contains('dark-mode') ? '☀️' : '🌙';
            localStorage.setItem('theme', document.body.classList.contains('dark-mode') ? 'dark' : 'light');
        }

        function loadTheme() {
            if (localStorage.getItem('theme') === 'dark') {
                document.body.classList.add('dark-mode');
                document.getElementById('theme-toggle').textContent = '☀️';
            }
        }

        function toggleSidebar() {
            const sidebar = document.getElementById('sidebar');
            sidebar.classList.toggle('open');
        }

        async function loadModels() {
            const modelSelect = document.getElementById('model-select');
            modelSelect.innerHTML = '<option value="">Loading models...</option>';
            
            try {
                const response = await fetch('/api/models');
                const data = await response.json();
                
                modelSelect.innerHTML = '<option value="">Select a model...</option>';
                
                if (data.status === 'success' && data.models.length > 0) {
                    data.models.forEach(model => {
                        const option = document.createElement('option');
                        option.value = model;
                        option.textContent = model;
                        if (model === 'gemma-3-270m') {
                            option.selected = true;
                        }
                        modelSelect.appendChild(option);
                    });
                    showStatus('Models loaded successfully', 'success');
                } else {
                    modelSelect.innerHTML = '<option value="">No models available</option>';
                    showStatus('No models available from API', 'error');
                }
            } catch (error) {
                console.error('Error loading models:', error);
                modelSelect.innerHTML = '<option value="">No models available</option>';
                showStatus('Failed to load models: ' + error.message, 'error');
            }
        }

        function loadChatHistory() {
            const chatHistoryList = document.getElementById('chat-history');
            chatHistoryList.innerHTML = '';
            conversationHistory.forEach((conv, index) => {
                const li = document.createElement('li');
                li.className = 'chat-history-item';
                li.textContent = `Chat ${index + 1} - ${new Date(conv.timestamp).toLocaleString()}`;
                li.dataset.convId = index;
                li.addEventListener('click', () => loadConversation(index));
                chatHistoryList.appendChild(li);
            });
            if (conversationHistory.length > 0) {
                loadConversation(conversationHistory.length - 1);
            }
        }

        function loadConversation(convId) {
            currentConversationId = convId;
            const chatContainer = document.getElementById('chat-container');
            chatContainer.innerHTML = '';
            const conversation = conversationHistory[convId];
            conversation.messages.forEach(msg => {
                const messageDiv = addMessage(msg.content, msg.role, false, false);
                if (msg.role === 'assistant') {
                    processCodeBlocks(messageDiv, msg.content);
                }
            });
            const historyItems = document.querySelectorAll('.chat-history-item');
            historyItems.forEach(item => item.classList.remove('active'));
            historyItems[convId].classList.add('active');
        }

        function saveConversation() {
            if (currentConversationId === null) {
                conversationHistory.push({
                    timestamp: Date.now(),
                    messages: []
                });
                currentConversationId = conversationHistory.length - 1;
            }
            localStorage.setItem('chatHistory', JSON.stringify(conversationHistory));
            loadChatHistory();
        }

        function setupEventListeners() {
            document.getElementById('chat-form').addEventListener('submit', handleSubmit);
            document.getElementById('temperature').addEventListener('input', function() {
                document.getElementById('temp-value').textContent = this.value;
            });
            document.getElementById('message-input').addEventListener('input', autoResizeTextarea);
            document.getElementById('theme-toggle').addEventListener('click', toggleTheme);
            document.getElementById('sidebar-toggle').addEventListener('click', toggleSidebar);
            loadTheme();
        }
        
        function autoResizeTextarea() {
            const textarea = document.getElementById('message-input');
            textarea.style.height = 'auto';
            textarea.style.height = Math.min(textarea.scrollHeight, 150) + 'px';
        }
        
        async function handleSubmit(e) {
            e.preventDefault();
            
            const messageInput = document.getElementById('message-input');
            const message = messageInput.value.trim();
            
            if (!message) return;
            
            const model = document.getElementById('model-select').value;
            const temperature = parseFloat(document.getElementById('temperature').value);
            const maxTokens = parseInt(document.getElementById('max-tokens').value);
            
            if (!model) {
                showStatus('Please select a model', 'error');
                return;
            }
            
            addMessage(message, 'user');
            if (currentConversationId === null) {
                saveConversation();
            }
            conversationHistory[currentConversationId].messages.push({ role: 'user', content: message });
            localStorage.setItem('chatHistory', JSON.stringify(conversationHistory));
            
            messageInput.value = '';
            autoResizeTextarea();
            showTypingIndicator(true);
            
            try {
                const response = await fetch('/api/chat', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ model, prompt: message, temperature, max_tokens: maxTokens, stream: true })
                });
                
                if (!response.ok) {
                    throw new Error('Network response was not ok');
                }
                
                showTypingIndicator(false);
                currentMessageDiv = addMessage('', 'assistant', true);
                currentCodeBlocks = [];
                
                const reader = response.body.getReader();
                const decoder = new TextDecoder();
                let accumulatedResponse = '';
                
                while (true) {
                    const { done, value } = await reader.read();
                    if (done) break;
                    
                    const chunk = decoder.decode(value);
                    const lines = chunk.split('\n');
                    
                    for (const line of lines) {
                        if (line.trim()) {
                            try {
                                const data = JSON.parse(line);
                                if (data.response) {
                                    accumulatedResponse += data.response;
                                    updateMessage(currentMessageDiv, accumulatedResponse);
                                }
                            } catch (e) {
                                console.error('Error parsing chunk:', e);
                            }
                        }
                    }
                }
                
                processCodeBlocks(currentMessageDiv, accumulatedResponse);
                conversationHistory[currentConversationId].messages.push({ role: 'assistant', content: accumulatedResponse });
                localStorage.setItem('chatHistory', JSON.stringify(conversationHistory));
                showStatus(`Response generated using ${model}`, 'success');
            } catch (error) {
                showTypingIndicator(false);
                if (currentMessageDiv) {
                    updateMessage(currentMessageDiv, 'Sorry, I encountered a network error.');
                    conversationHistory[currentConversationId].messages.push({ role: 'assistant', content: 'Sorry, I encountered a network error.' });
                    localStorage.setItem('chatHistory', JSON.stringify(conversationHistory));
                } else {
                    addMessage('Sorry, I encountered a network error.', 'assistant');
                }
                showStatus('Network error: ' + error.message, 'error');
            }
        }
        
        function addMessage(content, sender, isStreaming = false, save = true) {
            const chatContainer = document.getElementById('chat-container');
            const messageDiv = document.createElement('div');
            messageDiv.className = `message ${sender}`;
            
            const avatar = document.createElement('div');
            avatar.className = 'message-avatar';
            avatar.textContent = sender === 'user' ? 'U' : 'AI';
            
            const messageContent = document.createElement('div');
            messageContent.className = 'message-content';
            messageContent.textContent = content;
            
            messageDiv.appendChild(avatar);
            messageDiv.appendChild(messageContent);
            chatContainer.appendChild(messageDiv);
            chatContainer.scrollTop = chatContainer.scrollHeight;
            
            if (!isStreaming && save) {
                if (currentConversationId === null) {
                    saveConversation();
                }
                conversationHistory[currentConversationId].messages.push({ role: sender, content: content });
                localStorage.setItem('chatHistory', JSON.stringify(conversationHistory));
            }
            return messageDiv;
        }
        
        function updateMessage(messageDiv, content) {
            const messageContent = messageDiv.querySelector('.message-content');
            messageContent.textContent = content;
            const chatContainer = document.getElementById('chat-container');
            chatContainer.scrollTop = chatContainer.scrollHeight;
        }
        
        function processCodeBlocks(messageDiv, content) {
            const codeBlockRegex = /```(\w+)?\n([\s\S]*?)```/g;
            let match;
            let lastIndex = 0;
            const messageContent = messageDiv.querySelector('.message-content');
            const fragments = [];
            
            while ((match = codeBlockRegex.exec(content)) !== null) {
                const language = match[1] || 'javascript';
                const code = match[2].trim();
                const startIndex = match.index;
                
                if (startIndex > lastIndex) {
                    fragments.push({ type: 'text', content: content.slice(lastIndex, startIndex) });
                }
                
                fragments.push({ type: 'code', language, content: code });
                currentCodeBlocks.push({ language, content: code });
                lastIndex = codeBlockRegex.lastIndex;
            }
            
            if (lastIndex < content.length) {
                fragments.push({ type: 'text', content: content.slice(lastIndex) });
            }
            
            messageContent.innerHTML = '';
            fragments.forEach((fragment, index) => {
                if (fragment.type === 'text') {
                    const textSpan = document.createElement('span');
                    textSpan.textContent = fragment.content;
                    messageContent.appendChild(textSpan);
                } else if (fragment.type === 'code') {
                    const codeContainer = document.createElement('div');
                    codeContainer.className = 'sandbox-container';
                    
                    const actions = document.createElement('div');
                    actions.className = 'code-actions';
                    
                    const copyButton = document.createElement('button');
                    copyButton.className = 'code-button';
                    copyButton.textContent = 'Copy';
                    copyButton.onclick = () => {
                        navigator.clipboard.writeText(fragment.content);
                        showStatus('Code copied to clipboard', 'success');
                    };
                    
                    const downloadButton = document.createElement('button');
                    downloadButton.className = 'code-button';
                    downloadButton.textContent = 'Download';
                    downloadButton.onclick = () => {
                        const blob = new Blob([fragment.content], { type: 'text/plain' });
                        const url = URL.createObjectURL(blob);
                        const a = document.createElement('a');
                        a.href = url;
                        a.download = `code.${fragment.language || 'txt'}`;
                        a.click();
                        URL.revokeObjectURL(url);
                    };
                    
                    actions.appendChild(copyButton);
                    actions.appendChild(downloadButton);
                    messageContent.appendChild(actions);
                    
                    const script = document.createElement('script');
                    script.type = 'module';
                    script.textContent = `
                        import { Sandpack } from '@codesandbox/sandpack-react';
                        import { createRoot } from 'react-dom';
                        const root = createRoot(document.getElementById('sandpack-${currentConversationId}-${index}'));
                        root.render(
                            React.createElement(Sandpack, {
                                template: "${fragment.language === 'javascript' ? 'react' : fragment.language}",
                                files: {
                                    '/App.${fragment.language === 'javascript' ? 'jsx' : fragment.language}': \`${fragment.content.replace(/`/g, '\\`')}\`
                                },
                                options: {
                                    showNavigator: false,
                                    showTabs: true,
                                    showLineNumbers: true,
                                    editorHeight: 300
                                }
                            })
                        );
                    `;
                    
                    const sandboxDiv = document.createElement('div');
                    sandboxDiv.id = `sandpack-${currentConversationId}-${index}`;
                    codeContainer.appendChild(sandboxDiv);
                    codeContainer.appendChild(script);
                    messageContent.appendChild(codeContainer);
                }
            });
        }
        
        function showTypingIndicator(show) {
            const indicator = document.getElementById('typing-indicator');
            indicator.style.display = show ? 'block' : 'none';
            if (show) {
                const chatContainer = document.getElementById('chat-container');
                chatContainer.scrollTop = chatContainer.scrollHeight;
            }
        }
        
        function showStatus(message, type = '') {
            const statusDiv = document.getElementById('status');
            statusDiv.textContent = message;
            statusDiv.className = `status ${type}`;
            setTimeout(() => {
                statusDiv.textContent = '';
                statusDiv.className = 'status';
            }, 5000);
        }
    </script>
</body>
</html>
'''

@app.route('/')
def home():
    """Main chat interface."""
    return render_template_string(HTML_TEMPLATE, ollama_base_url=OLLAMA_BASE_URL, default_model=ALLOWED_MODELS)

@app.route('/api/chat', methods=['POST'])
def chat():
    """Chat API endpoint with streaming support."""
    try:
        data = request.get_json()
        if not data or 'prompt' not in data or 'model' not in data:
            logging.warning("Chat request missing 'prompt' or 'model' field")
            return jsonify({"status": "error", "message": "Prompt and model are required"}), 400
        
        prompt = data['prompt']
        model = data['model']
        temperature = data.get('temperature', TEMPERATURE)
        max_tokens = data.get('max_tokens', MAX_TOKENS)
        stream = data.get('stream', False)
        
        result = ollama_manager.generate(model, prompt, stream=stream, temperature=temperature, max_tokens=max_tokens)
        
        if stream and isinstance(result, requests.Response):
            def generate_stream():
                try:
                    for chunk in result.iter_content(chunk_size=None):
                        yield chunk
                except Exception as e:
                    logging.error(f"Streaming error: {e}")
                    yield json.dumps({"status": "error", "message": str(e)}).encode()
            return Response(generate_stream(), content_type='application/json')
        else:
            logging.info(f"Non-streaming chat response generated with model {model}")
            return jsonify(result), 200 if result["status"] == "success" else 500
    except Exception as e:
        logging.error(f"Chat endpoint error: {e}")
        return jsonify({"status": "error", "message": str(e)}), 500

@app.route('/api/models', methods=['GET'])
def get_models():
    """Get available models."""
    try:
        models = ollama_manager.list_models()
        logging.info(f"Returning models: {models}")
        return jsonify({
            "status": "success",
            "models": models,
            "count": len(models)
        })
    except Exception as e:
        logging.error(f"Models endpoint error: {e}")
        return jsonify({"status": "error", "message": str(e)}), 500

@app.route('/health', methods=['GET'])
def health_check():
    """Health check endpoint."""
    try:
        ollama_health = ollama_manager.health_check()
        return jsonify({
            "status": "healthy",
            "ollama_api": ollama_health,
            "timestamp": time.time()
        })
    except Exception as e:
        logging.error(f"Health check endpoint error: {e}")
        return jsonify({
            "status": "unhealthy",
            "error": str(e),
            "timestamp": time.time()
        }), 503

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860, debug=False)