flynn/src/models/openai.ts

import OpenAI from 'openai';
import type { ChatRequest, ChatResponse, ModelClient, MessageContentPart, TokenUsage } from './types.js';
import { getMessageTextWithTools } from './media.js';
import { ensureValidOpenAIAuth } from '../auth/openai.js';

export interface OpenAIClientConfig {
  apiKey?: string;
  model: string;
  maxTokens?: number;
  baseURL?: string;
  timeoutMs?: number;
  /** If true, use ChatGPT subscription OAuth via the Codex backend endpoint. */
  useOAuth?: boolean;
}

/**
 * Convert Flynn message content to OpenAI format.
 * OpenAI uses { type: 'text', text } and { type: 'image_url', image_url: { url } } parts.
 */
function toOpenAIContent(content: string | MessageContentPart[]): string | OpenAI.ChatCompletionContentPart[] {
  if (typeof content === 'string') {
    return content;
  }

  return content.map((part): OpenAI.ChatCompletionContentPart => {
    if (part.type === 'text') {
      return { type: 'text', text: part.text };
    }
    if (part.type === 'image') {
      // OpenAI accepts data URIs or regular URLs
      const url = part.source.type === 'base64'
        ? `data:${part.source.media_type};base64,${part.source.data!}`
        : part.source.url!;
      return { type: 'image_url', image_url: { url } };
    }
    if (part.type === 'audio') {
      // OpenAI native audio input via input_audio content part
      // Determine format from MIME type (OpenAI supports: wav, mp3, flac, opus, ogg, webm)
      const formatMap: Record<string, string> = {
        'audio/wav': 'wav',
        'audio/mpeg': 'mp3',
        'audio/mp3': 'mp3',
        'audio/ogg': 'ogg',
        'audio/webm': 'webm',
        'audio/mp4': 'mp4',
        'audio/x-m4a': 'mp4',
      };
      const format = formatMap[part.source.media_type] ?? 'wav';
      return {
        type: 'input_audio',
        input_audio: { data: part.source.data, format },
      } as unknown as OpenAI.ChatCompletionContentPart;
    }
    // Fallback — shouldn't happen
    return { type: 'text', text: JSON.stringify(part) };
  });
}

export class OpenAIClient implements ModelClient {
  private client?: OpenAI;
  private model: string;
  private defaultMaxTokens: number;
  private useOAuth: boolean;
  private baseURL?: string;

  constructor(config: OpenAIClientConfig) {
    const timeoutMs = config.timeoutMs ?? 20_000;
    this.useOAuth = Boolean(config.useOAuth);
    this.baseURL = config.baseURL;

    // OAuth mode uses a different backend (ChatGPT Codex) and a different API shape.
    // Only initialize the OpenAI SDK for API-key providers.
    if (!this.useOAuth) {
      this.client = new OpenAI({
        apiKey: config.apiKey,
        baseURL: config.baseURL,
        timeout: timeoutMs,
        maxRetries: 0,
      });
    }
    this.model = config.model;
    this.defaultMaxTokens = config.maxTokens ?? 4096;
  }

  private async chatViaOAuthCodex(request: ChatRequest): Promise<ChatResponse> {
    const CODEX_API_ENDPOINT = 'https://chatgpt.com/backend-api/codex/responses';

    const auth = await ensureValidOpenAIAuth();

    // Codex endpoint requires:
    // - instructions (non-empty)
    // - input must be a list
    // - store must be false
    // - stream must be true (SSE)
    const instructions = (request.system ?? '').trim() || 'You are helpful.';

    const input = request.messages
      .map((m) => {
        const text = getMessageTextWithTools(m);
        if (!text) {return null;}
        const contentType = m.role === 'assistant' ? 'output_text' : 'input_text';
        return {
          role: m.role,
          content: [{ type: contentType, text }],
        };
      })
      .filter((x): x is NonNullable<typeof x> => Boolean(x));

    const body = {
      model: this.model,
      instructions,
      store: false,
      stream: true,
      input,
      // Intentionally omit max_output_tokens: Codex endpoint rejects it.
      // Also omit tools/tool_choice for now.
    };

    const headers: Record<string, string> = {
      'content-type': 'application/json',
      'authorization': `Bearer ${auth.access_token}`,
      'originator': 'flynn',
      'user-agent': 'flynn/0.1',
      'session_id': `flynn-${Date.now()}`,
    };
    if (auth.account_id) {
      headers['ChatGPT-Account-Id'] = auth.account_id;
    }

    const res = await fetch(CODEX_API_ENDPOINT, {
      method: 'POST',
      headers,
      body: JSON.stringify(body),
    });

    if (!res.ok) {
      const text = await res.text();
      throw new Error(`${res.status} ${res.statusText}${text ? `: ${text}` : ''}`);
    }

    if (!res.body) {
      throw new Error('OpenAI OAuth request failed: missing response body');
    }

    let buffer = '';
    let outputText = '';
    let usage: TokenUsage | undefined;

    const reader = res.body.getReader();

    const processBlock = (block: string): void => {
      const lines = block.split('\n');
      let data = '';
      for (const line of lines) {
        if (line.startsWith('data:')) {
          data += line.slice('data:'.length).trim();
        }
      }
      if (!data) {return;}
      let obj: any;
      try {
        obj = JSON.parse(data);
      } catch {
        return;
      }

      if (obj.type === 'response.output_text.delta' && typeof obj.delta === 'string') {
        outputText += obj.delta;
      }

      if (obj.type === 'response.completed') {
        const u = obj.response?.usage;
        if (u) {
          usage = {
            inputTokens: u.input_tokens ?? 0,
            outputTokens: u.output_tokens ?? 0,
          };
        }
      }

      if (obj.type === 'response.failed') {
        const detail = obj.response?.error?.message ?? 'OpenAI OAuth response failed';
        throw new Error(detail);
      }
    };

    while (true) {
      const { value, done } = await reader.read();
      if (done) {break;}
      buffer += Buffer.from(value).toString('utf8');

      while (true) {
        const idx = buffer.indexOf('\n\n');
        if (idx === -1) {break;}
        const block = buffer.slice(0, idx);
        buffer = buffer.slice(idx + 2);
        processBlock(block);
      }
    }

    return {
      content: outputText,
      stopReason: 'end_turn',
      usage: usage ?? { inputTokens: 0, outputTokens: 0 },
    };
  }

  async chat(request: ChatRequest): Promise<ChatResponse> {
    if (this.useOAuth) {
      return this.chatViaOAuthCodex(request);
    }

    if (!this.client) {
      throw new Error('OpenAI client not initialized');
    }

    const messages: OpenAI.ChatCompletionMessageParam[] = [];

    if (request.system) {
      messages.push({ role: 'system', content: request.system });
    }

    for (const msg of request.messages) {
      messages.push({
        role: msg.role,
        content: toOpenAIContent(msg.content),
      } as OpenAI.ChatCompletionMessageParam);
    }

    // Build params, conditionally including tools
    const params: OpenAI.ChatCompletionCreateParamsNonStreaming = {
      model: this.model,
      max_tokens: request.maxTokens ?? this.defaultMaxTokens,
      messages,
    };

    if (request.tools && request.tools.length > 0) {
      params.tools = request.tools.map(t => ({
        type: 'function' as const,
        function: {
          name: t.name,
          description: t.description,
          parameters: t.input_schema as OpenAI.FunctionParameters,
        },
      }));
    }

    // Extended thinking/reasoning mode for o1/o3 models
    if (request.thinking) {
      (params as any).reasoning_effort = 'medium';
    }

    let response: OpenAI.ChatCompletion;
    try {
      response = await this.client.chat.completions.create(params);
    } catch (error) {
      const status = typeof (error as { status?: unknown })?.status === 'number'
        ? (error as { status: number }).status
        : undefined;
      const message = error instanceof Error ? error.message : String(error);

      const isZai = (this.baseURL ?? '').includes('api.z.ai');
      const missingModelRequestScope = message.includes('Missing scopes: model.request');

      if (isZai && status === 401) {
        const hint = missingModelRequestScope
          ? 'The key lacks `model.request` scope.'
          : 'The API key is invalid, expired, or not allowed for this model/endpoint.';
        throw new Error(
          `Z.AI authentication failed (401). ${hint} ` +
          'Run `flynn zai-auth` to update credentials, or set ZAI_API_KEY / ZHIPUAI_API_KEY / ZHIPUAI_AUTH_TOKEN.',
        );
      }

      throw error;
    }

    const choice = response.choices[0];
    const content = choice?.message?.content ?? '';

    // Parse tool_calls from the response if present
    const toolCalls = choice?.message?.tool_calls?.map((tc: OpenAI.ChatCompletionMessageToolCall) => ({
      id: tc.id,
      name: tc.function.name,
      args: JSON.parse(tc.function.arguments),
    })) ?? [];

    // Map OpenAI finish reasons to Flynn's stop reasons
    let stopReason: string;
    if (toolCalls.length > 0) {
      stopReason = 'tool_use';
    } else {
      const reason = choice?.finish_reason;
      if (reason === 'stop') {
        stopReason = 'end_turn';
      } else if (reason === 'length') {
        stopReason = 'max_tokens';
      } else if (reason === 'tool_calls') {
        // Edge case: finish_reason says tool_calls but none were parsed
        stopReason = 'end_turn';
      } else {
        stopReason = reason ?? 'end_turn';
      }
    }

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