import axios from 'axios';

/**
 * Generation parameter interface.
 */
export interface GenerateImageParams {
  prompt: string;
  nsfw_checker?: boolean;
  callBackUrl?: string;
}

/**
 * AI image generation helper class using the GPT Image-2 API.
 */
export class ImageGenerator {
  private readonly apiKey: string;
  private readonly createTaskUrl: string;
  private readonly queryTaskUrl: string;
  private readonly defaultModel: string;

  constructor(apiKey: string, model?: string) {
    this.apiKey = apiKey;
    this.createTaskUrl = 'https://api.kie.ai/api/v1/jobs/createTask';
    this.queryTaskUrl = 'https://api.kie.ai/api/v1/jobs/recordInfo';
    this.defaultModel = model || 'gpt-image-2-text-to-image';
  }

  /**
   * Create an image generation task.
   * @param params Generation parameters.
   */
  async createTask(params: GenerateImageParams): Promise<string> {
    try {
      console.error(`Creating image generation task...`);
      console.error(`Prompt: ${params.prompt}`);

      const requestData: {
        model: string;
        callBackUrl?: string;
        input: {
          prompt: string;
          nsfw_checker: boolean;
        };
      } = {
        model: this.defaultModel,
        input: {
          prompt: params.prompt,
          nsfw_checker: params.nsfw_checker ?? false
        }
      };

      if (params.callBackUrl) {
        requestData.callBackUrl = params.callBackUrl;
      }

      const response = await axios.post(this.createTaskUrl, requestData, {
        headers: {
          'Authorization': `Bearer ${this.apiKey}`,
          'Content-Type': 'application/json'
        },
        timeout: 30000
      });

      if (response.data.code === 200 && response.data.data?.taskId) {
        console.error(`Task created successfully, taskId: ${response.data.data.taskId}`);
        return response.data.data.taskId;
      } else {
        throw new Error(`Failed to create task: ${response.data.msg || 'unknown error'}`);
      }
    } catch (error) {
      console.error('Error while creating task:', error);
      if (error && typeof error === 'object' && 'response' in error) {
        const axiosError = error as any;
        console.error('Response status:', axiosError.response?.status);
        console.error('Response data:', axiosError.response?.data);
      }
      throw error;
    }
  }

  /**
   * Query task status.
   * @param taskId Task ID.
   */
  async queryTask(taskId: string): Promise<any> {
    try {
      const response = await axios.get(this.queryTaskUrl, {
        params: { taskId },
        headers: {
          'Authorization': `Bearer ${this.apiKey}`
        },
        timeout: 30000
      });

      if (response.data.code === 200) {
        return response.data.data;
      } else {
        throw new Error(`Failed to query task: ${response.data.msg || 'unknown error'}`);
      }
    } catch (error) {
      console.error('Error while querying task:', error);
      if (error && typeof error === 'object' && 'response' in error) {
        const axiosError = error as any;
        console.error('Response status:', axiosError.response?.status);
        console.error('Response data:', axiosError.response?.data);
      }
      throw error;
    }
  }

  /**
   * Generate an image with the complete flow: create task + poll status, with timeout control.
   * @param params Generation parameters.
   * @param onProgress Progress callback.
   * @param timeoutSeconds Timeout in seconds. Defaults to 2 seconds.
   * @returns Image URL on success, or the task ID when the timeout is reached.
   */
  async generateImageWithTimeout(
    params: GenerateImageParams,
    onProgress?: (message: string, elapsedSeconds: number) => void,
    timeoutSeconds: number = 2
  ): Promise<{ success: true; imageUrl: string } | { success: false; taskId: string; reason: 'timeout' }> {
    const startTime = Date.now();

    // 1. Create the task.
    const taskId = await this.createTask(params);

    // 2. Poll task status until completion or timeout.
    while (true) {
      await new Promise(resolve => setTimeout(resolve, 500)); // Short polling so results or taskId can return quickly within 2 seconds.

      const elapsedMs = Date.now() - startTime;
      const elapsed = Math.floor(elapsedMs / 1000);
      
      // Check for timeout.
      if (elapsedMs >= timeoutSeconds * 1000) {
        console.error(`Task ${taskId} exceeded ${timeoutSeconds} seconds; returning taskId.`);
        return { success: false, taskId, reason: 'timeout' };
      }

      try {
        const taskData = await this.queryTask(taskId);
        
        // Update progress.
        if (onProgress) {
          const statusMsg = taskData.state === 'waiting' 
            ? 'AI task is queued...' 
            : 'AI is generating the image...';
          onProgress(statusMsg, elapsed);
        }

        // Check task state.
        if (taskData.state === 'success') {
          console.error(`Task completed, total elapsed time: ${elapsed}s`);
          
          // Parse the result.
          const resultJson = JSON.parse(taskData.resultJson);
          if (resultJson.resultUrls && resultJson.resultUrls.length > 0) {
            return { success: true, imageUrl: resultJson.resultUrls[0] };
          } else {
            throw new Error('Task completed but did not return an image URL.');
          }
        } else if (taskData.state === 'fail') {
          throw new Error(`Task failed: ${taskData.failMsg || 'unknown error'}`);
        }
        
        // Continue waiting.
      } catch (error) {
        console.error(`Failed to query task status:`, error);
        // Retry query failures, which may be transient network issues.
        // Explicit task failures are thrown above.
        if (error instanceof Error && error.message.includes('Task failed')) {
          throw error;
        }
      }
    }
  }
}