import OpenAI from 'openai';
import { z } from 'zod';
import { estimateUsd, addUsage, ZERO_USAGE } from '../pricing.js';
import { recordCost } from '../costs.js';
import { AiNotConfiguredError, type ChatMessage, type Completion, type CompletionRequest, type ContentPart, type EmbedRequest, type EmbedResult, type ExtractRequest, type ExtractResult, type ModelProvider, type Role, type StreamDelta, type ToolAgentRequest, type Usage } from '../types.js';

export interface OpenAICompatibleOptions {
  id?: string;
  apiKey: string;
  baseURL?: string;
  models?: Partial<Record<Role, string>>;
  embeddingModel?: string;
  embeddingDimensions?: number;
}

const DEFAULT_MODELS: Record<Role, string> = {
  normalize: 'gpt-5-mini',
  classify: 'gpt-5-mini',
  summarize: 'gpt-5-mini',
  resolve: 'gpt-5-mini',
  research: 'gpt-5',
  vision: 'gpt-5',
  embed: 'text-embedding-3-small',
};

function toContent(content: string | ContentPart[]): string | OpenAI.Chat.Completions.ChatCompletionContentPart[] {
  if (typeof content === 'string') return content;
  return content.map((p): OpenAI.Chat.Completions.ChatCompletionContentPart => {
    if (p.type === 'text') return { type: 'text', text: p.text };
    const url = p.image.url ?? `data:${p.image.mediaType};base64,${p.image.data ?? ''}`;
    return { type: 'image_url', image_url: { url } };
  });
}

function toMessages(system: string | undefined, messages: ChatMessage[]): OpenAI.Chat.Completions.ChatCompletionMessageParam[] {
  const out: OpenAI.Chat.Completions.ChatCompletionMessageParam[] = [];
  if (system) out.push({ role: 'system', content: system });
  for (const m of messages) {
    if (m.role === 'user') out.push({ role: 'user', content: toContent(m.content) });
    else out.push({ role: 'assistant', content: typeof m.content === 'string' ? m.content : m.content.filter((p) => p.type === 'text').map((p) => (p as { text: string }).text).join('') });
  }
  return out;
}

function usageOf(u: OpenAI.Completions.CompletionUsage | null | undefined): Usage {
  return {
    inputTokens: (u?.prompt_tokens ?? 0) - (u?.prompt_tokens_details?.cached_tokens ?? 0),
    outputTokens: u?.completion_tokens ?? 0,
    cacheReadTokens: u?.prompt_tokens_details?.cached_tokens ?? 0,
    cacheWriteTokens: 0,
  };
}

/**
 * Provider for any OpenAI-compatible endpoint: OpenAI itself, or the local MacLustr gateway
 * (LOCAL_LLM_BASE_URL). Also the only embeddings provider for now.
 */
export class OpenAICompatibleProvider implements ModelProvider {
  readonly id: string;
  private readonly client: OpenAI;
  private readonly models: Record<Role, string>;
  private readonly embeddingModel: string;
  private readonly embeddingDimensions: number;
  private readonly isLocal: boolean;

  constructor(opts: OpenAICompatibleOptions) {
    if (!opts.apiKey && !opts.baseURL) throw new AiNotConfiguredError('OPENAI_API_KEY or LOCAL_LLM_BASE_URL missing');
    this.id = opts.id ?? (opts.baseURL ? 'openai-compatible' : 'openai');
    this.isLocal = Boolean(opts.baseURL);
    this.client = new OpenAI({ apiKey: opts.apiKey || 'local', baseURL: opts.baseURL, timeout: 5 * 60_000, maxRetries: 2 });
    this.models = { ...DEFAULT_MODELS, ...(opts.models ?? {}) };
    this.embeddingModel = opts.embeddingModel ?? this.models.embed;
    this.embeddingDimensions = opts.embeddingDimensions ?? 1536;
  }

  supports(role: Role): boolean {
    if (role === 'embed') return !this.isLocal || Boolean(this.embeddingModel);
    return true;
  }

  modelFor(role: Role): string {
    return this.models[role];
  }

  private async account(role: Role, model: string, usage: Usage, cost?: CompletionRequest['cost'], units?: number): Promise<number> {
    const usdEst = this.isLocal ? 0 : estimateUsd(model, usage);
    await recordCost({ provider: this.id, model, role, usage, usdEst, context: cost, units });
    return usdEst;
  }

  async complete(role: Role, req: CompletionRequest): Promise<Completion> {
    const model = this.modelFor(role);
    const res = await this.client.chat.completions.create({
      model,
      messages: toMessages(req.system, req.messages),
      max_completion_tokens: req.maxTokens ?? 4096,
      ...(req.temperature !== undefined ? { temperature: req.temperature } : {}),
      ...(req.json ? { response_format: { type: 'json_object' } } : {}),
      ...(req.stopSequences ? { stop: req.stopSequences } : {}),
    });
    const usage = usageOf(res.usage);
    const usdEst = await this.account(role, model, usage, req.cost);
    const choice = res.choices[0];
    return { text: choice?.message?.content ?? '', model, provider: this.id, usage, usdEst, stopReason: choice?.finish_reason ?? null, refusal: choice?.message?.refusal ? { category: null, explanation: choice.message.refusal } : null };
  }

  async *stream(role: Role, req: CompletionRequest): AsyncIterable<StreamDelta> {
    yield* this.runTools(role, { ...req, tools: [], execute: async () => null });
  }

  async *runTools(role: Role, req: ToolAgentRequest): AsyncIterable<StreamDelta> {
    const model = this.modelFor(role);
    const messages = toMessages(req.system, req.messages);
    const tools: OpenAI.Chat.Completions.ChatCompletionTool[] = req.tools.map((t) => ({ type: 'function', function: { name: t.name, description: t.description, parameters: t.inputSchema } }));
    let total: Usage = ZERO_USAGE;
    let stopReason: string | null = null;
    const maxIter = req.maxIterations ?? 8;
    try {
      for (let i = 0; i < maxIter; i++) {
        const stream = await this.client.chat.completions.create(
          { model, messages, max_completion_tokens: req.maxTokens ?? 8192, ...(tools.length ? { tools } : {}), stream: true, stream_options: { include_usage: true } },
          { signal: req.signal },
        );
        const calls = new Map<number, { id: string; name: string; args: string }>();
        let finish: string | null = null;
        for await (const chunk of stream) {
          if (chunk.usage) {
            const u = usageOf(chunk.usage);
            total = addUsage(total, u);
            yield { type: 'usage', usage: u, usdEst: this.isLocal ? 0 : estimateUsd(model, u), model };
          }
          const c = chunk.choices[0];
          if (!c) continue;
          if (c.delta.content) yield { type: 'text', text: c.delta.content };
          for (const tc of c.delta.tool_calls ?? []) {
            const cur = calls.get(tc.index) ?? { id: tc.id ?? '', name: '', args: '' };
            if (tc.id) cur.id = tc.id;
            if (tc.function?.name) cur.name += tc.function.name;
            if (tc.function?.arguments) cur.args += tc.function.arguments;
            calls.set(tc.index, cur);
          }
          if (c.finish_reason) finish = c.finish_reason;
        }
        stopReason = finish;
        if (finish !== 'tool_calls' || calls.size === 0) break;
        const toolCalls = [...calls.values()];
        messages.push({ role: 'assistant', content: null, tool_calls: toolCalls.map((tc) => ({ id: tc.id, type: 'function', function: { name: tc.name, arguments: tc.args } })) });
        for (const tc of toolCalls) {
          let input: unknown = {};
          try {
            input = tc.args ? JSON.parse(tc.args) : {};
          } catch {
            input = {};
          }
          yield { type: 'tool_call', id: tc.id, name: tc.name, input };
          const started = Date.now();
          let output: unknown;
          let isError = false;
          try {
            output = await req.execute(tc.name, input);
          } catch (err) {
            isError = true;
            output = { error: err instanceof Error ? err.message : String(err) };
          }
          yield { type: 'tool_result', id: tc.id, name: tc.name, output, isError, durationMs: Date.now() - started };
          messages.push({ role: 'tool', tool_call_id: tc.id, content: typeof output === 'string' ? output : JSON.stringify(output ?? null) });
        }
      }
    } finally {
      await this.account(role, model, total, req.cost);
    }
    yield { type: 'done', stopReason };
  }

  async extract<T>(role: Role, req: ExtractRequest<T>): Promise<ExtractResult<T>> {
    const model = this.modelFor(role);
    const schema = z.toJSONSchema(req.schema as z.ZodType, { target: 'draft-2020-12', io: 'output' }) as Record<string, unknown>;
    const content = typeof req.input === 'string' ? `${req.prompt}\n\n${req.input}` : [{ type: 'text', text: req.prompt } as ContentPart, ...req.input];
    const res = await this.client.chat.completions.create({
      model,
      messages: toMessages(req.system, [{ role: 'user', content }]),
      max_completion_tokens: req.maxTokens ?? 4096,
      response_format: { type: 'json_schema', json_schema: { name: 'extraction', schema, strict: false } },
    });
    const usage = usageOf(res.usage);
    const usdEst = await this.account(role, model, usage, req.cost);
    const text = res.choices[0]?.message?.content ?? '';
    const parsed = req.schema.safeParse(JSON.parse(text));
    if (!parsed.success) throw new Error(`extract: model output did not match schema: ${parsed.error.message}`);
    const data = parsed.data as T & { confidence?: number };
    return { data: parsed.data, confidence: typeof data.confidence === 'number' ? Math.max(0, Math.min(1, data.confidence)) : 1, model, provider: this.id, usage, usdEst };
  }

  vision<T>(req: ExtractRequest<T>): Promise<ExtractResult<T>> {
    return this.extract('vision', req);
  }

  async embed(req: EmbedRequest): Promise<EmbedResult> {
    if (!this.embeddingModel) throw new AiNotConfiguredError('no embedding model configured');
    const res = await this.client.embeddings.create({ model: this.embeddingModel, input: req.texts, ...(this.embeddingModel.startsWith('text-embedding-3') ? { dimensions: this.embeddingDimensions } : {}) });
    const usage: Usage = { ...ZERO_USAGE, inputTokens: res.usage?.prompt_tokens ?? 0 };
    const usdEst = await this.account('embed', this.embeddingModel, usage, req.cost, req.texts.length);
    const vectors = res.data.sort((a, b) => a.index - b.index).map((d) => d.embedding);
    return { vectors, model: this.embeddingModel, provider: this.id, dimensions: vectors[0]?.length ?? this.embeddingDimensions, usdEst };
  }
}