/**
 * Google Gemini streaming provider.
 *
 * Hand-written equivalent of the slice of pi-ai/providers/google.ts that bloby
 * needs — text streaming via `:streamGenerateContent?alt=sse`. Function-calling
 * is wired up in Phase 2; for now we drop tools and stream text only.
 *
 * Endpoint: POST {baseUrl}/models/{modelId}:streamGenerateContent?alt=sse&key={apiKey}
 * Stream:   SSE — each `data: {...}` is one candidate update.
 */
import crypto from 'crypto';
import { log } from '../../../../shared/logger.js';
import type {
  PiStreamRequest,
  PiStreamEvent,
  PiMessage,
  PiContentBlock,
  PiStopReason,
  PiUsage,
} from './types.js';
import { fetchWithRetry, readWithIdleTimeout } from './retry.js';
import { classifyPiError, classifyPiNetworkError } from './humanize-error.js';

/** Walk an SSE byte stream and yield each parsed JSON event. */
async function* parseSse(res: Response, dbg: { firstBytes: string }): AsyncIterable<any> {
  if (!res.body) return;
  const reader = res.body.getReader();
  const decoder = new TextDecoder();
  let buffer = '';
  let totalBytes = 0;
  try {
    while (true) {
      const { value, done } = await readWithIdleTimeout(reader, 'Google Gemini');
      if (done) break;
      if (value) totalBytes += value.byteLength;
      buffer += decoder.decode(value, { stream: true });
      if (!dbg.firstBytes && buffer.length > 0) {
        dbg.firstBytes = buffer.slice(0, 800);
      }
      // SSE event boundary is a blank line. Accept both LF and CRLF separators.
      let idx;
      while (
        (idx = (() => {
          const a = buffer.indexOf('\n\n');
          const b = buffer.indexOf('\r\n\r\n');
          if (a < 0) return b;
          if (b < 0) return a;
          return Math.min(a, b);
        })()) !== -1
      ) {
        const isCrlf = buffer.slice(idx, idx + 4) === '\r\n\r\n';
        const raw = buffer.slice(0, idx);
        buffer = buffer.slice(idx + (isCrlf ? 4 : 2));
        const parsed = parseSseEvent(raw);
        if (parsed !== undefined) yield parsed;
      }
    }
    // Flush whatever remains — Gemini's final event may not have a trailing blank line.
    buffer += decoder.decode();
    if (buffer.trim()) {
      const parsed = parseSseEvent(buffer);
      if (parsed !== undefined) yield parsed;
    }
  } finally {
    try { reader.releaseLock(); } catch {}
    dbg.firstBytes = dbg.firstBytes || `(zero bytes — total=${totalBytes})`;
  }
}

function parseSseEvent(raw: string): any | undefined {
  // Standard SSE: one or more `data:` lines per event. Concatenate their payloads.
  const lines = raw.split(/\r?\n/);
  const dataLines = lines
    .filter((l) => l.startsWith('data:'))
    .map((l) => l.slice(5).trimStart());
  if (!dataLines.length) {
    // Fallback: some servers omit the `data:` prefix and send pure JSON per event.
    const trimmed = raw.trim();
    if (!trimmed || trimmed === '[DONE]') return undefined;
    // Strip a leading JSON-array delimiter if Gemini is returning array-stream
    // instead of SSE (alt=sse not honored).
    const candidate = trimmed.replace(/^[\[,]/, '').replace(/[\],]$/, '').trim();
    if (!candidate) return undefined;
    try { return JSON.parse(candidate); } catch { return undefined; }
  }
  const data = dataLines.join('\n');
  if (!data || data === '[DONE]') return undefined;
  try { return JSON.parse(data); } catch { return undefined; }
}

function toGeminiRole(role: PiMessage['role']): 'user' | 'model' {
  // Tool results piggyback on the user role with a `functionResponse` part —
  // see Gemini function-calling docs.
  if (role === 'assistant') return 'model';
  return 'user';
}

function toGeminiParts(content: PiContentBlock[]): any[] {
  const parts: any[] = [];
  for (const b of content) {
    if (b.type === 'text') {
      parts.push({ text: b.text });
    } else if (b.type === 'image') {
      parts.push({ inlineData: { mimeType: b.mediaType, data: b.data } });
    } else if (b.type === 'document') {
      // Gemini ingests application/pdf inline via the same inlineData shape as
      // images (it OCRs/renders the document). buildUserMessage only routes a
      // document block here when the flavor supports it.
      parts.push({ inlineData: { mimeType: b.mediaType, data: b.data } });
    } else if (b.type === 'tool_use') {
      // Assistant turn: the model asked to invoke a tool. Thinking-capable
      // Gemini 3.x rejects (HTTP 400) any echoed functionCall whose
      // thoughtSignature is missing, so we forward it verbatim when present.
      const part: any = { functionCall: { name: b.name, args: b.input || {} } };
      if (b.thoughtSignature) part.thoughtSignature = b.thoughtSignature;
      parts.push(part);
    } else if (b.type === 'tool_result') {
      // Function responses can be strings, objects, or even error markers.
      // Wrap text in `{ output: ... }` (Gemini's docs use a flexible
      // `response` JSON map), with `isError` keying so the model can react.
      const response = b.isError ? { error: b.content } : { output: b.content };
      parts.push({ functionResponse: { name: extractToolName(b.toolUseId), response } });
    }
  }
  return parts;
}

/**
 * Gemini doesn't carry a tool-call id forward to the response; we encode the
 * tool name into the id we generate at tool-use time (`{name}::{uuid}`) so
 * we can recover it here. Falls back to the raw id if the prefix is missing.
 */
function extractToolName(toolUseId: string): string {
  const idx = toolUseId.indexOf('::');
  return idx > 0 ? toolUseId.slice(0, idx) : toolUseId;
}

function toGeminiTools(tools: { name: string; description: string; inputSchema: Record<string, any> }[]) {
  return [{
    functionDeclarations: tools.map((t) => ({
      name: t.name,
      description: t.description,
      // Gemini accepts plain JSON Schema for `parameters`.
      parameters: t.inputSchema,
    })),
  }];
}

function mapStopReason(reason?: string): PiStopReason {
  switch (reason) {
    case 'STOP':
    case 'FINISH_REASON_STOP':
    case undefined:
      return 'end_turn';
    case 'MAX_TOKENS':
      return 'max_tokens';
    case 'SAFETY':
    case 'RECITATION':
    case 'BLOCKLIST':
    case 'PROHIBITED_CONTENT':
    case 'SPII':
    case 'OTHER':
    case 'MALFORMED_FUNCTION_CALL':
      return 'error';
    default:
      return 'end_turn';
  }
}

function finishReasonMessage(reason?: string): string {
  switch (reason) {
    case 'MAX_TOKENS':
      return 'Response cut off — the model hit its output-token budget before finishing.';
    case 'SAFETY':
      return 'Response blocked by Gemini safety filters.';
    case 'RECITATION':
      return 'Response blocked by recitation policy.';
    case 'BLOCKLIST':
    case 'PROHIBITED_CONTENT':
    case 'SPII':
      return `Response blocked by Gemini policy (${reason}).`;
    case 'MALFORMED_FUNCTION_CALL':
      return 'Gemini emitted a malformed function call. Often means the model tried to invoke a tool that wasn\'t declared, or with arguments that failed schema validation.';
    case 'OTHER':
    default:
      return `Gemini stopped without producing output (finishReason=${reason || 'unknown'}).`;
  }
}

export async function* streamGoogle(req: PiStreamRequest): AsyncIterable<PiStreamEvent> {
  const url =
    `${req.baseUrl.replace(/\/+$/, '')}/models/${encodeURIComponent(req.modelId)}:streamGenerateContent` +
    `?alt=sse&key=${encodeURIComponent(req.apiKey)}`;

  // Filter out empty messages — Gemini rejects requests with no user content.
  const contents = req.messages
    .filter((m) => m.content.length > 0)
    .map((m) => ({ role: toGeminiRole(m.role), parts: toGeminiParts(m.content) }))
    .filter((m) => m.parts.length > 0);

  // Default to a generous cap because thinking-capable Gemini models (2.5+,
  // 3.x) consume `maxOutputTokens` for both reasoning AND final text — a small
  // cap silently truncates the answer to nothing. Pi's catalog lists 65 536
  // as the model's hard ceiling.
  const body: any = {
    contents,
    generationConfig: {
      maxOutputTokens: req.maxOutputTokens ?? 32768,
    },
  };
  // Thinking-capable families (2.5+/3.x): ask for thought summaries so the
  // harness can emit a liveness pulse — without this, Gemini 3 burns its
  // output budget on invisible reasoning and the chat looks hung. Gated by
  // model id; unknown/dynamic ids skip it (older models reject the field).
  // The rolling aliases (gemini-flash-latest / gemini-flash-lite-latest)
  // resolve to 2.5+/3.x thinking models too (review PI-D-2).
  if (/gemini-(2\.5|[3-9]|flash(-lite)?-latest)/i.test(req.modelId)) {
    body.generationConfig.thinkingConfig = { includeThoughts: true };
  }
  if (req.systemPrompt?.trim()) {
    body.systemInstruction = { parts: [{ text: req.systemPrompt }] };
  }
  if (req.tools && req.tools.length > 0) {
    body.tools = toGeminiTools(req.tools);
    // Round-cap wrap-up: forbid further function calls; tools stay declared so
    // functionCall/functionResponse parts in history remain valid.
    if (req.toolChoice === 'none') {
      body.toolConfig = { functionCallingConfig: { mode: 'NONE' } };
    }
  }

  let res: Response;
  try {
    res = await fetchWithRetry(url, {
      method: 'POST',
      headers: { 'content-type': 'application/json' },
      body: JSON.stringify(body),
      signal: req.signal,
    });
  } catch (err: any) {
    if (err?.name === 'AbortError') {
      yield { type: 'done', stopReason: 'aborted' };
      return;
    }
    const cls = classifyPiNetworkError('Google Gemini', err);
    yield { type: 'error', error: cls.message, kind: cls.kind, retryable: cls.retryable };
    return;
  }

  if (!res.ok) {
    let detail = '';
    try { detail = await res.text(); } catch {}
    const cls = classifyPiError('Google Gemini', res.status, res.statusText, detail);
    yield { type: 'error', error: cls.message, status: cls.status, kind: cls.kind, retryable: cls.retryable };
    return;
  }

  let accumulated = '';
  let toolCallCount = 0;
  let lastFinish: string | undefined;
  let promptBlockReason: string | undefined;
  let usage: PiUsage | undefined;
  // Debug counters — drop once this stabilises.
  let chunkCount = 0;
  let thoughtPartCount = 0;
  let emptyTextPartCount = 0;
  let firstChunkSummary = '';
  const dbg = { firstBytes: '' };

  try {
    for await (const chunk of parseSse(res, dbg)) {
      chunkCount++;
      if (chunkCount === 1) {
        try { firstChunkSummary = JSON.stringify(chunk).slice(0, 600); } catch {}
      }
      // The whole prompt can be rejected before we even get a candidate.
      if (chunk?.promptFeedback?.blockReason) {
        promptBlockReason = chunk.promptFeedback.blockReason;
      }
      const candidate = chunk?.candidates?.[0];
      const parts: any[] = candidate?.content?.parts || [];
      for (const part of parts) {
        // Thinking models emit reasoning parts with `thought: true`. They
        // shouldn't be shown to the user as part of the visible answer.
        if (part?.thought) {
          thoughtPartCount++;
          if (thoughtPartCount === 1) yield { type: 'thinking' };
          continue;
        }
        if (part?.functionCall && typeof part.functionCall.name === 'string') {
          // Gemini doesn't surface a tool-call id of its own; bake the tool
          // name into the id so the session can echo it back as a
          // `functionResponse` referencing the same name.
          const id = `${part.functionCall.name}::${crypto.randomUUID()}`;
          toolCallCount++;
          yield {
            type: 'tool_use',
            id,
            name: part.functionCall.name,
            input: part.functionCall.args || {},
            // Thinking-capable models attach a signature that we must echo
            // back unchanged on the next turn. Optional on non-thinking models.
            thoughtSignature: typeof part.thoughtSignature === 'string' ? part.thoughtSignature : undefined,
          };
          continue;
        }
        if (typeof part?.text === 'string' && part.text.length > 0) {
          accumulated += part.text;
          yield { type: 'text_delta', delta: part.text };
        } else {
          emptyTextPartCount++;
        }
      }
      if (candidate?.finishReason) lastFinish = candidate.finishReason;
      const usageMeta = chunk?.usageMetadata;
      if (usageMeta) {
        usage = {
          inputTokens: usageMeta.promptTokenCount,
          outputTokens: usageMeta.candidatesTokenCount,
        };
      }
    }
  } catch (err: any) {
    if (err?.name === 'AbortError') {
      yield { type: 'done', stopReason: 'aborted' };
      return;
    }
    const cls = classifyPiNetworkError('Google Gemini', err);
    yield { type: 'error', error: cls.message, kind: cls.kind, retryable: cls.retryable };
    return;
  }

  log.info(
    `[pi/google] stream done — chunks=${chunkCount} text=${accumulated.length} toolCalls=${toolCallCount} ` +
    `thoughtParts=${thoughtPartCount} emptyTextParts=${emptyTextPartCount} ` +
    `finishReason=${lastFinish || 'none'} ` +
    `promptTok=${usage?.inputTokens ?? '?'} outTok=${usage?.outputTokens ?? '?'}`,
  );
  if (chunkCount > 0 && !accumulated && toolCallCount === 0) {
    log.info(`[pi/google] first chunk (truncated): ${firstChunkSummary}`);
  } else if (chunkCount === 0) {
    log.warn(`[pi/google] SSE stream parsed zero chunks — content-type=${res.headers.get('content-type') || '?'}`);
    log.warn(`[pi/google] first raw bytes: ${JSON.stringify(dbg.firstBytes)}`);
  }

  // Prompt-level block: nothing came back at all.
  if (promptBlockReason) {
    yield { type: 'error', error: `Gemini blocked the prompt (${promptBlockReason}).` };
    yield { type: 'done', stopReason: 'error', usage };
    return;
  }

  // Tool-only round (Gemini fires functionCall parts with no text) is valid output —
  // the session will execute the tool, push the result, and re-stream.
  if (!accumulated && toolCallCount === 0) {
    const reason = lastFinish && lastFinish !== 'STOP' && lastFinish !== 'FINISH_REASON_STOP'
      ? lastFinish
      : undefined;
    const hint = thoughtPartCount > 0 && !lastFinish
      ? ' (model emitted thinking but never the final answer — try a non-thinking model like gemini-2.5-flash, or raise maxOutputTokens)'
      : '';
    yield { type: 'error', error: finishReasonMessage(reason) + hint };
    yield { type: 'done', stopReason: mapStopReason(lastFinish), usage };
    return;
  }

  if (accumulated) yield { type: 'text_end', text: accumulated };
  yield {
    type: 'done',
    stopReason: toolCallCount > 0 ? 'tool_use' : mapStopReason(lastFinish),
    usage,
  };
}