import React, { useState, useRef, useEffect, useMemo } from 'react';
import { Send, MessageSquare, Loader2, AlertTriangle, Bookmark, BookmarkCheck, BookmarkIcon } from 'lucide-react';
import { createFitnessChat } from '../services/geminiService';
import { Language } from '../types';
import { GenerateContentResponse } from '@google/genai';
import { t } from '../services/i18n';
import { generateId } from '../utils/uuid';
import { useAuth } from '../context/AuthContext';
import { useSession } from '../context/SessionContext';
import ReactMarkdown from 'react-markdown';
import remarkGfm from 'remark-gfm';
import { TopBar } from './ui/TopBar';
import SavedMessagesSheet from './SavedMessagesSheet';
import { createBookmark, deleteBookmark, SavedMessage, getBookmarks } from '../services/bookmarks';
import Snackbar from './Snackbar';

interface AICoachProps {
  lang: Language;
}

interface Message {
  id: string;
  role: 'user' | 'model';
  text: string;
  isBookmarked?: boolean;
  savedMessageId?: string;
  timestamp: number;
}

const STORAGE_KEY_PREFIX = 'ai_coach_history_';

const AICoach: React.FC<AICoachProps> = ({ lang }) => {
  const { currentUser } = useAuth();
  const { sessions: history, plans } = useSession();
  const userProfile = currentUser?.profile;
  const userId = currentUser?.id;

  // Load initial messages from local storage
  const [messages, setMessages] = useState<Message[]>(() => {
    if (!userId) return [];
    const saved = localStorage.getItem(`${STORAGE_KEY_PREFIX}${userId}`);
    if (saved) {
      try {
        return JSON.parse(saved);
      } catch (e) {
        console.error("Failed to parse saved chat history", e);
      }
    }
    return [{ id: 'intro', role: 'model', text: t('ai_intro', lang), timestamp: Date.now() }];
  });

  const [input, setInput] = useState('');
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [showSavedMessages, setShowSavedMessages] = useState(false);
  const [savedMessages, setSavedMessages] = useState<SavedMessage[]>([]);
  const [snackbar, setSnackbar] = useState<{ open: boolean, message: string }>({ open: false, message: '' });

  const chatSessionRef = useRef<any>(null); // Type 'Chat' is hard to import perfectly here without errors if not careful
  const messagesEndRef = useRef<HTMLDivElement>(null);

  // Sync with local storage
  useEffect(() => {
    if (userId) {
      localStorage.setItem(`${STORAGE_KEY_PREFIX}${userId}`, JSON.stringify(messages));
    }
  }, [messages, userId]);

  // Load bookmarks on mount to sync status
  useEffect(() => {
    if (!userId) return;
    const loadBookmarks = async () => {
      try {
        const bookmarks = await getBookmarks();
        setSavedMessages(bookmarks);
        // Update bookmarked status in local messages
        setMessages(prev => prev.map(msg => {
          const found = bookmarks.find(b => b.content === msg.text);
          if (found) {
            return { ...msg, isBookmarked: true, savedMessageId: found.id };
          }
          return msg;
        }));
      } catch (e: any) {
        if (e.message !== 'Unauthorized') {
          console.warn("Failed to load bookmarks:", e);
        }
      }
    };
    loadBookmarks();
  }, [userId, showSavedMessages]);

  useEffect(() => {
    try {
      const chat = createFitnessChat(history, lang, userProfile, plans);
      if (chat) {
        chatSessionRef.current = chat;
        // Restore history context
        // Note: Gemini SDK doesn't easily allow "restoring" state without re-sending history
        // This is a simplification; for full context restoration we'd need to rebuild history
        // For now, we start fresh session context but display old UI messages
      } else {
        setError(t('ai_error', lang));
      }
    } catch (err) {
      setError("Failed to initialize AI");
    }
  }, [history, lang, userProfile, plans]);

  const scrollToBottom = () => {
    messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
  };

  useEffect(() => {
    scrollToBottom();
  }, [messages]);

  const handleSend = async () => {
    if (!input.trim() || !chatSessionRef.current || loading) return;

    const userMsg: Message = { id: generateId(), role: 'user', text: input, timestamp: Date.now() };
    setMessages(prev => [...prev, userMsg]);
    setInput('');
    setLoading(true);

    try {
      const result: GenerateContentResponse = await chatSessionRef.current.sendMessage(userMsg.text);
      const text = result.text;

      const aiMsg: Message = {
        id: generateId(),
        role: 'model',
        text: text || "Error generating response.",
        timestamp: Date.now()
      };
      setMessages(prev => [...prev, aiMsg]);
    } catch (err) {
      console.error(err);
      let errorText = 'Connection error.';
      if (err instanceof Error) {
        errorText = err.message;
      }
      setMessages(prev => [...prev, { id: generateId(), role: 'model', text: errorText, timestamp: Date.now() }]);
    } finally {
      setLoading(false);
    }
  };

  const toggleBookmark = async (msg: Message) => {
    if (msg.role !== 'model') return;

    if (msg.isBookmarked && msg.savedMessageId) {
      // Unbookmark
      const success = await deleteBookmark(msg.savedMessageId);
      if (success) {
        setMessages(prev => prev.map(m => m.id === msg.id ? { ...m, isBookmarked: false, savedMessageId: undefined } : m));
        setSnackbar({ open: true, message: 'Bookmark removed' });
      }
    } else {
      // Bookmark
      const newBookmark = await createBookmark(msg.text, msg.role);
      if (newBookmark) {
        setMessages(prev => prev.map(m => m.id === msg.id ? { ...m, isBookmarked: true, savedMessageId: newBookmark.id } : m));
        setSnackbar({ open: true, message: 'Message saved' });
      }
    }
  };

  if (error) {
    return (
      <div className="flex flex-col items-center justify-center h-full p-6 text-center text-on-surface-variant">
        <AlertTriangle size={48} className="text-error mb-4" />
        <p>{error}</p>
      </div>
    )
  }

  return (
    <div className="flex flex-col h-full bg-surface">
      <TopBar
        title={t('ai_expert', lang)}
        icon={MessageSquare}
        actions={
          <button
            onClick={() => setShowSavedMessages(true)}
            className="p-2 text-on-surface hover:bg-surface-container-high rounded-full transition-colors"
            title="Saved Messages"
          >
            <BookmarkIcon size={20} />
          </button>
        }
      />

      {/* Messages */}
      <div className="flex-1 overflow-y-auto p-4 space-y-4 pb-4">
        {messages.map((msg) => (
          <div key={msg.id} className={`flex ${msg.role === 'user' ? 'justify-end' : 'justify-start'}`}>
            <div className={`max-w-[90%] sm:max-w-[85%] relative group`}>
              <div className={`p-4 rounded-[20px] text-sm leading-relaxed shadow-sm ${msg.role === 'user'
                ? 'bg-primary text-on-primary rounded-br-none'
                : 'bg-surface-container-high text-on-surface border border-outline-variant/20 rounded-bl-none'
                }`}>
                <div className="prose prose-invert prose-sm max-w-none">
                  <ReactMarkdown remarkPlugins={[remarkGfm]}>
                    {msg.text}
                  </ReactMarkdown>
                </div>
              </div>

              {/* Bookmark Action */}
              {msg.role === 'model' && (
                <button
                  onClick={() => toggleBookmark(msg)}
                  className={`absolute -right-8 top-2 p-1.5 rounded-full hover:bg-surface-container-high transition-colors text-on-surface-variant ${msg.isBookmarked ? 'text-primary opacity-100' : 'opacity-0 group-hover:opacity-100'}`}
                  title={msg.isBookmarked ? "Remove Bookmark" : "Bookmark Message"}
                >
                  {msg.isBookmarked ? <BookmarkCheck size={16} /> : <Bookmark size={16} />}
                </button>
              )}
            </div>
          </div>
        ))}
        {loading && (
          <div className="flex justify-start">
            <div className="bg-surface-container-high px-4 py-3 rounded-[20px] rounded-bl-none flex gap-2 items-center text-on-surface-variant text-sm">
              <Loader2 size={16} className="animate-spin" />
              {t('ai_typing', lang)}
            </div>
          </div>
        )}
        <div ref={messagesEndRef} />
      </div>

      {/* Input */}
      <div className="p-4 bg-surface-container mt-auto">
        <div className="flex gap-2 items-center bg-surface-container-high rounded-full border border-outline-variant px-2 py-1">
          <input
            type="text"
            className="flex-1 bg-transparent border-none px-4 py-3 text-on-surface focus:outline-none placeholder-on-surface-variant"
            placeholder={t('ai_placeholder', lang)}
            value={input}
            onChange={(e) => setInput(e.target.value)}
            onKeyDown={(e) => e.key === 'Enter' && handleSend()}
          />
          <button
            onClick={handleSend}
            disabled={loading || !input.trim()}
            className="w-10 h-10 bg-primary rounded-full flex items-center justify-center text-on-primary hover:opacity-90 disabled:opacity-50 transition-opacity shrink-0"
          >
            <Send size={18} />
          </button>
        </div>
      </div>

      <SavedMessagesSheet
        isOpen={showSavedMessages}
        onClose={() => setShowSavedMessages(false)}
        onUnbookmark={(id) => {
          setMessages(prev => prev.map(m => m.savedMessageId === id ? { ...m, isBookmarked: false, savedMessageId: undefined } : m));
        }}
      />

      <Snackbar
        isOpen={snackbar.open}
        message={snackbar.message}
        onClose={() => setSnackbar({ ...snackbar, open: false })}
        type="success"
      />
    </div>
  );
};

export default AICoach;
