import { Injectable, Logger } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import axios from 'axios';

export interface StreamChunk {
  type: 'token' | 'metadata' | 'done' | 'error';
  content?: string;
  metadata?: any;
  error?: string;
}

export type StreamCallback = (chunk: StreamChunk) => void;

/**
 * Streaming Service for AI Responses
 *
 * Handles Server-Sent Events (SSE) streaming for real-time AI responses
 * Supports both Azure OpenAI and OpenAI streaming APIs
 */
@Injectable()
export class StreamingService {
  private readonly logger = new Logger(StreamingService.name);
  private aiProvider: 'openai' | 'azure';
  private azureChatEndpoint: string;
  private azureChatDeployment: string;
  private azureChatApiVersion: string;
  private azureChatApiKey: string;

  constructor(private configService: ConfigService) {
    this.aiProvider = this.configService.get('AI_PROVIDER', 'azure') as any;
    this.azureChatEndpoint = this.configService.get('AZURE_OPENAI_CHAT_ENDPOINT');
    this.azureChatDeployment = this.configService.get('AZURE_OPENAI_CHAT_DEPLOYMENT');
    this.azureChatApiVersion = this.configService.get('AZURE_OPENAI_CHAT_API_VERSION');
    this.azureChatApiKey = this.configService.get('AZURE_OPENAI_CHAT_API_KEY');
  }

  /**
   * Stream Azure OpenAI completion
   */
  async streamAzureCompletion(
    messages: Array<{ role: string; content: string }>,
    callback: StreamCallback,
  ): Promise<void> {
    const url = `${this.azureChatEndpoint}/openai/deployments/${this.azureChatDeployment}/chat/completions?api-version=${this.azureChatApiVersion}`;

    const requestBody = {
      messages,
      max_tokens: 1000,
      temperature: 0.7,
      stream: true, // Enable streaming
    };

    try {
      const response = await axios.post(url, requestBody, {
        headers: {
          'Content-Type': 'application/json',
          'api-key': this.azureChatApiKey,
        },
        responseType: 'stream', // Important for streaming
        timeout: 60000,
      });

      let buffer = '';

      // Process the stream
      response.data.on('data', (chunk: Buffer) => {
        buffer += chunk.toString();
        const lines = buffer.split('\n');

        // Keep the last incomplete line in the buffer
        buffer = lines.pop() || '';

        for (const line of lines) {
          const trimmed = line.trim();

          // Skip empty lines and comments
          if (!trimmed || trimmed.startsWith(':')) {
            continue;
          }

          // Parse SSE format
          if (trimmed.startsWith('data: ')) {
            const data = trimmed.substring(6);

            // Check for completion marker
            if (data === '[DONE]') {
              callback({ type: 'done' });
              return;
            }

            try {
              const parsed = JSON.parse(data);

              // Extract the content delta
              if (parsed.choices && parsed.choices[0]?.delta?.content) {
                callback({
                  type: 'token',
                  content: parsed.choices[0].delta.content,
                });
              }

              // Check for finish reason
              if (parsed.choices && parsed.choices[0]?.finish_reason) {
                callback({
                  type: 'metadata',
                  metadata: {
                    finishReason: parsed.choices[0].finish_reason,
                    usage: parsed.usage,
                  },
                });
              }
            } catch (error) {
              this.logger.error('Failed to parse streaming chunk:', error);
            }
          }
        }
      });

      response.data.on('end', () => {
        callback({ type: 'done' });
      });

      response.data.on('error', (error: Error) => {
        callback({
          type: 'error',
          error: error.message,
        });
      });
    } catch (error) {
      this.logger.error('Azure streaming failed:', error);
      callback({
        type: 'error',
        error: error.message || 'Streaming failed',
      });
    }
  }

  /**
   * Stream OpenAI completion
   */
  async streamOpenAICompletion(
    messages: Array<{ role: string; content: string }>,
    callback: StreamCallback,
  ): Promise<void> {
    // TODO: Implement OpenAI streaming
    // For now, fall back to non-streaming
    callback({
      type: 'error',
      error: 'OpenAI streaming not yet implemented',
    });
  }

  /**
   * Main streaming method - routes to appropriate provider
   */
  async streamCompletion(
    messages: Array<{ role: string; content: string }>,
    callback: StreamCallback,
  ): Promise<void> {
    if (this.aiProvider === 'azure') {
      await this.streamAzureCompletion(messages, callback);
    } else {
      await this.streamOpenAICompletion(messages, callback);
    }
  }
}
