import type { ContentBlock, Message } from '../types';
import type { LlmAdapter, LlmCallbacks, LlmRequest, LlmTurnResult, StopReason } from './types';

// Minimal OpenAI-compatible Chat Completions types — covers what we use.
interface OAIToolCall {
  index?: number;
  id?: string;
  type?: 'function';
  function?: { name?: string; arguments?: string };
}

interface OAITool {
  type: 'function';
  function: {
    name: string;
    description: string;
    parameters: Record<string, unknown>;
  };
}

type OAIMessage =
  | { role: 'system'; content: string }
  | { role: 'user'; content: string }
  | { role: 'assistant'; content: string | null; tool_calls?: { id: string; type: 'function'; function: { name: string; arguments: string } }[] }
  | { role: 'tool'; tool_call_id: string; content: string };

interface OAIChunk {
  choices?: {
    index?: number;
    delta?: {
      role?: string;
      content?: string | null;
      tool_calls?: OAIToolCall[];
    };
    finish_reason?: string | null;
  }[];
}

function defaultBase(baseUrl: string): string {
  return (baseUrl || 'https://api.openai.com/v1').replace(/\/+$/, '');
}

function historyToApi(history: Message[]): OAIMessage[] {
  const out: OAIMessage[] = [];
  for (const m of history) {
    if (m.role === 'user') {
      // A user message in our model can be plain text OR a vehicle for tool_result blocks.
      const toolResults = m.blocks.filter((b): b is Extract<ContentBlock, { type: 'tool_result' }> => b.type === 'tool_result');
      if (toolResults.length > 0) {
        for (const r of toolResults) {
          out.push({ role: 'tool', tool_call_id: r.tool_use_id, content: r.content });
        }
      } else {
        const text = m.blocks
          .filter((b): b is Extract<ContentBlock, { type: 'text' }> => b.type === 'text')
          .map((b) => b.text)
          .join('\n');
        if (text.length > 0) out.push({ role: 'user', content: text });
      }
    } else {
      // assistant
      const text = m.blocks
        .filter((b): b is Extract<ContentBlock, { type: 'text' }> => b.type === 'text')
        .map((b) => b.text)
        .join('');
      const tcs = m.blocks
        .filter((b): b is Extract<ContentBlock, { type: 'tool_use' }> => b.type === 'tool_use')
        .map((b) => ({
          id: b.id,
          type: 'function' as const,
          function: { name: b.name, arguments: JSON.stringify(b.input ?? {}) },
        }));
      if (text.length === 0 && tcs.length === 0) continue;
      out.push({
        role: 'assistant',
        content: text.length > 0 ? text : null,
        ...(tcs.length > 0 ? { tool_calls: tcs } : {}),
      });
    }
  }
  return out;
}

function mapStop(r: string | null | undefined): StopReason {
  switch (r) {
    case 'stop':
      return 'end_turn';
    case 'tool_calls':
    case 'function_call':
      return 'tool_use';
    case 'length':
      return 'max_tokens';
    case 'content_filter':
      return 'refusal';
    default:
      return 'other';
  }
}

export const openaiAdapter: LlmAdapter = {
  async stream(req: LlmRequest, cb: LlmCallbacks): Promise<LlmTurnResult> {
    const url = defaultBase(req.baseUrl) + '/chat/completions';

    const tools: OAITool[] = req.tools.map((t) => ({
      type: 'function',
      function: {
        name: t.name,
        description: t.description,
        parameters: t.input_schema,
      },
    }));

    const body: Record<string, unknown> = {
      model: req.model,
      stream: true,
      max_tokens: req.maxTokens,
      messages: [{ role: 'system', content: req.system }, ...historyToApi(req.history)],
    };
    if (tools.length > 0) {
      body.tools = tools;
    }

    const resp = await fetch(url, {
      method: 'POST',
      headers: {
        'content-type': 'application/json',
        ...(req.apiKey ? { Authorization: `Bearer ${req.apiKey}` } : {}),
      },
      body: JSON.stringify(body),
    });

    if (!resp.ok || !resp.body) {
      const text = await resp.text().catch(() => '');
      throw new Error(`LLM ${resp.status}: ${text.slice(0, 400)}`);
    }

    const reader = resp.body.getReader();
    const decoder = new TextDecoder();

    let buffer = '';
    let accumText = '';
    // Tool call accumulators keyed by index (OpenAI streams partial JSON for arguments).
    const toolByIndex = new Map<number, { id: string; name: string; argsRaw: string }>();
    let finishReason: string | null | undefined;

    while (true) {
      const { value, done } = await reader.read();
      if (done) break;
      buffer += decoder.decode(value, { stream: true });

      // SSE events are separated by blank lines; each event has one or more lines starting with "data: ".
      const events = buffer.split('\n\n');
      buffer = events.pop() ?? '';

      for (const rawEvent of events) {
        for (const line of rawEvent.split('\n')) {
          if (!line.startsWith('data:')) continue;
          const payload = line.slice(5).trim();
          if (!payload || payload === '[DONE]') continue;
          let chunk: OAIChunk;
          try {
            chunk = JSON.parse(payload) as OAIChunk;
          } catch {
            continue;
          }
          const choice = chunk.choices?.[0];
          if (!choice) continue;
          if (choice.finish_reason) finishReason = choice.finish_reason;
          const delta = choice.delta;
          if (!delta) continue;
          if (typeof delta.content === 'string' && delta.content.length > 0) {
            accumText += delta.content;
            cb.onTextDelta(delta.content);
          }
          if (delta.tool_calls) {
            for (const tc of delta.tool_calls) {
              const idx = tc.index ?? 0;
              let agg = toolByIndex.get(idx);
              if (!agg) {
                agg = { id: tc.id ?? '', name: tc.function?.name ?? '', argsRaw: '' };
                toolByIndex.set(idx, agg);
              }
              if (tc.id && !agg.id) agg.id = tc.id;
              if (tc.function?.name && !agg.name) agg.name = tc.function.name;
              if (tc.function?.arguments) agg.argsRaw += tc.function.arguments;
            }
          }
        }
      }
    }

    const content: ContentBlock[] = [];
    if (accumText.length > 0) content.push({ type: 'text', text: accumText });

    // Sort tool calls by index for stable ordering.
    const tcKeys = Array.from(toolByIndex.keys()).sort((a, b) => a - b);
    for (const idx of tcKeys) {
      const agg = toolByIndex.get(idx);
      if (!agg) continue;
      let input: unknown = {};
      if (agg.argsRaw.length > 0) {
        try {
          input = JSON.parse(agg.argsRaw);
        } catch {
          input = { _raw: agg.argsRaw };
        }
      }
      const id = agg.id || `call_${idx}`;
      cb.onToolUse({ id, name: agg.name, input });
      content.push({ type: 'tool_use', id, name: agg.name, input });
    }

    return { stopReason: mapStop(finishReason), content };
  },
};
