import logger from '../../../../debugConfig';

export interface TranscriptionConfig {
  language?: string;
  task?: string;
  modelSize?: string;
  useVad?: boolean;
}

export interface ServerSegment {
  text: string;
  start: number;
  end: number;
}

type ServerSegmentsCallback = (segments: ServerSegment[], isLastLive: boolean) => void;

export class TranscriptionService {
  private socket: WebSocket | null = null;
  private stream: MediaStream | null = null;
  private audioContext: AudioContext | null = null;
  private mediaStreamSource: MediaStreamAudioSourceNode | null = null;
  private workletNode: AudioWorkletNode | null = null;
  private selectedDeviceId: string = '';
  private intentionalStop: boolean = false;
  private onServerSegments: ServerSegmentsCallback | null = null;
  private onDisconnect: (() => void) | null = null;

  constructor(deviceId: string = '') {
    this.selectedDeviceId = deviceId;
  }

  setServerSegmentsCallback(callback: ServerSegmentsCallback) {
    this.onServerSegments = callback;
  }

  setDisconnectCallback(callback: () => void) {
    this.onDisconnect = callback;
  }

  async startTranscription(config: TranscriptionConfig = {}) {
    console.log('🎙️ Starting transcription service...');
    this.intentionalStop = false;

    try {
      logger.info('transcription-service', '🔊 Requesting microphone access...');

      const audioConstraints: MediaTrackConstraints = this.selectedDeviceId
        ? { deviceId: { exact: this.selectedDeviceId } }
        : { echoCancellation: true, noiseSuppression: true };

      this.stream = await navigator.mediaDevices.getUserMedia({ audio: audioConstraints });
      console.log('Got audio stream');

      const uuid = crypto.randomUUID();
      const wsUrl = import.meta.env.VITE_WHISPERLIVE_URL;
      console.log('Connecting to WebSocket:', wsUrl);

      const ws = new WebSocket(wsUrl);
      this.socket = ws;

      const connectionTimeout = setTimeout(() => {
        if (ws.readyState !== WebSocket.OPEN) {
          logger.error('transcription-service', '⏰ WebSocket connection timed out');
          ws.close();
        }
      }, 20000);

      ws.onopen = () => {
        clearTimeout(connectionTimeout);
        logger.info('transcription-service', '✅ WebSocket connected');

        ws.send(JSON.stringify({
          uid: uuid,
          language: config.language || 'en',
          task: config.task || 'transcribe',
          model: config.modelSize || 'large-v3',
          use_vad: config.useVad ?? true,
          max_connection_time: 7200, // server default is 600 s — set to 2 h
        }));
      };

      ws.onerror = (error) => {
        logger.error('transcription-service', '❌ WebSocket error:', error);
      };

      ws.onclose = () => {
        logger.info('transcription-service', '🔌 WebSocket closed');
        const wasIntentional = this.intentionalStop;
        this.cleanup();
        if (!wasIntentional && this.onDisconnect) {
          this.onDisconnect();
        }
      };

      ws.onmessage = (event) => {
        const data = JSON.parse(event.data);
        if (data.uid !== uuid) return;

        if (data.message === 'SERVER_READY') {
          logger.info('transcription-service', '🟢 Server ready, starting audio capture');
          this.setupAudioProcessing();
          return;
        }

        if (data.status === 'WAIT') {
          logger.info('transcription-service', `⏳ Wait time: ${Math.round(data.message)} minutes`);
          this.intentionalStop = true;
          this.cleanup();
          return;
        }

        if (data.message === 'DISCONNECT') {
          logger.info('transcription-service', '🔕 Server requested disconnection');
          this.intentionalStop = true;
          this.cleanup();
          return;
        }

        // Pass the full segment window directly to the store — the store owns
        // all boundary and archival decisions, matching the WhisperLive reference
        // frontend which simply re-renders the server's authoritative segment list.
        if (this.onServerSegments && data.segments && data.segments.length > 0) {
          const segs: ServerSegment[] = data.segments.map((s: any) => ({
            text: String(s.text ?? ''),
            start: parseFloat(s.start ?? 0),
            end: parseFloat(s.end ?? 0),
          }));
          const isLastLive = !(data.segments[data.segments.length - 1]?.completed);
          this.onServerSegments(segs, isLastLive);
        }
      };
    } catch (error) {
      logger.error('transcription-service', '❌ Error starting transcription:', error);
      this.cleanup();
    }
  }

  private async setupAudioProcessing() {
    if (!this.stream || !this.socket) return;

    try {
      this.audioContext = new AudioContext({ sampleRate: 16000 });
      await this.audioContext.audioWorklet.addModule('/audioWorklet.js');

      this.mediaStreamSource = this.audioContext.createMediaStreamSource(this.stream);
      this.workletNode = new AudioWorkletNode(this.audioContext, 'audio-processor');

      this.workletNode.port.onmessage = (event) => {
        if (this.socket?.readyState === WebSocket.OPEN) {
          this.socket.send(event.data);
        }
      };

      this.mediaStreamSource.connect(this.workletNode);
      this.workletNode.connect(this.audioContext.destination);
    } catch (error) {
      console.error('Error setting up audio processing:', error);
    }
  }

  stopTranscription() {
    this.intentionalStop = true;
    if (this.socket?.readyState === WebSocket.OPEN) {
      this.socket.send('END_OF_AUDIO');
    }
    this.cleanup();
  }

  private cleanup() {
    if (this.workletNode) {
      this.workletNode.disconnect();
      this.workletNode = null;
    }
    if (this.mediaStreamSource) {
      this.mediaStreamSource.disconnect();
      this.mediaStreamSource = null;
    }
    if (this.audioContext) {
      this.audioContext.close();
      this.audioContext = null;
    }
    if (this.stream) {
      this.stream.getTracks().forEach(track => track.stop());
      this.stream = null;
    }
    if (this.socket) {
      this.socket.close();
      this.socket = null;
    }
  }
}
