import { startSpan } from '@/remote/monitor';
import { GoogleGenAI, type Part } from '@google/genai';
import { fetch as expoFetch } from 'expo/fetch';
import type { AIProvider } from './types';

const MODEL = 'gemini-3-flash-preview';
const API_KEY = process.env.EXPO_PUBLIC_GEMINI_API_KEY || '';
const STREAM_URL = `https://generativelanguage.googleapis.com/v1beta/models/${MODEL}:streamGenerateContent?alt=sse&key=${API_KEY}`;

const ai = new GoogleGenAI({ apiKey: API_KEY });

function buildParts(prompt: string, base64?: string): Part[] {
  const parts: Part[] = [{ text: prompt }];
  if (base64) {
    parts.push({ inlineData: { mimeType: 'image/jpeg', data: base64 } });
  }
  return parts;
}

export class GeminiProvider implements AIProvider {
  async query(prompt: string, base64?: string): Promise<string> {
    return startSpan({ name: 'ai.gemini', op: 'ai.inference' }, async span => {
      const start = Date.now();
      const result = await ai.models.generateContent({
        model: MODEL,
        contents: [{ role: 'user', parts: buildParts(prompt, base64) }],
        config: { thinkingConfig: { thinkingBudget: 0 } },
      });

      const usage = result.usageMetadata;
      span.setAttributes({
        'ai.model': MODEL,
        'ai.input_tokens': usage?.promptTokenCount ?? 0,
        'ai.output_tokens': usage?.candidatesTokenCount ?? 0,
        'ai.latency_ms': Date.now() - start,
        'ai.has_image': !!base64,
      });

      return result.text ?? '';
    });
  }

  async queryStream(
    prompt: string,
    callbacks: { onChunk: (accumulated: string) => void; signal?: AbortSignal },
    base64?: string
  ): Promise<string> {
    return startSpan({ name: 'ai.gemini.stream', op: 'ai.inference' }, async span => {
      const start = Date.now();

      const body = {
        contents: [{ role: 'user', parts: buildParts(prompt, base64) }],
        generationConfig: { thinkingConfig: { thinkingBudget: 0 } },
      };

      const response = await expoFetch(STREAM_URL, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(body),
      });

      if (!response.ok) {
        const text = await response.text();
        throw new Error(`Gemini stream error ${response.status}: ${text}`);
      }

      const reader = response.body?.getReader();
      if (!reader) throw new Error('No readable stream from response');

      const decoder = new TextDecoder();
      let accumulated = '';
      let buffer = '';

      const processLine = (line: string) => {
        if (!line.startsWith('data:')) return;
        const json = line.slice(5).trim();
        if (!json) return;
        try {
          const parsed = JSON.parse(json);
          const text = parsed.candidates?.[0]?.content?.parts?.[0]?.text;
          if (text) {
            accumulated += text;
            callbacks.onChunk(accumulated);
          }
        } catch {
          /* ignore */
        }
      };

      const drainBuffer = () => {
        const lines = buffer.split(/\r?\n/);
        buffer = lines.pop() ?? '';
        for (const line of lines) processLine(line);
      };

      try {
        while (true) {
          if (callbacks.signal?.aborted) break;
          const { done, value } = await reader.read();
          if (done) break;
          buffer += typeof value === 'string' ? value : decoder.decode(value, { stream: true });
          drainBuffer();
        }
        if (buffer.trim()) processLine(buffer);
      } finally {
        reader.releaseLock();
      }

      span.setAttributes({
        'ai.model': MODEL,
        'ai.latency_ms': Date.now() - start,
        'ai.has_image': !!base64,
        'ai.streaming': true,
      });

      return accumulated;
    });
  }
}