import { GoogleGenAI } from "@google/genai";
import { AstroData } from "../types";
import { ASTRO_RULES, DATA_EXTRACTION_PROMPT, buildVargaExtractionPrompt, buildAstroSummaryPrompt } from "./astroWorkflow";
import { parseJHoraGrid } from "./asciiParser";
import { 
  getHouseToSignMap, 
  formatHouseMapForAI,
  getLordshipMap,
  formatLordshipMapForAI,
  extractPlanetsFromLongitudes,
  calculateAtmakaraka,
  getPlanetHouse,
  getFunctionalRoles,
  formatFunctionalRolesForAI,
  getNaturalQualities,
  formatNaturalQualitiesForAI
} from "./vargaMapping";

const SYSTEM_BASE = `You are an expert Vedic Astrologer. You follow a strict 4-step interpretation workflow.
Your goal is to explain the user's astrological data based on the provided guide.

RULES OF ENGAGEMENT:
1. Be precise and technical, but explain concepts for a layperson.
2. Use the "Astro Context" summary as your primary source of truth.
3. If the summary is missing, refer to the raw data and create one.
`;

const SUPPORTED_MODELS_DEFAULT = [
  "gemini-2.5-flash",
  "gemini-3-flash-preview",
  "gemini-2.5-flash-tts",
  "gemini-2.5-flash-lite",
  "gemini-3.1-flash-lite-preview",
  "gemini-robotics-er-1.5-preview",
  "gemma-3-27b",
  "gemma-3-12b",
  "gemma-3-4b",
  "gemma-3-2b",
  "gemma-3-1b"
];

async function ollamaGenerate(contents: any[]) {
  const url = import.meta.env.VITE_LOCAL_LLM_URL || "http://localhost:11434/api/generate";
  const model = import.meta.env.VITE_LOCAL_LLM_MODEL || "llama3:8b";

  // Convert standard prompt structure to Ollama single prompt
  // Ollama generate API expects a simple string "prompt" usually, 
  // or we can use /api/chat if we want multi-turn. 
  // For extractions, /api/generate is fine.
  const prompt = contents.map(c => c.parts.map((p: any) => p.text).join('\n')).join('\n\n');

  console.log(`[AI Engine] Calling local model: ${model} via Ollama`);
  const startTime = Date.now();

  const response = await fetch(url, {
    method: 'POST',
    body: JSON.stringify({
      model: model,
      prompt: prompt,
      stream: false
    }),
  });

  if (!response.ok) {
    throw new Error(`Ollama error: ${response.statusText}`);
  }

  const data = await response.json();
  console.log(`[AI Engine] Local Success with ${model} in ${Date.now() - startTime}ms`);

  // Return a structure that matches the Gemini response object enough to work with .text
  return {
    text: data.response,
    model: model
  };
}

async function generateWithFallback(ai: GoogleGenAI, contents: any[], modelOverride?: string) {
  // If a specific model is requested by the user, we try ONLY that model
  if (modelOverride) {
    if (modelOverride === 'Local Ollama') {
      return await ollamaGenerate(contents);
    }

    try {
      console.log(`[AI Engine] Calling user-selected model: ${modelOverride}`);
      const response = await ai.models.generateContent({
        model: modelOverride,
        contents,
      });
      return {
        text: response.text,
        model: modelOverride
      };
    } catch (error: any) {
      console.error(`[AI Engine] Selected model ${modelOverride} failed:`, error.message);
      throw error; // Don't fall back if the user specifically asked for this one
    }
  }

  const useLocal = import.meta.env.VITE_USE_LOCAL_LLM === 'true';

  if (useLocal) {
    try {
      return await ollamaGenerate(contents);
    } catch (ollamaErr: any) {
      console.warn(`[AI Engine] Local Ollama failed, falling back to Gemini:`, ollamaErr.message);
    }
  }

  const envModel = import.meta.env.VITE_GEMINI_MODEL;
  const envFalls = import.meta.env.VITE_FALLBACK_MODELS;

  const fallbackList = envFalls
    ? envFalls.split(',').map((s: string) => s.trim())
    : SUPPORTED_MODELS_DEFAULT;

  const modelsToTry = envModel
    ? [envModel, ...fallbackList.filter(m => m !== envModel)]
    : fallbackList;

  let lastError = null;

  for (const model of modelsToTry) {
    try {
      console.log(`[AI Engine] Calling model: ${model}`);
      const response = await ai.models.generateContent({
        model,
        contents,
      });
      console.log(`[AI Engine] Success with ${model}`);
      return {
        text: response.text,
        model: model
      };
    } catch (error: any) {
      lastError = error;
      const errorMsg = error.message || String(error);
      console.warn(`[AI Engine] Model ${model} failed:`, errorMsg);

      const shouldFallback =
        errorMsg.includes('429') ||
        errorMsg.includes('quota') ||
        errorMsg.includes('Rate limit') ||
        errorMsg.includes('503') ||
        errorMsg.includes('500') ||
        errorMsg.includes('overloaded');

      if (!shouldFallback) {
        throw error;
      }
    }
  }

  throw new Error(`All models failed. Last error: ${lastError?.message || 'Unknown'}`);
}

export async function summarizeAstroData(userData: AstroData, modelOverride?: string) {
  const apiKey = import.meta.env.VITE_GEMINI_API_KEY;
  if (!apiKey) {
    throw new Error('API key not found. Please add `VITE_GEMINI_API_KEY` to your `.env` file.');
  }
  const ai = new GoogleGenAI({ apiKey });

  // 1. Pre-process the Rashi Grid for absolute determinism
  let parsedGrid = "";
  let houseMap = "";
  let deterministicFacts = "";

  if (userData.rasiChartText) {
    try {
      parsedGrid = parseJHoraGrid(userData.rasiChartText);

      // 1. PRIMARY: Extract Lagna sign from the absolute LONGITUDE TABLE (most reliable)
      const longitudeSignMatch = userData.longitudeTable.match(/(?:Lagna|Ascendant|As)\s+[\d. °'"]+\s+(Ar|Ta|Ge|Cn|Le|Vi|Li|Sc|Sg|Cp|Aq|Pi|Aries|Taurus|Gemini|Cancer|Leo|Virgo|Libra|Scorpio|Sagittarius|Capricorn|Aquarius|Pisces)\b/i);

      let lagnaSign = "";
      if (longitudeSignMatch && longitudeSignMatch[1]) {
        const signMap: Record<string, string> = {
          'Ar': 'ARIES', 'Ta': 'TAURUS', 'Ge': 'GEMINI', 'Cn': 'CANCER',
          'Le': 'LEO', 'Vi': 'VIRGO', 'Li': 'LIBRA', 'Sc': 'SCORPIO',
          'Sg': 'SAGITTARIUS', 'Cp': 'CAPRICORN', 'Aq': 'AQUARIUS', 'Pi': 'PISCES'
        };
        const rawMatch = longitudeSignMatch[1].trim();
        lagnaSign = signMap[rawMatch] || rawMatch.toUpperCase();
      }

      const gridLagnaMatch = parsedGrid.match(/The Ascendant \(As\) is in \*\*([A-Z]+)\*\*/i);
      if (!lagnaSign && gridLagnaMatch && gridLagnaMatch[1]) {
        lagnaSign = gridLagnaMatch[1];
      }

      if (lagnaSign) {
        // Simplified raw facts only - Let the AI do the logic as requested
        const planetPositions = extractPlanetsFromLongitudes(userData.longitudeTable);
        const ak = calculateAtmakaraka(planetPositions);
        
        const rawPlacementsList = Object.entries(planetPositions)
          .map(([name, pos]) => `${name}: ${pos.sign} (${pos.degree}° ${pos.minutes}' ${pos.seconds}")`)
          .join('\n');

        deterministicFacts = `
DETERMINISTIC LAGNA: ${lagnaSign}
ATMAKARAKA (AK): ${ak || 'Not Found'}

RAW PLANETARY POSITIONS (D1 RASHI):
${rawPlacementsList}

Note: Use the Lagna above to determine house numbers (e.g. ${lagnaSign} = House 1). 
Apply the Functional Role rules and Structural Evaluation logic from the Knowledge Base to these raw positions.
        `.trim();
      }
    } catch (e) {
      console.error("Failed to pre-parse Rashi grid:", e);
    }
  }


  const prompt = buildAstroSummaryPrompt(userData, parsedGrid, houseMap, deterministicFacts);

  const result = await generateWithFallback(ai, [
    { role: 'user', parts: [{ text: prompt }] }
  ], modelOverride);

  return {
    text: result.text,
    model: result.model
  };
}

function extractMemoryHook(text: string): string {
  const match = text.match(/# 📝 Memory Hook \(Internal\)\n([\s\S]*?)(?:\n---|$)/);
  if (match && match[1]) {
    return match[1].trim();
  }
  return ""; // Fallback if not found
}

export async function getAstroInsight(step: number, userData: AstroData, history: { role: string, text: string }[] = [], modelOverride?: string) {
  const apiKey = import.meta.env.VITE_GEMINI_API_KEY;
  if (!apiKey) {
    throw new Error('API key not found. Please add `VITE_GEMINI_API_KEY` to your `.env` file.');
  }
  const ai = new GoogleGenAI({ apiKey });
  const subStep = userData.currentSubStep || 0;

  // Resolve the correct rule based on step and subStep
  let rule = "";
  if (step === 0) rule = ASTRO_RULES.step0;
  else if (step === 1) rule = ASTRO_RULES.foundational;
  else if (step === 2) {
    const subStepRules = [ASTRO_RULES.step2_1, ASTRO_RULES.step2_2, ASTRO_RULES.step2_3, ASTRO_RULES.step2_4];
    rule = subStepRules[subStep] || ASTRO_RULES.step2_1;
  }
  else if (step === 3) rule = ASTRO_RULES.vargaVerification;
  else if (step === 4) rule = ASTRO_RULES.timing;

  const context = userData.summary ? `ASTRO CONTEXT (Pre-digested):\n${userData.summary}` : `RAW DATA:\n${JSON.stringify(userData, null, 2)}`;

  const previousInterpretations = userData.stepInterpretations
    ? Object.entries(userData.stepInterpretations)
      .filter(([s]) => parseInt(s) <= step) // Include current step if we are in sub-steps
      .map(([s, text]) => `PREVIOUS CONTEXT (Step ${s}):\n${text}`)
      .join('\n\n')
    : '';

  const contents = [];

  // First message contains system base, step rules, formatting rules, and context
  let firstMessageText = `${SYSTEM_BASE}\n\n${rule}\n\n${ASTRO_RULES.formatting}\n\n${context}`;
  if (previousInterpretations) {
    firstMessageText += `\n\n${previousInterpretations}`;
  }

  if (history.length === 0) {
    contents.push({
      role: 'user',
      parts: [{ text: `${firstMessageText}\n\nPlease provide the analysis. Follow the MANDATORY format.` }]
    });
  } else {
    contents.push({ role: 'user', parts: [{ text: firstMessageText }] });
    contents.push(...history.map(h => ({
      role: h.role === 'user' ? 'user' : 'model',
      parts: [{ text: h.text }]
    })));

    const lastRole = contents[contents.length - 1].role;
    if (lastRole === 'model') {
      contents.push({ role: 'user', parts: [{ text: "Please provide the analysis. Follow the MANDATORY format." }] });
    } else {
      const lastPart = contents[contents.length - 1].parts[0];
      if ('text' in lastPart) {
        lastPart.text += "\n\nPlease provide the analysis. Follow the MANDATORY format.";
      }
    }
  }

  const response = await generateWithFallback(ai, contents, modelOverride);

  const rawText = response.text;
  const memoryKey = step === 2 ? `2.${subStep}` : `${step}`;

  return {
    fullText: rawText,
    memoryHook: extractMemoryHook(rawText),
    memoryKey: memoryKey,
    model: response.model
  };
}

export async function chatWithAstro(step: number, userData: AstroData, question: string, history: { role: string, text: string }[], modelOverride?: string) {
  const apiKey = import.meta.env.VITE_GEMINI_API_KEY;
  if (!apiKey) {
    throw new Error('API key not found. Please add `VITE_GEMINI_API_KEY` to your `.env` file.');
  }
  const ai = new GoogleGenAI({ apiKey });
  const subStep = userData.currentSubStep || 0;

  let rule = "";
  if (step === 0) rule = ASTRO_RULES.step0;
  else if (step === 1) rule = ASTRO_RULES.foundational;
  else if (step === 2) {
    const subStepRules = [ASTRO_RULES.step2_1, ASTRO_RULES.step2_2, ASTRO_RULES.step2_3, ASTRO_RULES.step2_4];
    rule = subStepRules[subStep] || ASTRO_RULES.step2_1;
  }
  else if (step === 3) rule = ASTRO_RULES.vargaVerification;
  else if (step === 4) rule = ASTRO_RULES.timing;

  const context = userData.summary ? `ASTRO CONTEXT (Pre-digested):\n${userData.summary}` : `RAW DATA:\n${JSON.stringify(userData, null, 2)}`;
  const firstMessageText = `${SYSTEM_BASE}\n\n${rule}\n\n${ASTRO_RULES.formatting}\n\n${context}`;

  const contents = [];
  contents.push({ role: 'user', parts: [{ text: firstMessageText }] });

  contents.push(...history.map(h => ({
    role: h.role === 'user' ? 'user' : 'model',
    parts: [{ text: h.text }]
  })));

  const lastRole = contents[contents.length - 1].role;
  if (lastRole === 'model') {
    contents.push({ role: 'user', parts: [{ text: `Question related to Step ${step}: ${question}` }] });
  } else {
    const lastPart = contents[contents.length - 1].parts[0];
    if ('text' in lastPart) {
      lastPart.text += `\n\nQuestion related to Step ${step}: ${question}`;
    }
  }

  const response = await generateWithFallback(ai, contents, modelOverride);
  return {
    text: response.text,
    model: response.model
  };
}

export async function summarizeVargaData(
  varga: string,
  rawData: string,
  rasiChartText: string,
  modelOverride?: string
) {
  const apiKey = import.meta.env.VITE_GEMINI_API_KEY;
  if (!apiKey) {
    throw new Error('API key not found.');
  }
  const ai = new GoogleGenAI({ apiKey });

  const prompt = buildVargaExtractionPrompt(varga, rawData, rasiChartText);

  const result = await generateWithFallback(ai, [
    { role: 'user', parts: [{ text: prompt }] }
  ], modelOverride);

  return {
    text: result.text,
    model: result.model
  };
}

export function extractErrorMessage(error: any): string {
  const msg = error?.message || String(error);

  // Try to parse JSON if the error message is a stringified JSON (common in @google/genai)
  try {
    const parsed = JSON.parse(msg);
    if (parsed.error?.message) return parsed.error.message;
    if (parsed.message) return parsed.message;
  } catch (e) {
    // Not JSON, just use the string
  }

  if (msg.includes('429') || msg.includes('quota')) {
    return "API Quota Exceeded. Please try again in a few seconds or switch to a different model (e.g. Lite or Gemma).";
  }

  if (msg.includes('503') || msg.includes('overloaded')) {
    return "The model is currently overloaded. Please try again or switch models.";
  }

  return msg;
}