App.vue

<script setup>
import { ref, computed, onMounted, onUnmounted, watch, nextTick, triggerRef } from 'vue';
import Chat from './components/Chat.vue';
import ChatTabs from './components/ChatTabs.vue';
import Progress from './components/Progress.vue';
import Settings from './components/Settings.vue';
import QuickSettings from './components/QuickSettings.vue';
import ArrowRightIcon from './components/icons/ArrowRightIcon.vue';
import StopIcon from './components/icons/StopIcon.vue';
import SettingsIcon from './components/icons/SettingsIcon.vue';
import MicrophoneIcon from './components/icons/MicrophoneIcon.vue';
import workerUrl from './worker.js?url';
import apiWorkerUrl from './apiWorker.js?url';
import sttWorkerUrl from './stt-worker.js?url';
import chatStorage from './services/chatStorage.js';
import { MicVAD } from '@ricky0123/vad-web';
import { useTTS } from './composables/useTTS.js';
import { segmentText } from './utils/clause-chunker.js';
import { phonemizeText, buildTokenIdsFromPhonemes } from './utils/phoneme-utils.js';
import { cleanTextForTTS } from './utils/textCleaner.js';

const isWebGpuAvailable = ref(false);
const STICKY_SCROLL_THRESHOLD = 120;
const EXAMPLES = [
  'Solve the equation x^2 - 3x + 2 = 0',
  'Lily is three times older than her son. In 15 years, she will be twice as old as him. How old was she when he was born, and how old is she now?',
  'Write python code to compute the nth fibonacci number.',
];

const worker = ref(null);
const sttWorker = ref(null);
const textareaRef = ref(null);
const chatContainerRef = ref(null);
const chatRef = ref(null);
const bottomPaddingRef = ref(null);
const settingsOpen = ref(false);
const pendingClassification = ref(false);
const lastClassifiedText = ref('');
const vad = ref(null);
const isVoiceActive = ref(false);
const vadGate = ref(false);
const incompleteTimer = ref(null);
const lastIncompleteText = ref('');
const userStartedRecording = ref(false);
const currentSettings = ref({
  provider: 'local',
  localModel: 'onnx-community/Qwen3-4B-ONNX',
  customApis: [],
  selectedCustomApi: null,
  defaultCustomApi: null,
  systemPrompts: []
});

const status = ref(null);
const loadingMessage = ref('');
const progressItems = ref([]);
const isRunning = ref(false);

const input = ref('');
const messages = ref([]);
const tps = ref(null);
const numTokens = ref(null);
const lastAiMessageIndex = ref(-1);

const chats = ref([]);
const currentChatId = ref(null);
const sidebarCollapsed = ref(false);
const isLoadingHistory = ref(false);
const currentSystemPromptId = ref('');
const currentChatProvider = ref('local');
const currentChatSelectedCustomApi = ref(null);

// STT state
const isRecording = ref(false);
const sttStatus = ref(null);
const mediaRecorder = ref(null);
const audioStream = ref(null);
const audioChunks = ref([]);
const audioContext = ref(null);
const isSTTProcessing = ref(false);
const sttText = ref('');
// TTS state and helpers
const { init: initTTS, unlock: unlockTTS, flush: flushTTS, speakTokenIds, enqueueSilence, isReady: isTTSReady, isSpeaking: isTTSSpeaking } = useTTS();
const ttsModelSampleRate = 24000;
let currentTTSAbortController = null;

async function ensureTTSReady() {
  try {
    await initTTS({ modelUrl: '/tts-model/kitten_tts_nano_v0_1.onnx', voicesJsonUrl: '/tts-model/voices.json' });
    await unlockTTS();
  } catch (e) {
    console.error('TTS init error:', e);
  }
}

async function speakDelta(delta) {
  try {
    const clauses = segmentText(delta);
    for (const c of clauses) {
      const ph = await phonemizeText(c.text);
      const phonemeString = Array.isArray(ph) ? ph.join(' ') : ph;
      const tokenIds = buildTokenIdsFromPhonemes(phonemeString);
      speakTokenIds(tokenIds, { speed: 1.0, modelSampleRate: ttsModelSampleRate });
      if (c.pauseMs > 0) enqueueSilence(c.pauseMs);
    }
  } catch (e) {
    console.error('speakDelta error:', e);
  }
}

async function speakText(text) {
  try {
    // Cancel any existing TTS process
    if (currentTTSAbortController) {
      currentTTSAbortController.abort();
    }
    
    // Create new abort controller for this TTS session
    currentTTSAbortController = new AbortController();
    const abortSignal = currentTTSAbortController.signal;
    
    await ensureTTSReady();
    flushTTS();
    
    // Clean the text for TTS (remove markdown, emojis, unwanted symbols)
    const cleanedText = cleanTextForTTS(text);
    console.log('Original text:', text.substring(0, 200) + '...');
    console.log('Cleaned text:', cleanedText.substring(0, 200) + '...');
    
    if (!cleanedText.trim()) {
      console.log('No readable text found after cleaning');
      return;
    }
    
    const segments = segmentText(cleanedText);
    for (const c of segments) {
      // Check if we've been cancelled
      if (abortSignal.aborted) {
        console.log('TTS cancelled');
        return;
      }
      
      const ph = await phonemizeText(c.text);
      const phonemeString = Array.isArray(ph) ? ph.join(' ') : ph;
      const tokenIds = buildTokenIdsFromPhonemes(phonemeString);
      speakTokenIds(tokenIds, { speed: 1.0, modelSampleRate: ttsModelSampleRate });
      
      // Skip adding silence between segments to eliminate hiccups - natural pauses will come from punctuation
      // if (c.pauseMs > 0) enqueueSilence(Math.min(c.pauseMs, 50)); // Commented out to test hiccup elimination
    }
  } catch (error) {
    if (error.name === 'AbortError') {
      console.log('TTS was cancelled');
    } else {
      console.error('Error speaking text:', error);
      throw error;
    }
  }
}

function stopTTS() {
  // Cancel the TTS processing loop
  if (currentTTSAbortController) {
    currentTTSAbortController.abort();
    currentTTSAbortController = null;
  }
  
  // Flush the audio queue
  flushTTS();
}

const onEnter = async (message) => {
  const userMessage = { role: 'user', content: message };
  const userMessageIndex = messages.value.length;
  messages.value.push(userMessage);

  if (currentChatId.value) {
    try {
      await chatStorage.saveMessage(currentChatId.value, userMessage);
      await updateChatsList();
    } catch (error) {
      console.error('Failed to save user message:', error);
    }
  }

  tps.value = null;
  isRunning.value = true;
  input.value = '';

  // Stop recording and warm up STT worker before restarting
  if (isRecording.value) {
    stopContinuousRecording();

    // Reset and warm up STT worker
    if (sttWorker.value) {
      sttWorker.value.postMessage({ type: 'reset' });
      sttStatus.value = 'warming';
      sttWorker.value.postMessage({ type: 'warmup' });
    }
  }

  sttText.value = '';

  // Scroll to the user's message after it's added
  await nextTick();
  scrollToUserMessage(userMessageIndex);
};

const handleEnter = () => {
  if (input.value.length > 0 && !isRunning.value) {
    onEnter(input.value);
  }
};

const onInterrupt = () => {
  worker.value.postMessage({ type: 'interrupt' });
};

const resizeInput = () => {
  if (!textareaRef.value) return;
  const target = textareaRef.value;
  target.style.height = 'auto';
  const newHeight = Math.min(Math.max(target.scrollHeight, 24), 200);
  target.style.height = `${newHeight}px`;
};

const onMessageReceived = async (e) => {
  switch (e.data.status) {
    case 'loading':
      status.value = 'loading';
      loadingMessage.value = e.data.data;
      break;
    case 'initiate':
      progressItems.value.push(e.data);
      break;
    case 'progress':
      progressItems.value = progressItems.value.map((item) => {
        if (item.file === e.data.file) {
          return { ...item, ...e.data };
        }
        return item;
      });
      break;
    case 'done':
      progressItems.value = progressItems.value.filter(
        (item) => item.file !== e.data.file
      );
      break;
    case 'ready':
      status.value = 'ready';
      break;
    case 'start':
      messages.value.push({ role: 'assistant', content: '', reasoning: '' });
      // Update last AI message index to the new message
      lastAiMessageIndex.value = messages.value.length - 1;
      if (isRecording.value) {
        await ensureTTSReady();
        flushTTS();
      }
      break;
    case 'update':
      const { output, reasoning, fullReasoning, tps: newTps, numTokens: newNumTokens, state } = e.data;
      tps.value = newTps;
      numTokens.value = newNumTokens;
      const lastMessageIndex = messages.value.length - 1;
      const lastMessage = messages.value[lastMessageIndex];

      if (lastMessage.answerIndex === undefined && state === 'answering') {
        lastMessage.answerIndex = lastMessage.content.length;
      }

      // Try multiple approaches to force Vue reactivity
      const newContent = lastMessage.content + (output || '');
      const newReasoning = fullReasoning || (lastMessage.reasoning + (reasoning || ''));

      // Approach 1: Replace the entire array
      messages.value = [
        ...messages.value.slice(0, lastMessageIndex),
        {
          ...lastMessage,
          content: newContent,
          reasoning: newReasoning
        }
      ];

      // Force reactivity trigger
      triggerRef(messages);

      // Stream to TTS when mic is active
      if (isRecording.value && isTTSReady.value) {
        const delta = output || '';
        if (delta) {
          await speakDelta(delta);
        }
      }
      // Test: Force a DOM update by changing a simple value
      numTokens.value = newNumTokens;
      break;
    case 'complete':
      isRunning.value = false;
      if (isRecording.value && isTTSReady.value) enqueueSilence(120);
      if (currentChatId.value && messages.value.length > 0) {
        const lastMessage = messages.value[messages.value.length - 1];
        if (lastMessage.role === 'assistant' && lastMessage.content) {
          try {
            await chatStorage.saveMessage(currentChatId.value, lastMessage);
            await updateChatsList();
          } catch (error) {
            console.error('Failed to save assistant message:', error);
          }
        }
      }
      break;
    case 'error':
      isRunning.value = false;

      // Replace the last assistant message (the placeholder) with the error message
      const errorMessageIndex = messages.value.length - 1;
      if (errorMessageIndex >= 0 && messages.value[errorMessageIndex].role === 'assistant') {
        // Handle structured error from API worker
        if (typeof e.data.data === 'object' && e.data.data.message) {
          const errorData = e.data.data;
          const errorMessage = {
            role: 'assistant',
            content: errorData.message,
            isError: true,
            isRetryable: errorData.isRetryable,
            errorType: errorData.errorType,
            errorDetails: errorData.details
          };

          // Replace the placeholder message
          messages.value = [
            ...messages.value.slice(0, errorMessageIndex),
            errorMessage
          ];

          // Update last AI message index to the error message
          lastAiMessageIndex.value = errorMessageIndex;
        } else {
          // Handle simple string errors (backward compatibility)
          const errorMessage = {
            role: 'assistant',
            content: e.data.data || 'An unknown error occurred',
            isError: true,
            isRetryable: true,
            errorType: 'unknown'
          };

          // Replace the placeholder message
          messages.value = [
            ...messages.value.slice(0, errorMessageIndex),
            errorMessage
          ];

          // Update last AI message index to the error message
          lastAiMessageIndex.value = errorMessageIndex;
        }

        // Save the error message to IndexedDB
        if (currentChatId.value) {
          try {
            await chatStorage.saveMessage(currentChatId.value, messages.value[errorMessageIndex]);
            await updateChatsList();
          } catch (error) {
            console.error('Failed to save error message:', error);
          }
        }
      }
      break;
  }
};

const onErrorReceived = (e) => {
  console.error('Worker error:', e);
};

// STT message handlers
const onSTTMessageReceived = (e) => {
  switch (e.data.status) {
    case 'loading':
      sttStatus.value = 'loading';
      break;
    case 'ready':
      sttStatus.value = 'ready';

      // If user manually started recording and we were warming up after a message send, restart recording
      if (userStartedRecording.value && !isRecording.value) {
        startContinuousRecording();
      }
      break;
    case 'start':
      // STT generation started
      isSTTProcessing.value = true;
      break;
    case 'update':
      // Real-time STT text updates (streaming)
      const { tps: sttTps } = e.data;
      // Can show TPS if needed
      break;
    case 'complete':
      // Generation complete - replace text completely
      isSTTProcessing.value = false;
      const output = e.data.output;
      // Extract the text string from the array
      const text = Array.isArray(output) ? output[0] : output || '';
      sttText.value = text;

      // Update input box to show STT text
      if (isRecording.value) {
        input.value = sttText.value;
        resizeInput();

        // Trigger classification on the newly transcribed text
        if (sttText.value.trim()) {
          console.log('Triggering classification for new transcription:', sttText.value);
          tryClassification(sttText.value);
        }
      }
      break;
    case 'error':
      console.error('STT Worker error:', e.data.data);
      isSTTProcessing.value = false;
      break;
  }
};

const onSTTErrorReceived = (e) => {
  console.error('STT Worker error:', e);
};

// STT constants and functions
const WHISPER_SAMPLING_RATE = 16000;
const MAX_AUDIO_LENGTH = 30; // seconds
const MAX_SAMPLES = WHISPER_SAMPLING_RATE * MAX_AUDIO_LENGTH;

const startContinuousRecording = async () => {
  try {
    if (isRecording.value) return;

    // Initialize VAD-web
    vad.value = await MicVAD.new({
      onSpeechStart: () => {
        console.log('🎤 Voice detected - starting utterance recording');
        isVoiceActive.value = true;
        vadGate.value = true;

        // Clear any incomplete timeout
        if (incompleteTimer.value) {
          clearTimeout(incompleteTimer.value);
          incompleteTimer.value = null;
        }

        // Reset for new utterance
        sttText.value = '';
        audioChunks.value = [];

        // Start recording for this utterance
        startWhisperRecording();
      },

      onSpeechEnd: () => {
        console.log('🔇 Voice stopped - processing complete utterance');
        isVoiceActive.value = false;
        vadGate.value = false;

        // Stop recording and process the complete utterance
        stopAndProcessUtterance();
      },

      // Configuration
      preSpeechPadFrames: 70,
      positiveSpeechThreshold: 0.5,
      negativeSpeechThreshold: 0.35,
      redemptionFrames: 8,
    });

    // Start VAD
    vad.value.start();
    isRecording.value = true;
    console.log('VAD-web initialized and started');
    // Ensure TTS is unlocked and ready to speak
    await ensureTTSReady();

  } catch (error) {
    console.error('Error starting VAD recording:', error);
  }
};

const startWhisperRecording = async () => {
  if (mediaRecorder.value && mediaRecorder.value.state === 'recording') return;

  try {
    // Get user media for Whisper (if not already available)
    if (!audioStream.value) {
      audioStream.value = await navigator.mediaDevices.getUserMedia({
        audio: {
          channelCount: 1,
          echoCancellation: false,
          noiseSuppression: true,
          autoGainControl: false
        }
      });
    }

    // Create audio context with Whisper sampling rate (reuse if exists)
    if (!audioContext.value || audioContext.value.state === 'closed') {
      audioContext.value = new AudioContext({
        sampleRate: WHISPER_SAMPLING_RATE,
      });
    }

    // Create recorder with explicit options for better compatibility
    const options = {};
    if (MediaRecorder.isTypeSupported('audio/webm;codecs=opus')) {
      options.mimeType = 'audio/webm;codecs=opus';
    } else if (MediaRecorder.isTypeSupported('audio/webm')) {
      options.mimeType = 'audio/webm';
    } else if (MediaRecorder.isTypeSupported('audio/mp4')) {
      options.mimeType = 'audio/mp4';
    }

    mediaRecorder.value = new MediaRecorder(audioStream.value, options);

    mediaRecorder.value.onstart = () => {
      console.log('Recording utterance started');
    };

    mediaRecorder.value.ondataavailable = (e) => {
      // Collect all chunks for this utterance
      if (e.data && e.data.size > 0) {
        console.log('Audio chunk received:', e.data.size, 'bytes, type:', e.data.type);
        audioChunks.value.push(e.data);
      }
    };

    mediaRecorder.value.onerror = (e) => {
      console.error('MediaRecorder error:', e.error);
    };

    // Start recording without timeslice - get one complete container
    mediaRecorder.value.start();

  } catch (error) {
    console.error('Error starting Whisper recording:', error);
  }
};

const stopAndProcessUtterance = async () => {
  if (!mediaRecorder.value || mediaRecorder.value.state !== 'recording') return;

  // Set up the onstop handler before stopping
  mediaRecorder.value.onstop = async () => {
    try {
      console.log('Processing complete utterance with', audioChunks.value.length, 'chunks');

      if (audioChunks.value.length === 0) {
        console.warn('No audio chunks recorded');
        return;
      }

      const mimeType = mediaRecorder.value.mimeType || audioChunks.value[0]?.type || 'audio/webm;codecs=opus';
      const finalBlob = new Blob(audioChunks.value, { type: mimeType });

      console.log('Final utterance blob:', finalBlob.size, 'bytes, type:', finalBlob.type);

      if (finalBlob.size < 1000) {
        console.warn('Utterance blob too small:', finalBlob.size);
        return;
      }

      const arrayBuffer = await finalBlob.arrayBuffer();
      if (!arrayBuffer.byteLength) {
        console.warn('Empty utterance buffer');
        return;
      }

      const decoded = await audioContext.value.decodeAudioData(arrayBuffer);
      let audio = decoded.getChannelData(0);

      console.log('Decoded utterance:', audio.length, 'samples at', decoded.sampleRate, 'Hz');

      if (audio.length < 1600) {
        console.warn('Utterance too short:', audio.length, 'samples');
        return;
      }

      if (audio.length > MAX_SAMPLES) {
        audio = audio.slice(-MAX_SAMPLES);
      }

      console.log('Sending complete utterance to STT:', audio.length, 'samples');
      sttWorker.value.postMessage({
        type: 'generate',
        data: { audio, language: 'en' }
      });

    } catch (error) {
      console.error('Error processing utterance:', error);
    } finally {
      // Clean up for next utterance
      audioChunks.value = [];
    }
  };

  // Stop the recording
  mediaRecorder.value.stop();
};

const tryClassification = (text) => {
  if (text !== lastClassifiedText.value && !pendingClassification.value) {
    console.log('🔍 Classifying:', text);
    pendingClassification.value = true;
    lastClassifiedText.value = text;

    // Simple rule-based classification for complete sentences
    const isComplete = classifyMessageCompleteness(text);

    setTimeout(() => {
      pendingClassification.value = false;
      if (isComplete && text.trim().length > 2) {
        console.log('✅ Complete - auto-sending:', text);
        // Clear any incomplete timer
        if (incompleteTimer.value) {
          clearTimeout(incompleteTimer.value);
          incompleteTimer.value = null;
        }
        // Auto-send the message
        onEnter(text);
      } else {
        // Handle incomplete message with 3s timeout
        handleIncompleteMessage(text);
      }
    }, 500);
  }
};

// Simple rule-based classifier to determine if a message is complete
const classifyMessageCompleteness = (text) => {
  const trimmed = text.trim().toLowerCase();

  // Skip blank audio or very short texts
  if (trimmed.length < 2 || trimmed === '[blank_audio]' || trimmed === '') {
    return false;
  }

  // Check for sentence-ending punctuation
  if (/[.!?]\s*$/.test(trimmed)) {
    return true;
  }

  // Check for common complete phrases/greetings
  const completePatterns = [
    /^(hello|hi|hey|good morning|good afternoon|good evening)$/,
    /^(thank you|thanks|bye|goodbye|see you)$/,
    /^(yes|no|ok|okay|sure|alright)$/,
    /^(help|test|testing)$/,
    /^what is/,
    /^how do/,
    /^can you/,
    /^please/,
    /^tell me/,
    /^show me/,
    /^explain/,
    /^write/,
    /^create/,
    /^make/,
    /^solve/,
    /^calculate/,
    /^find/
  ];

  for (const pattern of completePatterns) {
    if (pattern.test(trimmed)) {
      return true;
    }
  }

  // Check for questions (even without question marks)
  if (/^(what|where|when|why|how|who|which|is|are|can|will|would|could|should|do|does|did)\b/.test(trimmed)) {
    return true;
  }

  // Check for commands/requests
  if (/^(let me|give me|send|open|close|start|stop|play|pause)\b/.test(trimmed)) {
    return true;
  }

  // If it's a longer phrase (5+ words), likely complete
  if (trimmed.split(/\s+/).length >= 5) {
    return true;
  }

  return false;
};

const handleIncompleteMessage = (text) => {
  console.log('❌ Message incomplete, setting 3s timeout...');
  lastIncompleteText.value = text;

  // Set 3-second timer to bypass classifier if no new speech
  incompleteTimer.value = setTimeout(() => {
    if (!isVoiceActive.value && lastIncompleteText.value === input.value) {
      console.log('⏰ 3s timeout - sending incomplete message anyway:', text);
      onEnter(text);
    }
  }, 3000);
};

const stopContinuousRecording = () => {
  if (vad.value) {
    vad.value.pause();
    vad.value = null;
  }
  if (mediaRecorder.value) {
    mediaRecorder.value.stop();
    mediaRecorder.value = null;
  }
  if (audioStream.value) {
    audioStream.value.getTracks().forEach(track => track.stop());
    audioStream.value = null;
  }
  if (audioContext.value) {
    audioContext.value.close().catch(() => {});
    audioContext.value = null;
  }
  if (incompleteTimer.value) {
    clearTimeout(incompleteTimer.value);
    incompleteTimer.value = null;
  }
  isRecording.value = false;
  isVoiceActive.value = false;
  vadGate.value = false;
};

const toggleRecording = () => {
  if (isRecording.value) {
    userStartedRecording.value = false;
    stopContinuousRecording();
    flushTTS();
  } else {
    userStartedRecording.value = true;
    startContinuousRecording();
  }
};

const initializeWorker = () => {
  if (worker.value) {
    worker.value.removeEventListener('message', onMessageReceived);
    worker.value.removeEventListener('error', onErrorReceived);
    worker.value.terminate();
  }

  // Choose worker based on provider type
  const workerType = effectiveSettings.value.provider === 'local' ? workerUrl : apiWorkerUrl;

  worker.value = new Worker(workerType, {
    type: 'module',
  });

  worker.value.addEventListener('message', onMessageReceived);
  worker.value.addEventListener('error', onErrorReceived);

  // Initialize STT worker if not already done
  if (!sttWorker.value) {
    sttWorker.value = new Worker(sttWorkerUrl, {
      type: 'module',
    });
    sttWorker.value.addEventListener('message', onSTTMessageReceived);
    sttWorker.value.addEventListener('error', onSTTErrorReceived);

    // Load STT model
    sttWorker.value.postMessage({ type: 'load' });
  }

  if (effectiveSettings.value.provider === 'local') {
    worker.value.postMessage({ type: 'check' });
    if (isWebGpuAvailable.value === true) {
      loadModel();
    }
  } else {
    // For API providers, initialize with settings including system prompt (convert to plain object)
    const settingsWithPrompt = {
      ...effectiveSettings.value,
      currentSystemPromptId: currentSystemPromptId.value,
      currentSystemPrompt: getCurrentSystemPrompt.value
    };
    worker.value.postMessage({ type: 'init', data: JSON.parse(JSON.stringify(settingsWithPrompt)) });
  }
};

const loadSettings = () => {
  const stored = localStorage.getItem('ai-chat-settings');
  if (stored) {
    try {
      const parsed = JSON.parse(stored);
      currentSettings.value = { ...currentSettings.value, ...parsed };
    } catch (e) {
      console.warn('Failed to parse stored settings:', e);
    }
  }
};

const onSettingsChanged = (newSettings) => {
  // Update current settings
  currentSettings.value = newSettings;

  // Always reinitialize worker when settings change from Settings modal
  // This ensures the new settings take effect immediately
  status.value = null;
  initializeWorker();
};

const onQuickSettingsChanged = async (newSettings) => {
  // Update per-chat provider settings
  currentChatProvider.value = newSettings.provider;
  currentChatSelectedCustomApi.value = newSettings.selectedCustomApi;

  // Save the provider settings for the current chat
  if (currentChatId.value) {
    try {
      await chatStorage.updateChatProvider(currentChatId.value, newSettings.provider, newSettings.selectedCustomApi);
    } catch (error) {
      console.error('Failed to save provider settings for chat:', error);
    }
  }

  // Reinitialize worker with the new settings
  status.value = null;
  initializeWorker();
};

onMounted(async () => {
  isWebGpuAvailable.value = !!navigator.gpu;
  loadSettings();
  await initializeChatStorage();
  // Don't initialize worker here - it will be initialized when the first chat is selected/created

  // Lazy-preload TTS model and voices; unlock will happen on mic start
  initTTS({ modelUrl: '/tts-model/kitten_tts_nano_v0_1.onnx', voicesJsonUrl: '/tts-model/voices.json' }).catch(() => {});

  // Scroll to bottom after initial load
  await nextTick(() => {
    scrollToBottom();
  });
});

onUnmounted(() => {
  if (worker.value) {
    worker.value.removeEventListener('message', onMessageReceived);
    worker.value.removeEventListener('error', onErrorReceived);
  }
  if (sttWorker.value) {
    sttWorker.value.removeEventListener('message', onSTTMessageReceived);
    sttWorker.value.removeEventListener('error', onSTTErrorReceived);
  }
  // Clean up recording
  stopContinuousRecording();
});

watch(messages, () => {
  // Don't trigger generation when loading history
  if (isLoadingHistory.value) {
    return;
  }

  if (messages.value.filter((x) => x.role === 'user').length === 0) {
    return;
  }

  if (messages.value[messages.value.length - 1].role === 'assistant') {
    return;
  }

  tps.value = null;
  worker.value.postMessage({ type: 'generate', data: JSON.parse(JSON.stringify(messages.value)) });
}, { deep: true });

// Watch for when we finish loading history to scroll to bottom
watch(isLoadingHistory, (newVal, oldVal) => {
  if (oldVal === true && newVal === false) {
    // We just finished loading history, scroll to bottom
    nextTick(() => {
      scrollToBottom();
    });
  }
});

// Watch for messages changes to adjust padding during AI response
watch(messages, () => {
  if (isRunning.value && chatContainerRef.value) {
    // Adjust padding as AI response grows
    adjustBottomPadding();
  }
}, { deep: true });

watch([messages, isRunning], () => {
  if (!chatContainerRef.value || !isRunning.value) return;

  nextTick(() => {
    const element = chatContainerRef.value;

    if (
      element.scrollHeight - element.scrollTop - element.clientHeight <
      STICKY_SCROLL_THRESHOLD
    ) {
      element.scrollTop = element.scrollHeight;
    }
  });
});

// Watch for chat changes to scroll to bottom
watch(currentChatId, () => {
  nextTick(() => {
    setTimeout(() => {
      scrollToBottom();
    }, 500);
  });
});

const loadModel = () => {
  const settingsWithPrompt = {
    ...effectiveSettings.value,
    currentSystemPromptId: currentSystemPromptId.value,
    currentSystemPrompt: getCurrentSystemPrompt.value
  };

  if (effectiveSettings.value.provider === 'local') {
    worker.value.postMessage({
      type: 'load',
      data: effectiveSettings.value.localModel,
      systemPrompt: JSON.parse(JSON.stringify(getCurrentSystemPrompt.value))
    });
  } else {
    worker.value.postMessage({
      type: 'init',
      data: JSON.parse(JSON.stringify(settingsWithPrompt))
    });
  }
  status.value = 'loading';
};

const reset = async () => {
  worker.value.postMessage({ type: 'reset' });
  messages.value = [];
  if (currentChatId.value) {
    await chatStorage.clearMessages(currentChatId.value);
  }
};

const initializeChatStorage = async () => {
  try {
    await chatStorage.init();
    await loadChats();
  } catch (error) {
    console.error('Failed to initialize chat storage:', error);
  }
};

const loadChats = async () => {
  try {
    chats.value = await chatStorage.getAllChats();
    if (chats.value.length === 0) {
      await createNewChat();
    } else {
      // Load the first chat and its settings
      await selectChat(chats.value[0].id);
      // Ensure we scroll to bottom after loading
      await nextTick();
      scrollToBottom();
    }
  } catch (error) {
    console.error('Failed to load chats:', error);
  }
};

const loadCurrentChat = async () => {
  if (!currentChatId.value) return;
  try {
    isLoadingHistory.value = true;
    messages.value = await chatStorage.getMessages(currentChatId.value);
    // Reset last AI message index
    lastAiMessageIndex.value = -1;
    // Find the last AI message if any
    for (let i = messages.value.length - 1; i >= 0; i--) {
      if (messages.value[i].role === 'assistant') {
        lastAiMessageIndex.value = i;
        break;
      }
    }
    isLoadingHistory.value = false;
    await nextTick();
    scrollToBottom();
  } catch (error) {
    console.error('Failed to load chat messages:', error);
    messages.value = [];
    lastAiMessageIndex.value = -1;
    isLoadingHistory.value = false;
  }
};

const scrollToBottom = () => {
  if (chatContainerRef.value) {
    // Use requestAnimationFrame to ensure DOM is fully updated
    requestAnimationFrame(() => {
      // Add a small delay to ensure content is fully rendered
      setTimeout(() => {
        if (chatContainerRef.value) {
          chatContainerRef.value.scrollTop = chatContainerRef.value.scrollHeight;
        }
      }, 10);
    });
  }
};

const scrollToUserMessage = (messageIndex) => {
  if (chatContainerRef.value && chatRef.value && chatRef.value.messageRefs[messageIndex]) {
    const messageElement = chatRef.value.messageRefs[messageIndex].messageRef;
    if (messageElement) {
      // Use requestAnimationFrame to ensure DOM is fully updated
      requestAnimationFrame(() => {
        // Add a small delay to ensure content is fully rendered
        setTimeout(() => {
          if (chatContainerRef.value && messageElement) {
            // Scroll to the user message
            chatContainerRef.value.scrollTop = messageElement.offsetTop;

            // Add padding to the bottom to ensure there's room for the AI response
            adjustBottomPadding();
          }
        }, 10);
      });
    }
  }
};

const adjustBottomPadding = () => {
  // Add padding to the bottom of the chat container to ensure there's room for the AI response
  if (bottomPaddingRef.value) {
    // Calculate the height of the viewport minus some space for the input area
    const viewportHeight = window.innerHeight;
    const inputAreaHeight = 300; // Approximate height of input area and stats
    const desiredPadding = Math.max(viewportHeight - inputAreaHeight, 200);

    // Set the height of the padding element
    bottomPaddingRef.value.style.height = `${desiredPadding}px`;
  }
};

// Reset padding when AI response is complete
watch(isRunning, (newVal) => {
  if (!newVal && bottomPaddingRef.value) {
    // Reset padding when AI response is complete
    bottomPaddingRef.value.style.height = '0px';
  }
});

const createNewChat = async () => {
  try {
    const defaultPromptId = currentSettings.value.defaultSystemPrompt || '';

    const newChat = await chatStorage.createChat('New Chat', defaultPromptId);
    chats.value.unshift(newChat);
    currentChatId.value = newChat.id;
    messages.value = [];
    lastAiMessageIndex.value = -1;
    tps.value = null;
    numTokens.value = null;

    // Set system prompt to default for new chats
    currentSystemPromptId.value = defaultPromptId;

    // Apply default provider and model settings for new chats
    if (currentSettings.value.defaultCustomApi !== null) {
      // Use default custom API if one is set
      currentChatProvider.value = 'custom';
      currentChatSelectedCustomApi.value = currentSettings.value.defaultCustomApi;
    } else {
      // Fall back to local model
      currentChatProvider.value = 'local';
      currentChatSelectedCustomApi.value = null;
    }

    // Initialize worker with default settings
    initializeWorker();

    // Scroll to bottom for new chat
    nextTick(() => {
      scrollToBottom();
    });
  } catch (error) {
    console.error('Failed to create new chat:', error);
  }
};

const selectChat = async (chatId) => {
  if (chatId === currentChatId.value) return;
  currentChatId.value = chatId;

  // Load the chat's settings (system prompt and provider/model)
  try {
    const chat = await chatStorage.getChat(chatId);
    if (chat) {
      // Load system prompt
      if (chat.systemPromptId !== undefined) {
        currentSystemPromptId.value = chat.systemPromptId;
      } else {
        currentSystemPromptId.value = currentSettings.value.defaultSystemPrompt || '';
      }

      // Load provider settings - use chat-specific if available, otherwise use defaults
      if (chat.provider !== undefined) {
        currentChatProvider.value = chat.provider;
        currentChatSelectedCustomApi.value = chat.selectedCustomApi || null;
      } else {
        // Use default settings for chats created before this feature
        if (currentSettings.value.defaultCustomApi !== null) {
          currentChatProvider.value = 'custom';
          currentChatSelectedCustomApi.value = currentSettings.value.defaultCustomApi;
        } else {
          currentChatProvider.value = 'local';
          currentChatSelectedCustomApi.value = null;
        }
      }
    } else {
      // Fallback to user's defaults
      currentSystemPromptId.value = currentSettings.value.defaultSystemPrompt || '';
      if (currentSettings.value.defaultCustomApi !== null) {
        currentChatProvider.value = 'custom';
        currentChatSelectedCustomApi.value = currentSettings.value.defaultCustomApi;
      } else {
        currentChatProvider.value = 'local';
        currentChatSelectedCustomApi.value = null;
      }
    }
  } catch (error) {
    console.error('Failed to load chat settings:', error);
    currentSystemPromptId.value = currentSettings.value.defaultSystemPrompt || '';
    if (currentSettings.value.defaultCustomApi !== null) {
      currentChatProvider.value = 'custom';
      currentChatSelectedCustomApi.value = currentSettings.value.defaultCustomApi;
    } else {
      currentChatProvider.value = 'local';
      currentChatSelectedCustomApi.value = null;
    }
  }

  await loadCurrentChat();
  tps.value = null;
  numTokens.value = null;

  // Reinitialize worker with chat-specific settings
  initializeWorker();

  // Scroll to bottom after selecting chat
  nextTick(() => {
    scrollToBottom();
  });
};

const deleteChat = async (chatId) => {
  if (chats.value.length <= 1) return;

  try {
    await chatStorage.deleteChat(chatId);
    chats.value = chats.value.filter(chat => chat.id !== chatId);

    if (currentChatId.value === chatId) {
      currentChatId.value = chats.value[0].id;
      await loadCurrentChat();
    }
  } catch (error) {
    console.error('Failed to delete chat:', error);
  }
};

const renameChat = async (chatId, newTitle) => {
  try {
    await chatStorage.updateChat(chatId, { title: newTitle });
    const chatIndex = chats.value.findIndex(chat => chat.id === chatId);
    if (chatIndex !== -1) {
      chats.value[chatIndex].title = newTitle;
    }
  } catch (error) {
    console.error('Failed to rename chat:', error);
  }
};

const toggleSidebar = () => {
  sidebarCollapsed.value = !sidebarCollapsed.value;
};

const updateChatsList = async () => {
  try {
    chats.value = await chatStorage.getAllChats();
  } catch (error) {
    console.error('Failed to update chats list:', error);
  }
};

// Built-in system prompt
const DEFAULT_SYSTEM_PROMPT = {
  id: '',
  name: 'Helpful Assistant',
  content: 'Be a helpful assistant',
  description: 'Generic helpful assistant',
  isBuiltIn: true
};

// System Prompt helpers
const getAllSystemPrompts = computed(() => {
  return [DEFAULT_SYSTEM_PROMPT, ...currentSettings.value.systemPrompts];
});

const getSystemPromptById = (id) => {
  return getAllSystemPrompts.value.find(p => p.id === id);
};

const getCurrentSystemPrompt = computed(() => {
  // First, try to get the prompt by the current chat's selection or the default setting
  const promptId = currentSystemPromptId.value || currentSettings.value.defaultSystemPrompt || '';
  const prompt = getSystemPromptById(promptId);

  if (prompt) {
    return prompt;
  }

  // Fallback to built-in if nothing else works
  return DEFAULT_SYSTEM_PROMPT;
});

const getCurrentModelName = computed(() => {
  if (currentChatProvider.value === 'local') {
    return 'Qwen3 4B (Local)';
  } else if (currentChatProvider.value === 'custom' && currentChatSelectedCustomApi.value !== null) {
    const api = currentSettings.value.customApis[currentChatSelectedCustomApi.value];
    return api ? api.name : 'Custom API';
  }
  return 'No Model';
});

const getCurrentProviderName = computed(() => {
  if (currentChatProvider.value === 'local') {
    return 'Local';
  } else if (currentChatProvider.value === 'custom' && currentChatSelectedCustomApi.value !== null) {
    const api = currentSettings.value.customApis[currentChatSelectedCustomApi.value];
    const providers = currentSettings.value.providers;
    const currentProvider = providers.filter((p) => p.id === api.providerId);

    return currentProvider.length ? currentProvider[0].name : 'No Provider';
  }
  return 'No Model';
});

// Computed property that combines global settings with per-chat overrides
const effectiveSettings = computed(() => {
  return {
    ...currentSettings.value,
    provider: currentChatProvider.value,
    selectedCustomApi: currentChatSelectedCustomApi.value
  };
});

const setSystemPrompt = async (promptId) => {
  currentSystemPromptId.value = promptId;

  // Save the system prompt for the current chat
  if (currentChatId.value) {
    try {
      await chatStorage.updateChatSystemPrompt(currentChatId.value, promptId);
    } catch (error) {
      console.error('Failed to save system prompt for chat:', error);
    }
  }
};

const retryLastMessage = async () => {
  if (isRunning.value || messages.value.length === 0) return;

  // Find the last error message and remove only that one
  const lastErrorIndex = messages.value.length - 1;
  if (lastErrorIndex >= 0 && messages.value[lastErrorIndex].isError) {
    // Remove only the last error message from the UI
    messages.value = messages.value.slice(0, lastErrorIndex);

    // Update the chat in IndexedDB by removing only the last error message
    if (currentChatId.value) {
      try {
        // Clear all messages and re-save the non-error messages
        await chatStorage.clearMessages(currentChatId.value);
        for (const message of messages.value) {
          await chatStorage.saveMessage(currentChatId.value, message);
        }
        await updateChatsList();
      } catch (error) {
        console.error('Failed to update chat during retry:', error);
      }
    }

    // Retry generation with the remaining messages
    tps.value = null;
    isRunning.value = true;
    worker.value.postMessage({ type: 'generate', data: JSON.parse(JSON.stringify(messages.value)) });
  }
};
</script>

<template>
  <div v-if="isWebGpuAvailable || effectiveSettings.provider !== 'local'" class="flex h-screen text-gray-200 bg-gray-900">
    <!-- Sidebar -->
    <ChatTabs
      :chats="chats"
      :currentChatId="currentChatId"
      :isCollapsed="sidebarCollapsed"
      @newChat="createNewChat"
      @selectChat="selectChat"
      @deleteChat="deleteChat"
      @renameChat="renameChat"
      @toggleSidebar="toggleSidebar"
    />

    <!-- Main Content Area -->
    <div class="flex-1 flex flex-col">
      <!-- Loading State -->
      <div v-if="status === 'loading'" class="flex items-center justify-center h-full">
        <div class="w-full max-w-[500px] text-left mx-auto p-6 glass-strong rounded-2xl shadow-2xl relative z-10">
          <p class="text-center mb-4 text-lg font-medium">{{ loadingMessage }}</p>
          <Progress v-for="(item, i) in progressItems" :key="i" :text="item.file" :percentage="item.progress" :total="item.total" />
        </div>
      </div>

      <!-- Chat Interface -->
      <div v-else-if="status === 'ready'" class="flex flex-col h-full">
        <!-- Header with Settings Button -->
        <div class="flex justify-end p-4">
          <button
            @click="settingsOpen = true"
            class="p-3 glass-strong rounded-xl hover:bg-white/10 transition-all duration-300 shadow-lg hover:shadow-xl border border-white/20 hover:border-white/30"
            title="Settings"
          >
            <SettingsIcon class="w-6 h-6 text-gray-300" />
          </button>
        </div>

        <!-- Chat Messages -->
        <div ref="chatContainerRef" class="flex-1 overflow-y-auto scrollbar-thin">
          <div class="max-w-[960px] mx-auto w-full">
            <Chat ref="chatRef" :messages="messages" :lastAiMessageIndex="lastAiMessageIndex" :onSpeak="speakText" :onStopTTS="stopTTS" :isTTSSpeaking="isTTSSpeaking" @retry="retryLastMessage" />
            <div ref="bottomPaddingRef" class="h-0"></div>

            <div v-if="messages.length === 0" class="flex flex-col items-center justify-center p-8">
              <div class="mb-8">
                <div v-for="(msg, i) in EXAMPLES" :key="i" class="m-2 glass-strong rounded-xl p-4 cursor-pointer hover:bg-white/10 transition-all duration-300 shadow-lg hover:shadow-xl hover-lift group border border-white/20 hover:border-white/30" @click="onEnter(msg)">
                  {{ msg }}
                </div>
              </div>
            </div>
          </div>
        </div>

        <!-- Stats -->
        <div class="px-4 py-2">
          <p class="text-center text-sm min-h-6 text-gray-300">
            <template v-if="tps && messages.length > 0">
              <span v-if="!isRunning">
                {{ numTokens }} tokens, {{ (numTokens / tps).toFixed(2) }} seconds
              </span>

              <span class="mr-1">
                <span class="font-medium text-white mr-1">({{ tps.toFixed(2) }}</span>
                <span class="font-normal text-gray-300">tokens/sec).</span>
              </span>

              <template v-if="!isRunning">
                <span class="underline cursor-pointer" @click="reset">Reset</span>
              </template>
            </template>
          </p>
        </div>


        <!-- Input Area -->
        <div class="p-4 pt-0">
          <div class="glass-strong rounded-2xl w-full max-w-[600px] mx-auto relative shadow-2xl hover:shadow-3xl transition-all duration-300 border border-white/20 hover:border-white/30 z-50">
            <!-- Textarea Container -->
            <div class="relative flex">
              <textarea
                id="prompt"
                ref="textareaRef"
                class="scrollbar-thin flex-1 px-4 py-4 bg-transparent border-none outline-none disabled:text-gray-400 text-gray-200 placeholder-gray-400 disabled:placeholder-gray-200 resize-none disabled:cursor-not-allowed"
                :placeholder="isRecording ? 'Listening...' : 'Type your message...'"
                rows="1"
                v-model="input"
                :disabled="status !== 'ready' || isRecording"
                :title="status === 'ready' ? 'Model is ready' : 'Model not loaded yet'"
                :class="{ 'text-green-400': isRecording }"
                @keydown.enter.prevent="handleEnter"
                @input="resizeInput"
                data-gramm="false"
                data-gramm_editor="false"
                data-enable-grammarly="false"
                spellcheck="false"
                style="max-height: 200px;"
              ></textarea>

              <!-- Microphone and Send Button Container -->
              <div class="flex items-end pb-3 pr-3 space-x-2">
                <!-- Microphone button -->
                <div class="cursor-pointer group" @click="toggleRecording" :disabled="sttStatus !== 'ready'">
                  <MicrophoneIcon
                    class="h-8 w-8 p-1 rounded-xl transition-all duration-200 shadow-md hover:shadow-xl"
                    :class="isRecording
                      ? 'bg-red-600 hover:bg-red-700 text-white animate-pulse shadow-red-600/25'
                      : sttStatus === 'ready'
                        ? 'bg-green-600 hover:bg-green-700 text-white shadow-green-600/25'
                        : 'bg-gray-600 text-gray-400 cursor-not-allowed'"
                  />
                </div>

                <!-- Send button -->
                <div v-if="isRunning" class="cursor-pointer group" @click="onInterrupt">
                  <StopIcon class="h-8 w-8 p-1 rounded-xl text-red-500 hover:text-red-600 transition-colors duration-200 hover:bg-red-900/20" />
                </div>
                <div v-else-if="input.length > 0" class="cursor-pointer group" @click="onEnter(input)">
                  <ArrowRightIcon class="h-8 w-8 p-1 bg-blue-600 hover:bg-blue-700 text-white rounded-xl transition-all duration-200 shadow-md hover:shadow-lg" />
                </div>
                <div v-else>
                  <ArrowRightIcon class="h-8 w-8 p-1 bg-gray-600 text-gray-400 rounded-xl" />
                </div>
              </div>
            </div>
            
            <!-- Bottom Controls Bar -->
            <div class="flex items-center justify-between px-4 pb-3 pt-1 border-t border-white/10">
              <div class="flex items-center gap-2">
                <QuickSettings
                  :currentSettings="effectiveSettings"
                  :currentSystemPromptId="currentSystemPromptId"
                  :getAllSystemPrompts="getAllSystemPrompts"
                  :getCurrentSystemPrompt="getCurrentSystemPrompt"
                  :setSystemPrompt="setSystemPrompt"
                  :loadModel="loadModel"
                  @settingsChanged="onQuickSettingsChanged"
                />
              </div>
              <div class="text-xs text-gray-500">
                {{ getCurrentSystemPrompt?.name || 'No prompt' }} • {{ getCurrentModelName }} • {{ getCurrentProviderName }}
              </div>
            </div>
          </div>

          <p class="text-xs text-gray-400 text-center mt-4 glass rounded-full px-4 py-2 border border-white/10 mx-auto max-w-fit relative z-[1]">
            ⚠️ Generated content may be inaccurate or false.
          </p>
        </div>
      </div>
    </div>

    <!-- Settings Modal -->
    <Settings
      :modelOpen="settingsOpen"
      @close="settingsOpen = false"
      @settingsChanged="onSettingsChanged"
    />
  </div>

  <div v-else class="fixed w-screen h-screen bg-gradient-to-br from-red-900 to-gray-900 z-10 text-white text-2xl font-semibold flex flex-col justify-center items-center text-center p-8">
    <div class="bg-white/10 backdrop-blur-md rounded-2xl p-8 border border-white/20">
      <div class="text-4xl mb-4">🚫</div>
      <div class="text-xl mb-2">WebGPU Not Supported</div>
      <div class="text-gray-300 text-base mb-4">Your browser doesn't support WebGPU technology for local models</div>
      <button
        @click="settingsOpen = true"
        class="px-4 py-2 bg-blue-600 hover:bg-blue-700 text-white rounded-lg transition-colors"
      >
        Use API Provider Instead
      </button>
    </div>

    <!-- Settings Modal for fallback case -->
    <Settings
      :modelOpen="settingsOpen"
      @close="settingsOpen = false"
      @settingsChanged="onSettingsChanged"
    />
  </div>
</template>