flare-ui/src/app/components/chat/chat.component.ts

import { Component, OnInit, OnDestroy, ViewChild, ElementRef, AfterViewChecked } from '@angular/core';
import { CommonModule } from '@angular/common';
import { FormBuilder, ReactiveFormsModule, Validators, FormsModule } from '@angular/forms';
import { MatButtonModule } from '@angular/material/button';
import { MatIconModule } from '@angular/material/icon';
import { MatFormFieldModule } from '@angular/material/form-field';
import { MatInputModule } from '@angular/material/input';
import { MatCardModule } from '@angular/material/card';
import { MatSelectModule } from '@angular/material/select';
import { MatDividerModule } from '@angular/material/divider';
import { MatTooltipModule } from '@angular/material/tooltip';
import { MatProgressSpinnerModule } from '@angular/material/progress-spinner';
import { MatCheckboxModule } from '@angular/material/checkbox';
import { Subscription } from 'rxjs';

import { ApiService } from '../../services/api.service';
import { EnvironmentService } from '../../services/environment.service';

interface ChatMessage {
  author: 'user' | 'assistant';
  text: string;
  timestamp?: Date;
  audioUrl?: string;  // For TTS audio
}

@Component({
  selector: 'app-chat',
  standalone: true,
  imports: [
    CommonModule,
    FormsModule,
    ReactiveFormsModule,
    MatButtonModule,
    MatIconModule,
    MatFormFieldModule,
    MatInputModule,
    MatCardModule,
    MatSelectModule,
    MatDividerModule,
    MatTooltipModule,
    MatProgressSpinnerModule,
    MatCheckboxModule
  ],
  templateUrl: './chat.component.html',
  styleUrls: ['./chat.component.scss']
})
export class ChatComponent implements OnInit, OnDestroy, AfterViewChecked {
  @ViewChild('scrollMe') private myScrollContainer!: ElementRef;
  @ViewChild('audioPlayer') private audioPlayer!: ElementRef<HTMLAudioElement>;
  @ViewChild('waveformCanvas') private waveformCanvas!: ElementRef<HTMLCanvasElement>;

  projects: string[] = [];
  selectedProject: string | null = null;
  useTTS = false;
  ttsAvailable = false;

  sessionId: string | null = null;
  messages: ChatMessage[] = [];
  input = this.fb.control('', Validators.required);
  
  loading = false;
  error = '';
  playingAudio = false;
  useSTT = false;
  sttAvailable = false;
  isListening = false;
  
  // Audio visualization
  audioContext?: AudioContext;
  analyser?: AnalyserNode;
  animationId?: number;
  
  private subs = new Subscription();
  private shouldScroll = false;

  constructor(
    private fb: FormBuilder, 
    private api: ApiService,
    private environmentService: EnvironmentService
  ) {}

  ngOnInit(): void {
    this.loadProjects();
    this.checkTTSAvailability();
    this.checkSTTAvailability();
    
    // Initialize Audio Context
    this.audioContext = new (window.AudioContext || (window as any).webkitAudioContext)();
  }

  ngAfterViewChecked() {
    if (this.shouldScroll) {
      this.scrollToBottom();
      this.shouldScroll = false;
    }
  }

  ngOnDestroy(): void {
    this.subs.unsubscribe();
    if (this.animationId) {
      cancelAnimationFrame(this.animationId);
    }
    if (this.audioContext) {
      this.audioContext.close();
    }
  }

  loadProjects(): void {
    this.loading = true;
    const sub = this.api.getChatProjects().subscribe({
      next: projects => {
        this.projects = projects;
        this.loading = false;
        if (projects.length === 0) {
          this.error = 'No enabled projects found. Please enable a project with published version.';
        }
      },
      error: (err) => {
        this.error = 'Failed to load projects';
        this.loading = false;
        console.error('Load projects error:', err);
      }
    });
    this.subs.add(sub);
  }

  checkTTSAvailability(): void {
    const sub = this.environmentService.environment$.subscribe(env => {
      if (env) {
        this.ttsAvailable = env.tts_engine !== 'no_tts';
        if (!this.ttsAvailable) {
          this.useTTS = false;
        }
      }
    });
    this.subs.add(sub);
    
    // Also get current environment
    this.api.getEnvironment().subscribe({
      next: (env) => {
        this.ttsAvailable = env.tts_engine !== 'no_tts';
        if (!this.ttsAvailable) {
          this.useTTS = false;
        }
      }
    });
  }

  checkSTTAvailability(): void {
    const sub = this.environmentService.environment$.subscribe(env => {
      if (env) {
        this.sttAvailable = env.stt_engine !== 'no_stt';
        if (!this.sttAvailable) {
          this.useSTT = false;
        }
      }
    });
    this.subs.add(sub);
    
    // Also get current environment
    this.api.getEnvironment().subscribe({
      next: (env) => {
        this.sttAvailable = env.stt_engine !== 'no_stt';
        if (!this.sttAvailable) {
          this.useSTT = false;
        }
      }
    });
  }
  
  startChat(): void {
    if (!this.selectedProject) return;
    
    this.loading = true;
    this.error = '';
    
    const sub = this.api.startChat(this.selectedProject).subscribe({
      next: res => {
        this.sessionId = res.session_id;
        const message: ChatMessage = { 
          author: 'assistant', 
          text: res.answer,
          timestamp: new Date()
        };
        
        this.messages = [message];
        this.loading = false;
        this.shouldScroll = true;
        
        // Generate TTS if enabled
        if (this.useTTS) {
          this.generateTTS(res.answer, this.messages.length - 1);
        }
      },
      error: (err) => {
        this.error = err.error?.detail || 'Failed to start session';
        this.loading = false;
        console.error('Start chat error:', err);
      }
    });
    this.subs.add(sub);
  }

  send(): void {
    if (!this.sessionId || this.input.invalid) return;
    
    const text = this.input.value!.trim();
    if (!text) return;

    // Add user message
    this.messages.push({ 
      author: 'user', 
      text,
      timestamp: new Date()
    });
    
    this.input.reset();
    this.loading = true;
    this.shouldScroll = true;

    // Send to backend
    const sub = this.api.chat(this.sessionId, text).subscribe({
      next: res => {
        const message: ChatMessage = { 
          author: 'assistant', 
          text: res.answer,
          timestamp: new Date()
        };
        
        this.messages.push(message);
        this.loading = false;
        this.shouldScroll = true;
        
        // Generate TTS if enabled
        if (this.useTTS) {
          this.generateTTS(res.answer, this.messages.length - 1);
        }
      },
      error: (err) => {
        this.messages.push({
          author: 'assistant',
          text: '⚠️ ' + (err.error?.detail || 'Failed to send message. Please try again.'),
          timestamp: new Date()
        });
        this.loading = false;
        this.shouldScroll = true;
        console.error('Chat error:', err);
      }
    });
    this.subs.add(sub);
  }

  generateTTS(text: string, messageIndex: number): void {
    const sub = this.api.generateTTS(text).subscribe({
      next: (audioBlob) => {
        const audioUrl = URL.createObjectURL(audioBlob);
        this.messages[messageIndex].audioUrl = audioUrl;
        
        // Auto-play the latest message
        if (messageIndex === this.messages.length - 1) {
          setTimeout(() => this.playAudio(audioUrl), 100);
        }
      },
      error: (err) => {
        console.error('TTS generation error:', err);
      }
    });
    this.subs.add(sub);
  }

  playAudio(audioUrl: string): void {
    if (!this.audioPlayer) return;
    
    const audio = this.audioPlayer.nativeElement;
    audio.src = audioUrl;
    
    // Set up audio visualization
    this.setupAudioVisualization(audio);
    
    audio.play().then(() => {
      this.playingAudio = true;
    }).catch(err => {
      console.error('Audio play error:', err);
    });
    
    audio.onended = () => {
      this.playingAudio = false;
      if (this.animationId) {
        cancelAnimationFrame(this.animationId);
        this.clearWaveform();
      }
    };
  }

  setupAudioVisualization(audio: HTMLAudioElement): void {
    if (!this.audioContext || !this.waveformCanvas) return;
    
    // Create audio source and analyser
    const source = this.audioContext.createMediaElementSource(audio);
    this.analyser = this.audioContext.createAnalyser();
    this.analyser.fftSize = 256;
    
    // Connect nodes
    source.connect(this.analyser);
    this.analyser.connect(this.audioContext.destination);
    
    // Start visualization
    this.drawWaveform();
  }

  drawWaveform(): void {
    if (!this.analyser || !this.waveformCanvas) return;
    
    const canvas = this.waveformCanvas.nativeElement;
    const ctx = canvas.getContext('2d');
    if (!ctx) return;
    
    const bufferLength = this.analyser.frequencyBinCount;
    const dataArray = new Uint8Array(bufferLength);
    
    const draw = () => {
      this.animationId = requestAnimationFrame(draw);
      
      this.analyser!.getByteFrequencyData(dataArray);
      
      ctx.fillStyle = 'rgb(240, 240, 240)';
      ctx.fillRect(0, 0, canvas.width, canvas.height);
      
      const barWidth = (canvas.width / bufferLength) * 2.5;
      let barHeight;
      let x = 0;
      
      for (let i = 0; i < bufferLength; i++) {
        barHeight = (dataArray[i] / 255) * canvas.height * 0.8;
        
        ctx.fillStyle = `rgb(63, 81, 181)`;
        ctx.fillRect(x, canvas.height - barHeight, barWidth, barHeight);
        
        x += barWidth + 1;
      }
    };
    
    draw();
  }

  clearWaveform(): void {
    if (!this.waveformCanvas) return;
    
    const canvas = this.waveformCanvas.nativeElement;
    const ctx = canvas.getContext('2d');
    if (!ctx) return;
    
    ctx.fillStyle = 'rgb(240, 240, 240)';
    ctx.fillRect(0, 0, canvas.width, canvas.height);
  }

  endSession(): void {
    this.sessionId = null;
    this.messages = [];
    this.selectedProject = null;
    this.input.reset();
    this.error = '';
    
    // Clean up audio
    if (this.audioPlayer) {
      this.audioPlayer.nativeElement.pause();
    }
    if (this.animationId) {
      cancelAnimationFrame(this.animationId);
    }
    this.clearWaveform();
  }

  private scrollToBottom(): void {
    try {
      if (this.myScrollContainer) {
        this.myScrollContainer.nativeElement.scrollTop = 
          this.myScrollContainer.nativeElement.scrollHeight;
      }
    } catch(err) {
      console.error('Scroll error:', err);
    }
  }
}