import React, { useState, useEffect, useRef } from 'react';
import { AstroData } from '../types';
import { getAstroInsight, chatWithAstro, summarizeAstroData } from '../services/geminiService';
import ReactMarkdown from 'react-markdown';
import { Send, Loader2, MessageSquare, ChevronRight, ChevronLeft, Plus, X, Download, Copy, Check } from 'lucide-react';
import { motion, AnimatePresence } from 'motion/react';

interface Props {
  step: number;
  userData: AstroData;
  onNext: () => void;
  onPrev: () => void;
  onUpdateData: (newData: AstroData) => void;
}

const VARGA_INQUIRIES = [
  { id: 'D9', name: 'Marriage & Path', varga: 'D9', desc: 'Navamsa: Deep dive into Dharma, committed partnerships, and inner devotion.' },
  { id: 'D10', name: 'Career & Status', varga: 'D10', desc: 'Dashamsha: Analyze your capacity to impact the world and achieve public honor.' },
  { id: 'D2', name: 'Wealth & Resources', varga: 'D2', desc: 'Hora: Capacity to manage resources and maintain physical security.' },
  { id: 'D30', name: 'Health & Misfortune', varga: 'D30', desc: 'Trishamsha: Analyze fated diseases and deep subconscious karmic blocks.' },
  { id: 'D4', name: 'Home & Assets', varga: 'D4', desc: 'Chaturthamsha: Fixed assets, property, land, and foundational happiness.' },
  { id: 'D7', name: 'Children & Creativity', varga: 'D7', desc: 'Saptamsha: Creative output and co-creative relationships.' },
  { id: 'D16', name: 'Luxuries & Contentment', varga: 'D16', desc: 'Shodashamsha: Inner emotional contentment and vehicles.' },
];

export default function AstroStep({ step, userData, onNext, onPrev, onUpdateData }: Props) {
  const [insight, setInsight] = useState<string>('');
  const [technicalAudit, setTechnicalAudit] = useState<string>('');
  const [showTechnical, setShowTechnical] = useState(false);
  const [subStep, setSubStep] = useState(userData.currentSubStep || 0);
  const [loading, setLoading] = useState(true);
  const [question, setQuestion] = useState('');
  const [chatHistory, setChatHistory] = useState<{ role: string, text: string }[]>([]);
  const [chatLoading, setChatLoading] = useState(false);
  const [showAddChart, setShowAddChart] = useState(false);
  const [copied, setCopied] = useState(false);
  const [newChartName, setNewChartName] = useState('');
  const [newChartText, setNewChartText] = useState('');
  const chatEndRef = useRef<HTMLDivElement>(null);

  const lastFetchedStepRef = useRef<number | null>(null);

  useEffect(() => {
    setSubStep(userData.currentSubStep || 0);
  }, [userData.currentSubStep]);

  useEffect(() => {
    const fetchInitialInsight = async () => {
      // Check cache first
      const cacheKey = step === 2 ? `2.${subStep}` : `${step}`;
      if (userData.fullInterpretations?.[cacheKey]) {
        setInsight(userData.fullInterpretations[cacheKey].guidance);
        setTechnicalAudit(userData.fullInterpretations[cacheKey].audit);
        setLoading(false);
        return;
      }

      setLoading(true);
      setInsight('');
      setTechnicalAudit('');
      try {
        let currentData = { ...userData, currentSubStep: subStep };
        
        if (!currentData.summary) {
          const summary = await summarizeAstroData(currentData);
          currentData.summary = summary.text;
          onUpdateData(currentData);
        }

        const result = await getAstroInsight(step, currentData);
        
        // Split result into Guidance and Technical Audit
        const guidancePart = result.fullText.split('---')[0] || result.fullText;
        const auditPart = result.fullText.split('---')[1] || "";
        
        setInsight(guidancePart);
        setTechnicalAudit(auditPart);
        setChatHistory([{ role: 'model', text: result.fullText }]);
        
        // Save BOTH memory hook (for AI) and full interpretation (for UI cache)
        const updatedData = {
          ...currentData,
          stepInterpretations: {
            ...currentData.stepInterpretations,
            [result.memoryKey]: result.memoryHook
          },
          fullInterpretations: {
            ...currentData.fullInterpretations,
            [result.memoryKey]: { guidance: guidancePart, audit: auditPart }
          }
        };
        onUpdateData(updatedData);
        lastFetchedStepRef.current = step;
      } catch (error: any) {
        console.error(error);
        if (error?.message?.includes('429')) {
          setInsight('**Sky is Overcast (Rate Limit reached)**:\n\nYou have reached the daily free limit of your Gemini API key. Please wait until tomorrow, or consider setting up a billing account in Google AI Studio to increase your capacity.');
        } else if (error?.message?.includes('API_KEY_INVALID') || error?.message?.includes('API key not found')) {
          setInsight('Failed to load astrological insight. Please check your **VITE_GEMINI_API_KEY** in the `.env` file.');
        } else {
          setInsight('Failed to load astrological insight. There was an error consulting the stars.');
        }
      } finally {
        setLoading(false);
      }
    };
    fetchInitialInsight();
  }, [step, subStep]);

  useEffect(() => {
    chatEndRef.current?.scrollIntoView({ behavior: 'smooth' });
  }, [chatHistory]);

  const handleChat = async (e: React.FormEvent) => {
    e.preventDefault();
    if (!question.trim() || chatLoading) return;

    const userMsg = { role: 'user', text: question };
    setChatHistory(prev => [...prev, userMsg]);
    setQuestion('');
    setChatLoading(true);

    try {
      const response = await chatWithAstro(step, userData, question, chatHistory);
      setChatHistory(prev => [...prev, { role: 'model', text: response.text }]);
    } catch (error: any) {
      console.error(error);
      let errorMsg = 'Error communicating with the stars...';
      if (error?.message?.includes('API_KEY_INVALID') || error?.message?.includes('API key not found')) {
        errorMsg += ' Please check your `GEMINI_API_KEY` in the Secrets menu.';
      }
      setChatHistory(prev => [...prev, { role: 'model', text: errorMsg }]);
    } finally {
      setChatLoading(false);
    }
  };

  const handleDeepDive = async (inquiry: typeof VARGA_INQUIRIES[0]) => {
    if (!userData.divisionalCharts[inquiry.varga]) {
      setNewChartName(inquiry.varga);
      setShowAddChart(true);
      return;
    }

    setChatLoading(true);
    const deepDiveQuestion = `DEEP DIVE REQUEST: I want to analyze my ${inquiry.name} using the ${inquiry.varga} chart. Please perform Step 3.2-3.5 analysis as defined in the rules: Evaluate Varga Dignity Shifts and specific Karakas.`;
    
    // Add prompt to chat
    setChatHistory(prev => [...prev, { role: 'user', text: `Deep Dive: ${inquiry.name}` }]);
    
    try {
      const response = await chatWithAstro(step, userData, deepDiveQuestion, chatHistory);
      setChatHistory(prev => [...prev, { role: 'model', text: response.text }]);
    } catch (err) {
      console.error(err);
      setChatHistory(prev => [...prev, { role: 'model', text: 'Error performing deep dive calculation.' }]);
    } finally {
      setChatLoading(false);
    }
  };

  const handleAddChart = () => {
    if (!newChartName || !newChartText) return;
    const newData = {
      ...userData,
      divisionalCharts: {
        ...userData.divisionalCharts,
        [newChartName.toUpperCase()]: newChartText
      }
    };
    onUpdateData(newData);
    setShowAddChart(false);
    setNewChartName('');
    setNewChartText('');
  };

  const handleCopy = () => {
    navigator.clipboard.writeText(insight);
    setCopied(true);
    setTimeout(() => setCopied(false), 2000);
  };

  const handleExport = () => {
    const blob = new Blob([insight], { type: 'text/markdown' });
    const url = URL.createObjectURL(blob);
    const a = document.createElement('a');
    a.href = url;
    a.download = `Vedic_Astro_Step_${step}_Interpretation.md`;
    document.body.appendChild(a);
    a.click();
    document.body.removeChild(a);
    URL.revokeObjectURL(url);
  };

  return (
    <div className="max-w-6xl mx-auto space-y-8 p-4 md:p-8">
      <div className="flex flex-col md:flex-row justify-between items-start md:items-center border-b border-stone-200 pb-4 gap-4">
        <div className="space-y-1">
          <h2 className="text-4xl font-serif italic text-stone-900">Step {step}</h2>
          <p className="text-xs font-mono uppercase tracking-[0.2em] text-stone-400">
            {step === 1 && "Foundational Analysis"}
            {step === 2 && "Rashi Chart Promise"}
            {step === 3 && "Varga Verification"}
            {step === 4 && "Event Timing"}
          </p>
        </div>
        <div className="flex items-center gap-4">
          {step === 3 && (
            <button 
              onClick={() => setShowAddChart(true)}
              className="flex items-center gap-2 text-xs font-mono bg-stone-100 text-stone-600 px-4 py-2 rounded-full hover:bg-stone-200 transition-colors"
            >
              <Plus size={14} /> Add Varga Chart
            </button>
          )}
          <div className="flex gap-2">
            <button 
              onClick={onPrev}
              className="p-2 rounded-full border border-stone-200 hover:bg-stone-50 transition-colors"
            >
              <ChevronLeft size={24} />
            </button>
            <button 
              onClick={onNext}
              disabled={step === 4}
              className="p-2 rounded-full border border-stone-200 hover:bg-stone-50 transition-colors disabled:opacity-30"
            >
              <ChevronRight size={24} />
            </button>
          </div>
        </div>
      </div>

      <div className="grid grid-cols-1 lg:grid-cols-3 gap-8">
        {/* Main Insight Area */}
        <div className="lg:col-span-2 space-y-6">
          {loading ? (
            <div className="flex flex-col items-center justify-center h-64 space-y-4">
              <Loader2 className="animate-spin text-stone-400" size={48} />
              <p className="font-serif italic text-stone-500">Consulting the celestial alignment...</p>
            </div>
          ) : (
            <motion.div 
              initial={{ opacity: 0, y: 20 }}
              animate={{ opacity: 1, y: 0 }}
              className="prose prose-stone max-w-none bg-white p-10 rounded-3xl border border-stone-100 shadow-xl relative group leading-relaxed text-stone-800"
            >
              <div className="absolute top-6 right-6 flex gap-2 opacity-0 group-hover:opacity-100 transition-opacity">
                <button 
                  onClick={() => setShowTechnical(!showTechnical)}
                  className={`px-4 py-2 rounded-lg text-xs font-mono transition-colors ${showTechnical ? 'bg-stone-900 text-white' : 'bg-stone-100 text-stone-600 hover:bg-stone-200'}`}
                >
                  {showTechnical ? 'View Guidance' : 'View Audit'}
                </button>
                <button 
                  onClick={handleCopy}
                  className="p-2 bg-stone-100 text-stone-600 rounded-lg hover:bg-stone-200 transition-colors"
                >
                  {copied ? <Check size={16} /> : <Copy size={16} />}
                </button>
              </div>

              <div className="space-y-4">
                {showTechnical ? (
                  <div className="bg-stone-50 p-6 rounded-2xl border border-stone-100 font-mono text-sm overflow-x-auto">
                    <ReactMarkdown>{technicalAudit || insight}</ReactMarkdown>
                  </div>
                ) : (
                  <div className="space-y-6 pt-4">
                    {step === 3 && (
                      <div className="mb-12 space-y-6">
                        <div className="p-4 bg-stone-900 text-stone-100 rounded-xl">
                          <h3 className="text-lg font-serif italic mb-2 text-white">Step 3.1: Selective Deep Dive</h3>
                          <p className="text-xs text-stone-400">Choose an area of life to examine through the mathematical micro-lens of Divisional Charts (Vargas).</p>
                        </div>
                        <div className="grid grid-cols-1 md:grid-cols-2 gap-4">
                          {VARGA_INQUIRIES.map((inquiry) => {
                            const hasChart = !!userData.divisionalCharts[inquiry.varga];
                            return (
                              <button
                                key={inquiry.id}
                                onClick={() => handleDeepDive(inquiry)}
                                className={`text-left p-4 rounded-2xl border transition-all group ${
                                  hasChart 
                                    ? 'bg-stone-50 border-stone-200 hover:border-stone-900 shadow-sm' 
                                    : 'bg-white border-dashed border-stone-300 opacity-60 hover:opacity-100'
                                }`}
                              >
                                <div className="flex justify-between items-start mb-2">
                                  <span className={`px-2 py-0.5 rounded text-[10px] font-mono ${hasChart ? 'bg-stone-900 text-white' : 'bg-stone-200 text-stone-600'}`}>
                                    {inquiry.varga}
                                  </span>
                                  {hasChart ? <ChevronRight size={14} className="text-stone-400 group-hover:text-stone-900" /> : <Plus size={14} />}
                                </div>
                                <h4 className="font-serif italic text-stone-900 text-sm mb-1">{inquiry.name}</h4>
                                <p className="text-[10px] text-stone-500 line-clamp-2">{inquiry.desc}</p>
                              </button>
                            );
                          })}
                        </div>
                        <div className="h-px bg-stone-100 w-full my-8" />
                      </div>
                    )}
                    <ReactMarkdown>{insight}</ReactMarkdown>
                  </div>
                )}
              </div>
              
              {!loading && (
                <div className="mt-16 pt-8 border-t border-stone-100 flex flex-col items-center gap-6">
                  {step === 2 && subStep < 3 ? (
                    <button 
                      onClick={() => {
                        const nextSub = subStep + 1;
                        setSubStep(nextSub);
                        onUpdateData({ ...userData, currentSubStep: nextSub });
                      }}
                      className="flex items-center gap-3 bg-stone-900 text-white px-10 py-5 rounded-full font-serif italic text-2xl hover:bg-stone-800 transition-all shadow-xl hover:shadow-2xl active:scale-95"
                    >
                      Deepen This Step (Part {subStep + 2}/4)
                      <Plus size={24} />
                    </button>
                  ) : step < 4 && (
                    <button 
                      onClick={() => {
                        setSubStep(0); // Reset for next step
                        onNext();
                      }}
                      className="flex items-center gap-3 bg-stone-900 text-white px-10 py-5 rounded-full font-serif italic text-2xl hover:bg-stone-800 transition-all shadow-xl hover:shadow-2xl active:scale-95"
                    >
                      Proceed to {step === 3 ? "Final Timing (Step 4)" : `Step ${step + 1}`}
                      <ChevronRight size={24} />
                    </button>
                  )}
                  
                  {step === 2 && (
                    <div className="flex gap-2">
                      {[0, 1, 2, 3].map(i => (
                        <div 
                          key={i} 
                          className={`w-2 h-2 rounded-full ${i === subStep ? 'bg-stone-900' : 'bg-stone-200'}`} 
                        />
                      ))}
                    </div>
                  )}
                </div>
              )}
            </motion.div>
          )}

          {/* Display current charts info */}
          <div className="bg-stone-50 p-4 rounded-xl border border-stone-200">
            <h4 className="text-[10px] font-mono uppercase tracking-widest text-stone-400 mb-2">Available Charts</h4>
            <div className="flex flex-wrap gap-2">
              <span className="px-2 py-1 bg-stone-200 text-stone-600 text-[10px] font-mono rounded">D1 (Rasi)</span>
              {Object.keys(userData.divisionalCharts).map(key => (
                <span key={key} className="px-2 py-1 bg-stone-900 text-stone-50 text-[10px] font-mono rounded">{key}</span>
              ))}
            </div>
          </div>
        </div>

        {/* Deep Dive Chat */}
        <div className="lg:col-span-1 flex flex-col h-[650px] bg-stone-900 rounded-2xl overflow-hidden shadow-2xl border border-stone-800">
          <div className="p-4 border-b border-stone-800 flex items-center gap-2">
            <MessageSquare size={18} className="text-stone-400" />
            <h3 className="text-sm font-mono uppercase tracking-widest text-stone-300">Deep Dive Chat</h3>
          </div>
          
          <div className="flex-1 overflow-y-auto p-4 space-y-4 custom-scrollbar">
            {chatHistory.slice(1).map((msg, idx) => (
              <div key={idx} className={`flex ${msg.role === 'user' ? 'justify-end' : 'justify-start'}`}>
                <div className={`max-w-[85%] p-3 rounded-xl text-sm ${
                  msg.role === 'user' 
                    ? 'bg-stone-700 text-stone-100' 
                    : 'bg-stone-800 text-stone-300 border border-stone-700'
                }`}>
                  <ReactMarkdown>{msg.text}</ReactMarkdown>
                </div>
              </div>
            ))}
            {chatLoading && (
              <div className="flex justify-start">
                <div className="bg-stone-800 p-3 rounded-xl">
                  <Loader2 className="animate-spin text-stone-500" size={16} />
                </div>
              </div>
            )}
            <div ref={chatEndRef} />
          </div>

          <form onSubmit={handleChat} className="p-4 bg-stone-950 border-t border-stone-800">
            <div className="relative">
              <input 
                type="text"
                placeholder="Ask a deep dive question..."
                className="w-full bg-stone-900 border border-stone-800 rounded-full py-3 pl-4 pr-12 text-sm text-stone-200 focus:outline-none focus:ring-1 focus:ring-stone-600 transition-all"
                value={question}
                onChange={e => setQuestion(e.target.value)}
              />
              <button 
                type="submit"
                className="absolute right-2 top-1/2 -translate-y-1/2 p-2 text-stone-400 hover:text-stone-100 transition-colors"
              >
                <Send size={18} />
              </button>
            </div>
          </form>
        </div>
      </div>

      {/* Add Chart Modal */}
      <AnimatePresence>
        {showAddChart && (
          <div className="fixed inset-0 z-50 flex items-center justify-center p-4 bg-stone-900/60 backdrop-blur-sm">
            <motion.div 
              initial={{ opacity: 0, scale: 0.95 }}
              animate={{ opacity: 1, scale: 1 }}
              exit={{ opacity: 0, scale: 0.95 }}
              className="bg-white w-full max-w-lg rounded-2xl shadow-2xl overflow-hidden"
            >
              <div className="p-6 border-b border-stone-100 flex justify-between items-center">
                <h3 className="text-xl font-serif italic">Add Divisional Chart</h3>
                <button onClick={() => setShowAddChart(false)} className="text-stone-400 hover:text-stone-900">
                  <X size={20} />
                </button>
              </div>
              <div className="p-6 space-y-4">
                <div className="space-y-1">
                  <label className="text-[10px] font-mono uppercase text-stone-400">Chart Name (e.g. D30)</label>
                  <input 
                    type="text"
                    className="w-full p-3 bg-stone-50 border border-stone-200 rounded-lg focus:ring-1 focus:ring-stone-400 outline-none"
                    value={newChartName}
                    onChange={e => setNewChartName(e.target.value)}
                  />
                </div>
                <div className="space-y-1">
                  <label className="text-[10px] font-mono uppercase text-stone-400">Chart Details (Text)</label>
                  <textarea 
                    className="w-full p-3 bg-stone-50 border border-stone-200 rounded-lg h-48 font-mono text-xs focus:ring-1 focus:ring-stone-400 outline-none"
                    value={newChartText}
                    onChange={e => setNewChartText(e.target.value)}
                  />
                </div>
              </div>
              <div className="p-6 bg-stone-50 flex justify-end gap-3">
                <button 
                  onClick={() => setShowAddChart(false)}
                  className="px-6 py-2 text-sm font-mono text-stone-500 hover:text-stone-900"
                >
                  Cancel
                </button>
                <button 
                  onClick={handleAddChart}
                  className="px-6 py-2 bg-stone-900 text-white text-sm font-mono rounded-full hover:bg-stone-800"
                >
                  Add Chart
                </button>
              </div>
            </motion.div>
          </div>
        )}
      </AnimatePresence>
    </div>
  );
}