services/geminiService.ts

import { GoogleGenAI, Type, Schema } from "@google/genai";
import { BentoCardData, BlogSection, ChatMessage, GeminiModel, ChartData, PaperStructure, SectionPlan, ValidationStatus, ValidationResult } from "../types";

// ============================================================================
// RETRY UTILITY WITH EXPONENTIAL BACKOFF
// ============================================================================

interface RetryConfig {
  maxRetries: number;
  baseDelay: number; // in ms
  maxDelay: number; // in ms
  retryableErrors: string[];
}

const DEFAULT_RETRY_CONFIG: RetryConfig = {
  maxRetries: 3,
  baseDelay: 1000,
  maxDelay: 10000,
  retryableErrors: [
    'RESOURCE_EXHAUSTED',
    'UNAVAILABLE', 
    'DEADLINE_EXCEEDED',
    'INTERNAL',
    'rate limit',
    'quota',
    '429',
    '500',
    '502',
    '503',
    '504',
    'timeout',
    'network',
    'ECONNRESET',
    'ETIMEDOUT'
  ]
};

/**
 * Check if an error is retryable
 */
const isRetryableError = (error: any, config: RetryConfig = DEFAULT_RETRY_CONFIG): boolean => {
  const errorMessage = error?.message?.toLowerCase() || '';
  const errorCode = error?.code?.toLowerCase() || '';
  const statusCode = error?.status?.toString() || '';
  
  return config.retryableErrors.some(pattern => 
    errorMessage.includes(pattern.toLowerCase()) ||
    errorCode.includes(pattern.toLowerCase()) ||
    statusCode.includes(pattern)
  );
};

/**
 * Sleep for a given duration
 */
const sleep = (ms: number): Promise<void> => new Promise(resolve => setTimeout(resolve, ms));

/**
 * Calculate delay with exponential backoff and jitter
 */
const calculateBackoff = (attempt: number, config: RetryConfig): number => {
  const exponentialDelay = config.baseDelay * Math.pow(2, attempt);
  const jitter = Math.random() * 0.3 * exponentialDelay; // Add 0-30% jitter
  return Math.min(exponentialDelay + jitter, config.maxDelay);
};

/**
 * Execute a function with retry logic
 */
async function withRetry<T>(
  fn: () => Promise<T>,
  operationName: string,
  config: RetryConfig = DEFAULT_RETRY_CONFIG
): Promise<T> {
  let lastError: any;
  
  for (let attempt = 0; attempt <= config.maxRetries; attempt++) {
    try {
      return await fn();
    } catch (error: any) {
      lastError = error;
      
      if (attempt === config.maxRetries || !isRetryableError(error, config)) {
        // Don't retry if max attempts reached or error is not retryable
        console.error(`[${operationName}] Failed after ${attempt + 1} attempt(s):`, error.message);
        throw new Error(formatUserFriendlyError(error, operationName));
      }
      
      const delay = calculateBackoff(attempt, config);
      console.warn(`[${operationName}] Attempt ${attempt + 1} failed, retrying in ${Math.round(delay)}ms...`, error.message);
      await sleep(delay);
    }
  }
  
  throw lastError;
}

/**
 * Format error message for user display
 */
const formatUserFriendlyError = (error: any, operation: string): string => {
  const errorMsg = error?.message?.toLowerCase() || '';
  
  if (errorMsg.includes('api key') || errorMsg.includes('invalid key') || errorMsg.includes('unauthorized')) {
    return 'Invalid API key. Please check your Gemini API key and try again.';
  }
  
  if (errorMsg.includes('rate limit') || errorMsg.includes('quota') || errorMsg.includes('429')) {
    return 'API rate limit exceeded. Please wait a moment and try again.';
  }
  
  if (errorMsg.includes('timeout') || errorMsg.includes('deadline')) {
    return `Request timed out during ${operation}. The paper may be too large or complex. Try again or use a smaller document.`;
  }
  
  if (errorMsg.includes('network') || errorMsg.includes('econnreset') || errorMsg.includes('etimedout')) {
    return 'Network error. Please check your connection and try again.';
  }
  
  if (errorMsg.includes('resource_exhausted')) {
    return 'Server resources are temporarily exhausted. Please wait a few seconds and try again.';
  }
  
  if (errorMsg.includes('unavailable') || errorMsg.includes('500') || errorMsg.includes('502') || errorMsg.includes('503')) {
    return 'The AI service is temporarily unavailable. Please try again in a moment.';
  }
  
  // Default message with original error
  return error.message || `An error occurred during ${operation}. Please try again.`;
};

const RESPONSE_SCHEMA = {
  type: Type.ARRAY,
  items: {
    type: Type.OBJECT,
    properties: {
      title: { type: Type.STRING },
      summary: { type: Type.STRING, description: "Concise summary using Markdown for emphasis (bold/italic) but no headings." },
      type: { type: Type.STRING, enum: ['stat', 'concept', 'quote', 'insight', 'process'] },
      colSpan: { type: Type.INTEGER },
      rowSpan: { type: Type.INTEGER },
      detailPrompt: { type: Type.STRING },
      mermaid: { type: Type.STRING, description: "A valid Mermaid.js graph definition (e.g. 'graph TD...') if this card describes a process, workflow, or architecture found in the PDF. Leave empty otherwise." }
    },
    required: ['title', 'summary', 'type', 'colSpan', 'rowSpan', 'detailPrompt']
  }
};

export const generateBentoCards = async (
  apiKey: string,
  model: GeminiModel,
  content: string,
  isPdf: boolean = false,
  useThinking: boolean = false
): Promise<BentoCardData[]> => {
  if (!apiKey) throw new Error("API Key is missing");

  const ai = new GoogleGenAI({ apiKey });
  
  let promptParts: any[] = [];
  
  if (isPdf) {
    promptParts.push({
      inlineData: {
        data: content, // content is base64 here
        mimeType: "application/pdf",
      },
    });
    promptParts.push({ text: "Analyze this research paper, paying close attention to any embedded figures, diagrams, and charts." });
  } else {
    promptParts.push({ text: `Analyze the following research paper text/content: \n\n${content.substring(0, 40000)}...` }); 
  }

  const thinkingInstruction = useThinking 
    ? "You are in THINKING MODE. Use your extended reasoning budget to deeply analyze the methodology and theoretical underpinnings before summarizing. Identify subtle connections and foundational axioms." 
    : "";

  promptParts.push({
    text: `
    ${thinkingInstruction}
    Create a highly structured, dynamic summary of this paper designed for a rich Bento Grid interface.
    
    RULES:
    1. Generate between 6 to 9 unique cards. Do not generate fewer than 6.
    2. The 'colSpan' (1-4) and 'rowSpan' (1-2) must vary based on the visual hierarchy and importance of the point. 
       - Use 2x2 or 3x2 for the "Main Result" or "Core Concept".
       - Use 1x1 for specific stats or quick quotes.
       - Use 2x1 or 4x2 for process steps, architectures, or lists.
       - IF you provide a 'mermaid' diagram, you MUST set colSpan >= 2 and rowSpan >= 2.
       - Ensure the total grid packs densely (Total grid width is 4 columns).
    3. 'summary' field: Use clear Markdown. You can use **bold** for emphasis. Keep it concise.
    4. 'detailPrompt' must be a technical follow-up prompt that asks to explain the "First Principles" of this card's content.
    5. 'mermaid' field: Look at the figures in the PDF. If there is a flowchart, architecture diagram, or framework overview, translate it into a valid Mermaid.js graph ('graph TD' or 'flowchart LR'). This is highly preferred for 'process' or 'concept' cards.
    
    CONTENT TO EXTRACT:
    - The Main Innovation/Contribution (High importance).
    - Key Quantitative Results (Accuracy, Speedup, etc.).
    - The "Secret Sauce" (Methodology details) -> Visualize this with Mermaid if a diagram exists in the PDF.
    - A provoking Quote from the text.
    - Limitations or Future Work.
    - First Principles or Theory involved.
    
    Return a valid JSON array.
    `
  });

  // Thinking Mode Configuration
  let effectiveModel = model;
  let requestConfig: any = {
    responseMimeType: "application/json",
    responseSchema: RESPONSE_SCHEMA as any, 
    systemInstruction: "You are a senior research scientist summarizing papers for a high-tech dashboard. You prioritize density of information, visual hierarchy, and architectural clarity.",
  };

  if (useThinking) {
    effectiveModel = 'gemini-3-pro-preview';
    requestConfig.thinkingConfig = { thinkingBudget: 32768 };
    // IMPORTANT: Do not set maxOutputTokens when thinking is enabled for 2.5/3.0 models in this context if not strictly needed, 
    // but ensuring responseSchema is present usually works. 
  }

  return withRetry(async () => {
    const response = await ai.models.generateContent({
      model: effectiveModel,
      contents: { parts: promptParts },
      config: requestConfig
    });

    const text = response.text;
    if (!text) throw new Error("No response from Gemini");

    // If thinking text is included in response.text (depending on SDK version/flags), we might need to clean it, 
    // but typically response.text returns the model output.
    // For JSON schema mode, it should be just JSON.
    
    // Attempt to find JSON if the model outputted extra chat:
    let jsonStr = text;
    const jsonMatch = text.match(/\[.*\]/s);
    if (jsonMatch) {
      jsonStr = jsonMatch[0];
    }

    const parsed = JSON.parse(jsonStr);
    
    return parsed.map((item: any, index: number) => ({
      ...item,
      id: `card-${index}-${Date.now()}`,
      expandedContent: undefined,
      isLoadingDetails: false
    }));
  }, 'generateBentoCards');
};

export const expandBentoCard = async (
  apiKey: string,
  model: GeminiModel,
  topic: string,
  detailPrompt: string,
  originalContext: string,
  useThinking: boolean = false
): Promise<string> => {
  const ai = new GoogleGenAI({ apiKey });

  const prompt = `
  Context: The user is reading a summary card about "${topic}".
  Original Paper Context (Excerpt): ${originalContext.substring(0, 5000)}...

  Task: ${detailPrompt}
  
  CRITICAL INSTRUCTION: Explain this concept from FIRST PRINCIPLES. 
  - Return the response in RICH MARKDOWN format.
  - Use Headers (###), Lists, Bold text, and Code blocks if necessary.
  - Derive the conclusion from fundamental axioms, physical laws, or basic mathematical truths.
  - Explain the "mechanistic why" — how does it actually work at the lowest level?
  - If it's a result, explain the statistical validity and the dataset composition.
  `;

  let effectiveModel = model;
  let requestConfig: any = {};

  if (useThinking) {
    effectiveModel = 'gemini-3-pro-preview';
    requestConfig.thinkingConfig = { thinkingBudget: 32768 };
  }

  return withRetry(async () => {
    const response = await ai.models.generateContent({
      model: effectiveModel,
      contents: prompt,
      config: requestConfig
    });

    return response.text || "Could not generate details.";
  }, 'expandBentoCard');
};

export const chatWithDocument = async (
  apiKey: string,
  model: GeminiModel,
  history: ChatMessage[],
  newMessage: string,
  context: string
): Promise<string> => {
  if (!apiKey) throw new Error("API Key is missing");
  
  const ai = new GoogleGenAI({ apiKey });
  
  const chatHistory = history.map(h => ({
    role: h.role === 'user' ? 'user' : 'model',
    parts: [{ text: h.text }]
  }));

  return withRetry(async () => {
    const chat = ai.chats.create({
      model: model,
      history: chatHistory,
      config: {
        systemInstruction: `You are a helpful research assistant. You have read the following paper content/summary: ${context.substring(0, 20000)}. Answer the user's questions accurately based on this context.`
      }
    });

    const result = await chat.sendMessage({ message: newMessage });
    return result.text || "";
  }, 'chatWithDocument');
};

// Paper Structure Analysis Schema
const PAPER_STRUCTURE_SCHEMA = {
  type: Type.OBJECT,
  properties: {
    paperTitle: { type: Type.STRING, description: "The title of the paper" },
    paperAbstract: { type: Type.STRING, description: "A 2-3 sentence summary of the paper" },
    mainContribution: { type: Type.STRING, description: "The key innovation or finding in one sentence" },
    keyTerms: {
      type: Type.ARRAY,
      items: { type: Type.STRING },
      description: "5-10 important technical terms used throughout the paper"
    },
    sections: {
      type: Type.ARRAY,
      items: {
        type: Type.OBJECT,
        properties: {
          title: { type: Type.STRING, description: "Engaging section title" },
          sectionType: { 
            type: Type.STRING, 
            enum: ['intro', 'background', 'methodology', 'results', 'analysis', 'applications', 'conclusion', 'custom']
          },
          focusPoints: {
            type: Type.ARRAY,
            items: { type: Type.STRING },
            description: "3-5 key points this section should cover"
          },
          suggestedVisualization: {
            type: Type.STRING,
            enum: ['mermaid', 'chart', 'equation', 'none']
          },
          visualizationHint: { 
            type: Type.STRING, 
            description: "What specifically should be visualized (e.g., 'architecture diagram from Figure 2', 'accuracy comparison from Table 1')"
          },
          estimatedLength: { type: Type.STRING, enum: ['short', 'medium', 'long'] }
        },
        required: ['title', 'sectionType', 'focusPoints', 'suggestedVisualization', 'estimatedLength']
      }
    }
  },
  required: ['paperTitle', 'paperAbstract', 'mainContribution', 'sections', 'keyTerms']
};

// Single Blog Section Schema with Progressive Disclosure Layers
const SINGLE_SECTION_SCHEMA = {
  type: Type.OBJECT,
  properties: {
    // Progressive Disclosure Layers (ADHD-Friendly)
    layers: {
      type: Type.OBJECT,
      properties: {
        thesis: { 
          type: Type.STRING, 
          description: "Tweet-length core thesis of this section (< 280 characters). The one key insight." 
        },
        takeaways: {
          type: Type.ARRAY,
          items: { type: Type.STRING },
          description: "3-5 key bullet points that summarize this section"
        },
        detailed: { 
          type: Type.STRING, 
          description: "Full detailed content in rich Markdown (300-500 words)" 
        },
        eli5Content: {
          type: Type.STRING,
          description: "Same content but explained using simple everyday language and analogies, avoiding jargon"
        }
      },
      required: ['thesis', 'takeaways', 'detailed']
    },
    content: { type: Type.STRING, description: "Rich Markdown content (300-500 words). Use **bold**, *italics*, bullet points, numbered lists, and ### subheadings." },
    visualizationType: { 
      type: Type.STRING, 
      enum: ['mermaid', 'chart', 'equation', 'none']
    },
    visualizationData: { 
      type: Type.STRING, 
      description: "For 'mermaid': valid Mermaid.js graph. For 'equation': LaTeX-style string. Empty for 'chart' or 'none'."
    },
    chartData: {
      type: Type.OBJECT,
      nullable: true,
      properties: {
        type: { type: Type.STRING, enum: ['bar', 'line', 'pie', 'area'] },
        title: { type: Type.STRING },
        data: {
          type: Type.ARRAY,
          items: {
            type: Type.OBJECT,
            properties: {
              label: { type: Type.STRING },
              value: { type: Type.NUMBER }
            },
            required: ['label', 'value']
          }
        },
        xAxis: { type: Type.STRING },
        yAxis: { type: Type.STRING }
      }
    },
    marginNotes: {
      type: Type.ARRAY,
      items: {
        type: Type.OBJECT,
        properties: {
          text: { type: Type.STRING },
          icon: { type: Type.STRING, enum: ['info', 'warning', 'tip', 'note'] }
        },
        required: ['text']
      }
    },
    technicalTerms: {
      type: Type.ARRAY,
      items: {
        type: Type.OBJECT,
        properties: {
          term: { type: Type.STRING },
          definition: { type: Type.STRING }
        },
        required: ['term', 'definition']
      }
    },
    collapsibleSections: {
      type: Type.ARRAY,
      items: {
        type: Type.OBJECT,
        properties: {
          title: { type: Type.STRING },
          content: { type: Type.STRING }
        },
        required: ['title', 'content']
      }
    }
  },
  required: ['content', 'visualizationType']
};

/**
 * Step 1: Analyze the paper and determine the best structure
 */
export const analyzePaperStructure = async (
  apiKey: string,
  model: GeminiModel,
  content: string,
  isPdf: boolean = false,
  useThinking: boolean = false
): Promise<PaperStructure> => {
  if (!apiKey) throw new Error("API Key is missing");

  const ai = new GoogleGenAI({ apiKey });
  
  let promptParts: any[] = [];
  
  if (isPdf) {
    promptParts.push({
      inlineData: {
        data: content,
        mimeType: "application/pdf",
      },
    });
    promptParts.push({ text: "Carefully analyze this research paper, including all figures, tables, diagrams, and equations." });
  } else {
    promptParts.push({ text: `Analyze this research paper: \n\n${content.substring(0, 50000)}` });
  }

  const thinkingInstruction = useThinking 
    ? "Use deep reasoning to understand the paper's structure, contributions, and the best way to present it to a general technical audience." 
    : "";

  promptParts.push({
    text: `
    ${thinkingInstruction}
    
    You are planning a distill.pub-style interactive article. Analyze this paper and create an OPTIMAL STRUCTURE for presenting it.
    
    YOUR TASK:
    1. Identify the paper's core contribution and why it matters
    2. Determine 5-8 sections that best tell the story of this paper
    3. For each section, identify what visualization would be most impactful
    
    SECTION PLANNING GUIDELINES:
    - Start with a hook that explains WHY this research matters
    - Include background only if essential for understanding
    - The methodology section should explain the "how" clearly
    - Results sections should reference SPECIFIC figures/tables from the paper
    - Include practical applications or implications
    - End with limitations and future directions
    
    VISUALIZATION PLANNING:
    - 'mermaid': Use for architectures, pipelines, workflows, decision trees
    - 'chart': Use when there are quantitative comparisons (accuracy, speed, etc.)
    - 'equation': Use for key mathematical formulations
    - 'none': Use for narrative sections
    
    Look at the paper's figures and tables - reference them in visualizationHint so we can recreate them!
    
    Return a structured JSON plan.
    `
  });

  let effectiveModel = model;
  let requestConfig: any = {
    responseMimeType: "application/json",
    responseSchema: PAPER_STRUCTURE_SCHEMA as any,
    systemInstruction: "You are a senior science editor who plans engaging, accessible articles from complex research. You identify the narrative arc and visual opportunities in papers.",
  };

  if (useThinking) {
    effectiveModel = 'gemini-3-pro-preview';
    requestConfig.thinkingConfig = { thinkingBudget: 16384 };
  }

  return withRetry(async () => {
    const response = await ai.models.generateContent({
      model: effectiveModel,
      contents: { parts: promptParts },
      config: requestConfig
    });

    const text = response.text;
    if (!text) throw new Error("No response from Gemini");

    let jsonStr = text;
    const jsonMatch = text.match(/\{.*\}/s);
    if (jsonMatch) {
      jsonStr = jsonMatch[0];
    }

    const parsed = JSON.parse(jsonStr);
    
    return {
      ...parsed,
      sections: parsed.sections.map((section: any, index: number) => ({
        ...section,
        id: `plan-${index}-${Date.now()}`
      }))
    };
  }, 'analyzePaperStructure');
};

/**
 * Step 2: Generate a single section based on the plan
 */
export const generateSingleBlogSection = async (
  apiKey: string,
  model: GeminiModel,
  content: string,
  sectionPlan: SectionPlan,
  sectionIndex: number,
  totalSections: number,
  paperContext: { title: string; abstract: string; mainContribution: string; keyTerms: string[] },
  isPdf: boolean = false,
  useThinking: boolean = false
): Promise<BlogSection> => {
  if (!apiKey) throw new Error("API Key is missing");

  const ai = new GoogleGenAI({ apiKey });
  
  let promptParts: any[] = [];
  
  if (isPdf) {
    promptParts.push({
      inlineData: {
        data: content,
        mimeType: "application/pdf",
      },
    });
  } else {
    promptParts.push({ text: `Paper content: \n\n${content.substring(0, 40000)}` });
  }

  const lengthGuide = {
    'short': '200-300 words',
    'medium': '300-450 words', 
    'long': '450-600 words'
  };

  const sectionTypeGuide: Record<string, string> = {
    'intro': 'Write an engaging opening that hooks the reader. Explain WHY this research matters to the world. Be provocative and inspiring.',
    'background': 'Provide essential context. Explain prerequisite concepts clearly. Define technical terms for a smart but non-expert reader.',
    'methodology': 'Explain HOW it works step by step. Use clear analogies. Break down complex processes into digestible parts.',
    'results': 'Present the key findings with specifics. Use actual numbers from the paper. Compare to baselines meaningfully.',
    'analysis': 'Go deeper. Discuss implications, trade-offs, and surprising findings. Address limitations honestly.',
    'applications': 'Describe real-world use cases. Be concrete about how this could be applied. Include potential impact.',
    'conclusion': 'Synthesize the key takeaways. Discuss open questions and future directions. End with a forward-looking statement.',
    'custom': 'Write compelling content following the focus points provided.'
  };

  promptParts.push({
    text: `
    PAPER CONTEXT:
    - Title: ${paperContext.title}
    - Main Contribution: ${paperContext.mainContribution}
    - Key Terms: ${paperContext.keyTerms.join(', ')}
    
    SECTION TO GENERATE (${sectionIndex + 1} of ${totalSections}):
    - Title: "${sectionPlan.title}"
    - Type: ${sectionPlan.sectionType}
    - Focus Points: ${sectionPlan.focusPoints.map((p, i) => `\n  ${i + 1}. ${p}`).join('')}
    - Target Length: ${lengthGuide[sectionPlan.estimatedLength]}
    - Visualization: ${sectionPlan.suggestedVisualization}${sectionPlan.visualizationHint ? ` - ${sectionPlan.visualizationHint}` : ''}
    
    WRITING INSTRUCTIONS:
    ${sectionTypeGuide[sectionPlan.sectionType]}
    
    ===== CRITICAL: PROGRESSIVE DISCLOSURE LAYERS (ADHD-FRIENDLY) =====
    You MUST generate content in THREE LAYERS for this section:
    
    LAYER 1 - THESIS (Tweet-Length < 280 chars):
    Write the ONE key insight of this section that a reader absolutely must know.
    Example: "Attention mechanisms let models focus on relevant parts of input, like a spotlight on important words."
    
    LAYER 2 - TAKEAWAYS (3-5 Bullet Points):
    Write 3-5 key bullet points that expand on the thesis.
    Each bullet should be one actionable/memorable piece of information.
    
    LAYER 3 - DETAILED (Full Markdown Content):
    The full explanation with all the context, examples, and depth.
    
    BONUS - ELI5 CONTENT (Simple Language Version):
    Rewrite the detailed content using everyday language, analogies, and no jargon.
    Imagine explaining to a smart 12-year-old.
    =================================================================
    
    CONTENT REQUIREMENTS:
    - Write in distill.pub style: accessible yet technically accurate
    - Use Markdown formatting: **bold** for emphasis, bullet points for lists, ### for subheadings
    - Include 1-2 margin notes with helpful asides (use 'tip' for pro tips, 'info' for context, 'warning' for caveats)
    - Define 2-3 technical terms that readers might not know
    - Add 1 collapsible "deep dive" section for readers who want more detail
    
    VISUALIZATION REQUIREMENTS:
    ${sectionPlan.suggestedVisualization === 'mermaid' ? `
    Create a Mermaid.js diagram. Use this format:
    - Start with 'graph TD' (top-down) or 'graph LR' (left-right)
    - Use descriptive node labels: A[Input Data] --> B[Process Step]
    - Keep it clear and not too complex (5-10 nodes max)
    ` : ''}
    ${sectionPlan.suggestedVisualization === 'chart' ? `
    Extract ACTUAL data from the paper to create a chart:
    - Use real numbers from tables/figures
    - Choose the right chart type (bar for comparisons, line for trends)
    - Include axis labels
    - Make it interactive: if appropriate, add slider config for "what-if" exploration
    ` : ''}
    ${sectionPlan.suggestedVisualization === 'equation' ? `
    Write the key equation in LaTeX-style notation:
    - Use \\frac{}{} for fractions
    - Use \\sum, \\prod for summations/products
    - Use Greek letters: \\alpha, \\beta, \\theta, etc.
    ` : ''}
    
    Generate the section content now.
    `
  });

  let effectiveModel = model;
  let requestConfig: any = {
    responseMimeType: "application/json",
    responseSchema: SINGLE_SECTION_SCHEMA as any,
    systemInstruction: "You are an expert science writer for distill.pub. You write clear, engaging, and technically accurate content that makes complex research accessible.",
  };

  if (useThinking) {
    effectiveModel = 'gemini-3-pro-preview';
    requestConfig.thinkingConfig = { thinkingBudget: 8192 };
  }

  return withRetry(async () => {
    const response = await ai.models.generateContent({
      model: effectiveModel,
      contents: { parts: promptParts },
      config: requestConfig
    });

    const text = response.text;
    if (!text) throw new Error("No response from Gemini");

    let jsonStr = text;
    const jsonMatch = text.match(/\{.*\}/s);
    if (jsonMatch) {
      jsonStr = jsonMatch[0];
    }

    const parsed = JSON.parse(jsonStr);
    
    return {
      id: sectionPlan.id,
      title: sectionPlan.title,
      content: parsed.content || parsed.layers?.detailed || '',
      layers: parsed.layers ? {
        thesis: parsed.layers.thesis || '',
        takeaways: parsed.layers.takeaways || [],
        detailed: parsed.layers.detailed || parsed.content || '',
        eli5Content: parsed.layers.eli5Content
      } : undefined,
      visualizationType: parsed.visualizationType,
      visualizationData: parsed.visualizationData,
      chartData: parsed.chartData,
      marginNotes: (parsed.marginNotes || []).map((note: any, noteIdx: number) => ({
        ...note,
        id: `note-${sectionIndex}-${noteIdx}-${Date.now()}`
      })),
      technicalTerms: parsed.technicalTerms || [],
      collapsibleSections: (parsed.collapsibleSections || []).map((section: any, secIdx: number) => ({
        ...section,
        id: `collapse-${sectionIndex}-${secIdx}-${Date.now()}`
      }))
    };
  }, `generateSection:${sectionPlan.title}`);
};

/**
 * Legacy function - kept for compatibility but now uses the progressive approach internally
 */
export const generateBlogContent = async (
  apiKey: string,
  model: GeminiModel,
  content: string,
  isPdf: boolean = false,
  useThinking: boolean = false
): Promise<BlogSection[]> => {
  // First, analyze the structure
  const structure = await analyzePaperStructure(apiKey, model, content, isPdf, useThinking);
  
  // Then generate each section
  const sections: BlogSection[] = [];
  const paperContext = {
    title: structure.paperTitle,
    abstract: structure.paperAbstract,
    mainContribution: structure.mainContribution,
    keyTerms: structure.keyTerms
  };

  for (let i = 0; i < structure.sections.length; i++) {
    const section = await generateSingleBlogSection(
      apiKey,
      model,
      content,
      structure.sections[i],
      i,
      structure.sections.length,
      paperContext,
      isPdf,
      useThinking
    );
    sections.push(section);
  }

  return sections;
};

// ============================================================================
// VALIDATION SYSTEM - Focused on Coded Visualizations Only
// ============================================================================

/**
 * Validate Mermaid syntax locally (basic check)
 */
const validateMermaidSyntax = (mermaidCode: string): { valid: boolean; errors: string[] } => {
  const errors: string[] = [];
  
  if (!mermaidCode || mermaidCode.trim() === '') {
    return { valid: true, errors: [] };
  }

  const code = mermaidCode.trim();
  
  // Check for valid start
  const validStarts = ['graph ', 'flowchart ', 'sequenceDiagram', 'classDiagram', 'stateDiagram', 'erDiagram', 'gantt', 'pie', 'mindmap'];
  const hasValidStart = validStarts.some(start => code.toLowerCase().startsWith(start.toLowerCase()));
  
  if (!hasValidStart) {
    errors.push('Mermaid diagram must start with a valid type (graph, flowchart, sequenceDiagram, etc.)');
  }

  // Check for basic structure in graph/flowchart
  if (code.toLowerCase().startsWith('graph') || code.toLowerCase().startsWith('flowchart')) {
    const hasNodes = /\w+\[.+\]/.test(code) || /\w+\(.+\)/.test(code) || /\w+\{.+\}/.test(code);
    if (!hasNodes && !code.includes('-->') && !code.includes('---')) {
      errors.push('Graph should contain node definitions or connections');
    }
    
    // Check for unbalanced brackets
    const openBrackets = (code.match(/\[/g) || []).length;
    const closeBrackets = (code.match(/\]/g) || []).length;
    if (openBrackets !== closeBrackets) {
      errors.push(`Unbalanced square brackets: ${openBrackets} open, ${closeBrackets} close`);
    }
  }

  // Check for markdown code fences (common error)
  if (code.includes('```')) {
    errors.push('Mermaid code should not contain markdown code fences');
  }

  return { valid: errors.length === 0, errors };
};

/**
 * Validate chart data structure
 */
const validateChartData = (chartData: ChartData | undefined): { valid: boolean; errors: string[] } => {
  const errors: string[] = [];
  
  if (!chartData) {
    return { valid: true, errors: [] };
  }

  const validTypes = ['bar', 'line', 'pie', 'area', 'scatter'];
  if (!validTypes.includes(chartData.type)) {
    errors.push(`Invalid chart type: ${chartData.type}. Must be one of: ${validTypes.join(', ')}`);
  }

  if (!chartData.data || !Array.isArray(chartData.data)) {
    errors.push('Chart data must be an array');
  } else if (chartData.data.length === 0) {
    errors.push('Chart data cannot be empty');
  } else {
    chartData.data.forEach((point, index) => {
      if (typeof point.label !== 'string') {
        errors.push(`Data point ${index}: label must be a string`);
      }
      if (typeof point.value !== 'number' || isNaN(point.value)) {
        errors.push(`Data point ${index}: value must be a valid number`);
      }
    });

    if (chartData.type === 'pie') {
      const total = chartData.data.reduce((sum, p) => sum + (p.value || 0), 0);
      if (total <= 0) {
        errors.push('Pie chart values must sum to a positive number');
      }
    }
  }

  return { valid: errors.length === 0, errors };
};

/**
 * Validate equation syntax (basic LaTeX check)
 */
const validateEquationSyntax = (equation: string): { valid: boolean; errors: string[] } => {
  const errors: string[] = [];
  
  if (!equation || equation.trim() === '') {
    return { valid: true, errors: [] };
  }

  // Check for unbalanced braces
  let braceCount = 0;
  for (const char of equation) {
    if (char === '{') braceCount++;
    if (char === '}') braceCount--;
    if (braceCount < 0) {
      errors.push('Unbalanced curly braces: closing brace without opening');
      break;
    }
  }
  if (braceCount > 0) {
    errors.push(`Unbalanced curly braces: ${braceCount} unclosed`);
  }

  return { valid: errors.length === 0, errors };
};

/**
 * Validate visualization only - no content validation needed
 */
export const validateVisualization = (section: BlogSection): ValidationStatus => {
  // Only validate if there's a coded visualization
  if (section.visualizationType === 'none' || !section.visualizationType) {
    return {
      isValidated: true,
      contentRelevance: { passed: true, score: 100, issues: [] },
      visualizationValidity: { passed: true, score: 100, issues: [] },
      overallScore: 100
    };
  }

  let vizValidation = { valid: true, errors: [] as string[] };

  switch (section.visualizationType) {
    case 'mermaid':
      vizValidation = validateMermaidSyntax(section.visualizationData || '');
      break;
    case 'chart':
      vizValidation = validateChartData(section.chartData);
      break;
    case 'equation':
      vizValidation = validateEquationSyntax(section.visualizationData || '');
      break;
  }

  return {
    isValidated: true,
    contentRelevance: { passed: true, score: 100, issues: [] },
    visualizationValidity: {
      passed: vizValidation.valid,
      score: vizValidation.valid ? 100 : 30,
      issues: vizValidation.errors
    },
    overallScore: vizValidation.valid ? 100 : 50
  };
};

// Visualization Repair Schema
const VISUALIZATION_REPAIR_SCHEMA = {
  type: Type.OBJECT,
  properties: {
    visualizationData: { 
      type: Type.STRING, 
      description: "The corrected visualization code (Mermaid, LaTeX equation, etc.)"
    },
    chartData: {
      type: Type.OBJECT,
      nullable: true,
      properties: {
        type: { type: Type.STRING, enum: ['bar', 'line', 'pie', 'area'] },
        title: { type: Type.STRING },
        data: {
          type: Type.ARRAY,
          items: {
            type: Type.OBJECT,
            properties: {
              label: { type: Type.STRING },
              value: { type: Type.NUMBER }
            },
            required: ['label', 'value']
          }
        },
        xAxis: { type: Type.STRING },
        yAxis: { type: Type.STRING }
      }
    }
  },
  required: ['visualizationData']
};

/**
 * Repair only the visualization - focused and direct
 */
export const repairVisualization = async (
  apiKey: string,
  model: GeminiModel,
  section: BlogSection,
  validationErrors: string[],
  paperContent: string,
  isPdf: boolean = false
): Promise<{ visualizationData?: string; chartData?: ChartData }> => {
  if (!apiKey) throw new Error("API Key is missing");

  const ai = new GoogleGenAI({ apiKey });

  let promptParts: any[] = [];
  
  if (isPdf) {
    promptParts.push({
      inlineData: { data: paperContent, mimeType: "application/pdf" },
    });
  } else {
    promptParts.push({ text: `Paper context: ${paperContent.substring(0, 20000)}` });
  }

  promptParts.push({
    text: `
    FIX THIS ${section.visualizationType?.toUpperCase()} VISUALIZATION.

    ORIGINAL CODE:
    """
    ${section.visualizationType === 'chart' ? JSON.stringify(section.chartData, null, 2) : section.visualizationData}
    """

    ERRORS FOUND:
    ${validationErrors.map(e => `- ${e}`).join('\n')}

    ${section.visualizationType === 'mermaid' ? `
    MERMAID REQUIREMENTS:
    - Start with 'graph TD' or 'flowchart LR'
    - Use proper syntax: A[Label] --> B[Label]
    - No markdown code fences
    - Balance all brackets
    ` : ''}

    ${section.visualizationType === 'chart' ? `
    CHART REQUIREMENTS:
    - Valid type: bar, line, pie, or area
    - Data array with {label: string, value: number} objects
    - Non-empty data array
    ` : ''}

    ${section.visualizationType === 'equation' ? `
    EQUATION REQUIREMENTS:
    - Valid LaTeX syntax
    - Balanced curly braces
    - Use \\frac{}{}, \\sum, \\alpha, etc.
    ` : ''}

    Return ONLY the fixed visualization.
    `
  });

  const requestConfig: any = {
    responseMimeType: "application/json",
    responseSchema: VISUALIZATION_REPAIR_SCHEMA as any,
    systemInstruction: "You fix visualization syntax errors. Return only the corrected code.",
  };

  return withRetry(async () => {
    const response = await ai.models.generateContent({
      model: model,
      contents: { parts: promptParts },
      config: requestConfig
    });

    const text = response.text;
    if (!text) throw new Error("No repair response");

    let jsonStr = text;
    const jsonMatch = text.match(/\{.*\}/s);
    if (jsonMatch) {
      jsonStr = jsonMatch[0];
    }

    return JSON.parse(jsonStr);
  }, 'repairVisualization');
};

/**
 * Generate section content and validate/repair visualization only
 */
export const generateAndValidateSection = async (
  apiKey: string,
  model: GeminiModel,
  content: string,
  sectionPlan: SectionPlan,
  sectionIndex: number,
  totalSections: number,
  paperContext: { title: string; abstract: string; mainContribution: string; keyTerms: string[] },
  isPdf: boolean = false,
  useThinking: boolean = false,
  maxRepairAttempts: number = 2,
  onStatusUpdate?: (status: 'generating' | 'validating' | 'repairing' | 'complete', message: string) => void
): Promise<BlogSection> => {
  
  // Step 1: Generate section content
  onStatusUpdate?.('generating', `Generating section: ${sectionPlan.title}`);
  let section = await generateSingleBlogSection(
    apiKey, model, content, sectionPlan, sectionIndex, totalSections, paperContext, isPdf, useThinking
  );

  // Step 2: Validate visualization only (no content validation - trust the LLM)
  const validation = validateVisualization(section);
  section.validationStatus = validation;

  // Step 3: Repair visualization if needed
  if (!validation.visualizationValidity.passed && section.visualizationType !== 'none') {
    let attempts = 0;
    while (!validation.visualizationValidity.passed && attempts < maxRepairAttempts) {
      attempts++;
      onStatusUpdate?.('repairing', `Fixing visualization (attempt ${attempts}/${maxRepairAttempts})...`);
      
      try {
        const repaired = await repairVisualization(
          apiKey, model, section, validation.visualizationValidity.issues, content, isPdf
        );
        
        // Apply the repaired visualization
        if (section.visualizationType === 'chart' && repaired.chartData) {
          section.chartData = repaired.chartData;
        } else if (repaired.visualizationData) {
          section.visualizationData = repaired.visualizationData;
        }
        
        // Re-validate
        const newValidation = validateVisualization(section);
        section.validationStatus = {
          ...newValidation,
          wasRepaired: true,
          repairAttempts: attempts
        };
        
        if (newValidation.visualizationValidity.passed) break;
      } catch (error) {
        console.error('Visualization repair failed:', error);
        break;
      }
    }
  }

  onStatusUpdate?.('complete', `Section complete`);
  return section;
};

// Legacy exports for backward compatibility
export const validateBlogSection = async (
  apiKey: string,
  model: GeminiModel,
  section: BlogSection,
  _sectionPlan: SectionPlan,
  _paperContext: { title: string; abstract: string; mainContribution: string; keyTerms: string[] },
  _paperContent: string,
  _isPdf: boolean = false
): Promise<ValidationStatus> => {
  // Simplified: just do local validation
  return validateVisualization(section);
};

export const repairBlogSection = async (
  apiKey: string,
  model: GeminiModel,
  section: BlogSection,
  validationStatus: ValidationStatus,
  _sectionPlan: SectionPlan,
  _paperContext: { title: string; abstract: string; mainContribution: string; keyTerms: string[] },
  paperContent: string,
  isPdf: boolean = false
): Promise<BlogSection> => {
  // Only repair visualization
  if (validationStatus.visualizationValidity.passed) {
    return section;
  }
  
  try {
    const repaired = await repairVisualization(
      apiKey, model, section, validationStatus.visualizationValidity.issues, paperContent, isPdf
    );
    
    return {
      ...section,
      visualizationData: repaired.visualizationData || section.visualizationData,
      chartData: repaired.chartData || section.chartData,
      validationStatus: {
        ...validationStatus,
        wasRepaired: true,
        repairAttempts: (section.validationStatus?.repairAttempts || 0) + 1
      }
    };
  } catch (error) {
    return section;
  }
};