import { ALL_NUTRITION_KEYS, USA_GOV_DGA_KEYS } from '@/common/constants/nutrition';
import { strings } from '@/common/strings';
import { startSpan } from '@/remote/monitor';
import type { BoundingBox, CandidateResult, DetectedFoodItem } from '@/types/foodlog';
import { z } from 'zod';
import { GeminiProvider } from './providers';

const provider = new GeminiProvider();

const CandidateResultSchema = z.object({
  candidates: z.array(
    z.object({
      foodName: z.string(),
      hpPredIngredients: z.array(z.string()),
      hrPredIngredients: z.array(z.string()),
    })
  ),
});

const IngredientsResultSchema = z.object({
  hpPredIngredients: z.array(z.string()),
  hrPredIngredients: z.array(z.string()),
});

const MultipleFoodItemsSchema = z.object({
  items: z.array(
    z.object({
      foodName: z.string(),
      hpPredIngredients: z.array(z.string()),
      hrPredIngredients: z.array(z.string()),
      boundingBox: z
        .object({
          x: z.number(),
          y: z.number(),
          width: z.number(),
          height: z.number(),
        })
        .optional(),
      confidence: z.number(),
    })
  ),
});

const NutritionResultSchema = z.object({
  novaClass: z.number().optional(),
  nutrition: z.record(z.string(), z.number()).optional(),
});

import { AIResStreamParser } from './aiResStreamParser';

function normalizeBBox(box: BoundingBox | undefined): BoundingBox | undefined {
  if (!box) return undefined;
  if (box.x > 1 || box.y > 1 || box.width > 1 || box.height > 1) {
    return { x: box.x / 1000, y: box.y / 1000, width: box.width / 1000, height: box.height / 1000 };
  }
  return box;
}

function extractJSON(text: string): string {
  const parser = new AIResStreamParser();
  const parsed = parser.parse(text);
  return parsed ? JSON.stringify(parsed) : text;
}

export function extractStreamingItems(text: string): StreamingUpdate {
  const completedItems: StreamingFoodItem[] = [];
  let partialItem: Partial<StreamingFoodItem> | undefined;

  const itemsIdx = text.indexOf('"items"');
  if (itemsIdx === -1) return { completedItems };
  const bracketIdx = text.indexOf('[', itemsIdx);
  if (bracketIdx === -1) return { completedItems };

  let i = bracketIdx + 1;
  let objStart = -1;
  let depth = 0;
  let inStr = false;
  let esc = false;

  while (i < text.length) {
    const ch = text[i];
    if (esc) {
      esc = false;
      i++;
      continue;
    }
    if (ch === '\\' && inStr) {
      esc = true;
      i++;
      continue;
    }
    if (ch === '"') {
      inStr = !inStr;
      i++;
      continue;
    }
    if (inStr) {
      i++;
      continue;
    }
    if (ch === '{') {
      if (depth === 0) objStart = i;
      depth++;
    } else if (ch === '}') {
      depth--;
      if (depth === 0 && objStart !== -1) {
        try {
          const obj = JSON.parse(text.substring(objStart, i + 1));
          if (obj.foodName) {
            completedItems.push({
              foodName: obj.foodName,
              hpPredIngredients: obj.hpPredIngredients ?? [],
              hrPredIngredients: obj.hrPredIngredients ?? [],
              boundingBox: normalizeBBox(obj.boundingBox),
              confidence: obj.confidence ?? 0.5,
            });
          }
        } catch {
          /* malformed object */
        }
        objStart = -1;
      }
    }
    i++;
  }

  if (depth > 0 && objStart !== -1) {
    const tail = text.substring(objStart);
    const nameMatch = tail.match(/"foodName"\s*:\s*"([^"]*)"/);
    if (nameMatch) {
      partialItem = { foodName: nameMatch[1] };
      const hpMatch = tail.match(/"hpPredIngredients"\s*:\s*\[([^\]]*)\]/);
      if (hpMatch)
        try {
          partialItem.hpPredIngredients = JSON.parse(`[${hpMatch[1]}]`);
        } catch {
          /* partial */
        }
      const hrMatch = tail.match(/"hrPredIngredients"\s*:\s*\[([^\]]*)\]/);
      if (hrMatch)
        try {
          partialItem.hrPredIngredients = JSON.parse(`[${hrMatch[1]}]`);
        } catch {
          /* partial */
        }
    }
  }

  return { completedItems, partialItem };
}

export type StreamingFoodItem = {
  foodName: string;
  hpPredIngredients: string[];
  hrPredIngredients: string[];
  boundingBox?: BoundingBox;
  confidence: number;
};

export type StreamingUpdate = {
  completedItems: StreamingFoodItem[];
  partialItem?: Partial<StreamingFoodItem>;
};

export async function identifyFoodStreaming(input: {
  base64: string;
  onUpdate: (update: StreamingUpdate) => void;
  signal?: AbortSignal;
}): Promise<StreamingFoodItem[]> {
  return startSpan(
    {
      name: 'ai_provider_gemini_stream',
      op: 'ai.inference',
      attributes: { provider: 'gemini', mode: 'stream' },
    },
    async () => {
      const callbacks = {
        onChunk: accumulated => {
          const result = extractStreamingItems(accumulated);
          if (result.completedItems.length > 0 || result.partialItem?.foodName) {
            input.onUpdate(result);
          }
        },
        signal: input.signal,
      };

      const finalText = await provider.queryStream(
        strings.ai.identifyFoodFromPhoto,
        callbacks,
        input.base64
      );
      if (!finalText.trim()) return [];

      let parsed: { items?: StreamingFoodItem[] };
      try {
        parsed = JSON.parse(extractJSON(finalText));
      } catch {
        console.warn('Failed to parse streaming response:', finalText.slice(0, 200));
        return [];
      }

      return (parsed.items ?? [])
        .filter((item: Partial<StreamingFoodItem>) => item.foodName)
        .map((item: StreamingFoodItem) => ({
          foodName: item.foodName,
          hpPredIngredients: item.hpPredIngredients ?? [],
          hrPredIngredients: item.hrPredIngredients ?? [],
          boundingBox: normalizeBBox(item.boundingBox),
          confidence: item.confidence ?? 0.5,
        }));
    }
  );
}

async function identifyWithGemini(input: {
  base64?: string;
  text?: string;
  foodName?: string;
}): Promise<CandidateResult> {
  return startSpan(
    {
      name: 'ai_provider_gemini',
      op: 'ai.inference',
      attributes: { provider: 'gemini' },
    },
    async () => {
      const { base64, text, foodName } = input;

      let prompt: string;
      if (foodName) {
        prompt = strings.ai.identifyIngredients(foodName, !!base64);
      } else {
        prompt =
          strings.ai.identifyItemFromCrop +
          '\n\n' +
          (text || strings.addLog.detection.identifyThisFood);
      }

      const responseText = await provider.query(prompt, base64);
      const parsed = JSON.parse(extractJSON(responseText));

      if (foodName) {
        const ingredients = IngredientsResultSchema.parse(parsed);
        return {
          candidates: [
            {
              foodName,
              hpPredIngredients: ingredients.hpPredIngredients,
              hrPredIngredients: ingredients.hrPredIngredients,
            },
          ],
        };
      }

      const parsedCandidates = CandidateResultSchema.parse(parsed);
      return {
        candidates: parsedCandidates.candidates.map(candidate => ({
          foodName: candidate.foodName,
          hpPredIngredients: candidate.hpPredIngredients,
          hrPredIngredients: candidate.hrPredIngredients,
        })),
      };
    }
  );
}

export async function identifyFood(input: {
  base64?: string;
  text?: string;
  foodName?: string;
}): Promise<{ result: CandidateResult | null; error?: string }> {
  const { base64, text, foodName } = input;

  if (!base64 && !text && !foodName) {
    throw new Error('Must provide base64, text, or foodName');
  }

  try {
    const result = await identifyWithGemini({ base64, text, foodName });
    return { result };
  } catch (e) {
    return { result: null, error: String(e) };
  }
}

export async function estimateNutrition(input: {
  foodName: string;
  ingredients: string[];
}): Promise<{ novaClass?: number; nutrition?: Record<string, number> }> {
  const prompt = strings.ai.estimateNutrition(
    input.foodName,
    input.ingredients,
    ALL_NUTRITION_KEYS.join(', '),
    USA_GOV_DGA_KEYS.join(', ')
  );
  const responseText = await provider.query(prompt);
  const parsed = JSON.parse(extractJSON(responseText));
  return NutritionResultSchema.parse(parsed);
}

export async function identifyMultipleFoods(input: {
  base64: string;
}): Promise<DetectedFoodItem[]> {
  return startSpan(
    {
      name: 'ai_provider_gemini_multi',
      op: 'ai.inference',
      attributes: { provider: 'gemini', mode: 'multi' },
    },
    async () => {
      let responseText: string;
      try {
        responseText = await provider.query(strings.ai.identifyFoodFromPhoto, input.base64);
      } catch (e) {
        console.error('identifyMultipleFoods API error:', e);
        return [];
      }

      try {
        const parsed = JSON.parse(extractJSON(responseText));
        const items = Array.isArray(parsed) ? parsed : parsed.items;
        const parsedItems = MultipleFoodItemsSchema.parse({ items }).items;
        return parsedItems.map(item => ({
          foodName: item.foodName,
          hpPredIngredients: item.hpPredIngredients,
          hrPredIngredients: item.hrPredIngredients,
          boundingBox: normalizeBBox(item.boundingBox),
          confidence: item.confidence,
        }));
      } catch (e) {
        console.error(
          'identifyMultipleFoods parse error:',
          e,
          'response:',
          responseText?.slice(0, 200)
        );
        return [];
      }
    }
  );
}