import OpenAI from 'openai';
import {
  AIProvider,
  ProviderOptions,
  ProviderResult,
  StreamChunk,
  ProviderConfig,
  ProviderError,
  ProviderUnavailableError,
  ProviderRateLimitError,
  ToolDefinition,
  ToolCall,
} from './provider.js';

export class OpenAIProvider implements AIProvider {
  readonly name = 'OpenAI';
  readonly type = 'openai' as const;

  private client: OpenAI;
  private config: ProviderConfig;
  private defaultModel: string;

  constructor(config: ProviderConfig = {}) {
    this.config = config;
    this.defaultModel = config.model || process.env.OPENAI_MODEL || 'gpt-4o';

    this.client = new OpenAI({
      apiKey: config.apiKey || process.env.OPENAI_API_KEY,
      baseURL: config.baseUrl,
      organization: config.organizationId || process.env.OPENAI_ORG_ID,
      timeout: config.timeout || 120000,
      maxRetries: config.maxRetries || 2,
    });
  }

  async execute(
    prompt: string,
    options?: ProviderOptions,
  ): Promise<ProviderResult> {
    if (!this.isAvailable()) {
      throw new ProviderUnavailableError(this.name, 'API key not configured');
    }

    try {
      const response = await this.client.chat.completions.create({
        model: options?.model || this.defaultModel,
        max_tokens: options?.maxTokens || 4096,
        temperature: options?.temperature,
        messages: options?.systemPrompt
          ? [
              { role: 'system', content: options.systemPrompt },
              { role: 'user', content: prompt },
            ]
          : [{ role: 'user', content: prompt }],
        tools: options?.tools ? this.convertTools(options.tools) : undefined,
        stop: options?.stopSequences,
      });

      return this.parseResponse(response);
    } catch (error) {
      throw this.handleError(error);
    }
  }

  async *stream(
    prompt: string,
    options?: ProviderOptions,
  ): AsyncIterable<StreamChunk> {
    if (!this.isAvailable()) {
      throw new ProviderUnavailableError(this.name, 'API key not configured');
    }

    try {
      const stream = await this.client.chat.completions.create({
        model: options?.model || this.defaultModel,
        max_tokens: options?.maxTokens || 4096,
        temperature: options?.temperature,
        messages: options?.systemPrompt
          ? [
              { role: 'system', content: options.systemPrompt },
              { role: 'user', content: prompt },
            ]
          : [{ role: 'user', content: prompt }],
        tools: options?.tools ? this.convertTools(options.tools) : undefined,
        stop: options?.stopSequences,
        stream: true,
      });

      let inputTokens = 0;
      let outputTokens = 0;

      for await (const chunk of stream) {
        const delta = chunk.choices[0]?.delta;

        if (delta?.content) {
          yield { type: 'text', delta: delta.content };
        }

        if (delta?.tool_calls) {
          for (const toolCall of delta.tool_calls) {
            if (toolCall.function?.name) {
              yield {
                type: 'tool_use',
                toolCall: {
                  id: toolCall.id || '',
                  name: toolCall.function.name,
                  arguments: JSON.parse(toolCall.function.arguments || '{}'),
                },
              };
            }
          }
        }

        if (chunk.usage) {
          inputTokens = chunk.usage.prompt_tokens;
          outputTokens = chunk.usage.completion_tokens;
        }
      }

      yield { type: 'usage', usage: { inputTokens, outputTokens } };
    } catch (error) {
      throw this.handleError(error);
    }
  }

  isAvailable(): boolean {
    return !!(this.config.apiKey || process.env.OPENAI_API_KEY);
  }

  async getModels(): Promise<string[]> {
    return [
      'gpt-4o',
      'gpt-4o-mini',
      'gpt-4-turbo',
      'gpt-4',
      'gpt-3.5-turbo',
      'o1-preview',
      'o1-mini',
    ];
  }

  getDefaultModel(): string {
    return this.defaultModel;
  }

  private convertTools(
    tools: ToolDefinition[],
  ): OpenAI.Chat.ChatCompletionTool[] {
    return tools.map((tool) => ({
      type: 'function' as const,
      function: {
        name: tool.name,
        description: tool.description,
        parameters: tool.parameters,
      },
    }));
  }

  private parseResponse(response: OpenAI.Chat.ChatCompletion): ProviderResult {
    const choice = response.choices[0];
    let content = '';
    const toolCalls: ToolCall[] = [];

    if (choice.message.content) {
      content = choice.message.content;
    }

    if (choice.message.tool_calls) {
      for (const toolCall of choice.message.tool_calls) {
        toolCalls.push({
          id: toolCall.id,
          name: toolCall.function.name,
          arguments: JSON.parse(toolCall.function.arguments),
        });
      }
    }

    return {
      content,
      toolCalls: toolCalls.length > 0 ? toolCalls : undefined,
      usage: response.usage
        ? {
            inputTokens: response.usage.prompt_tokens,
            outputTokens: response.usage.completion_tokens,
          }
        : undefined,
      stopReason: this.mapStopReason(choice.finish_reason),
    };
  }

  private mapStopReason(reason: string | null): ProviderResult['stopReason'] {
    switch (reason) {
      case 'stop':
        return 'end_turn';
      case 'length':
        return 'max_tokens';
      case 'tool_calls':
        return 'tool_use';
      default:
        return 'end_turn';
    }
  }

  private handleError(error: unknown): Error {
    if (error instanceof OpenAI.APIError) {
      if (error.status === 429) {
        const retryAfter = error.headers?.['retry-after'];
        return new ProviderRateLimitError(
          this.name,
          retryAfter ? parseInt(retryAfter, 10) : undefined,
        );
      }
      return new ProviderError(
        error.message,
        this.name,
        error.code,
        error.status,
      );
    }
    if (error instanceof Error) {
      return new ProviderError(error.message, this.name);
    }
    return new ProviderError('Unknown error', this.name);
  }
}
