import React, { useState, useEffect, useRef } from 'react';
import ReactMarkdown from 'react-markdown';
import remarkMath from 'remark-math';
import rehypeKatex from 'rehype-katex';
import { Send, Bot, User, Sparkles } from 'lucide-react';
import { querySystem } from '../api';
import { motion } from 'framer-motion';

const Notebook = ({ messages, setMessages, onCitationClick, notebookId }) => { // Accept notebookId
    //   const [messages, setMessages] = useState([ ... ]); // Removed internal state
    const [input, setInput] = useState('');
    const [loading, setLoading] = useState(false);
    const scrollRef = useRef(null);

    useEffect(() => {
        if (scrollRef.current) {
            scrollRef.current.scrollTop = scrollRef.current.scrollHeight;
        }
    }, [messages]);

    const handleSend = async () => {
        if (!input.trim()) return;

        const userMsg = { role: 'user', content: input };
        setMessages(prev => [...prev, userMsg]);
        setInput('');
        setLoading(true);

        try {
            // Pass notebookId to querySystem
            const response = await querySystem(userMsg.content, notebookId);
            const botMsg = {
                role: 'assistant',
                content: response.answer,
                citations: response.citations
            };
            setMessages(prev => [...prev, botMsg]);
        } catch (err) {
            setMessages(prev => [...prev, { role: 'assistant', content: 'Connection error. Please ensure background services are running.' }]);
        }
        setLoading(false);
    };



    return (
        <div className="flex flex-col h-full bg-white relative">
            {/* Scrollable Notebook Area */}
            <div className="flex-1 overflow-y-auto px-12 py-8" ref={scrollRef}>
                <div className="max-w-3xl mx-auto space-y-10 pb-32">
                    {messages.map((msg, idx) => (
                        <motion.div
                            initial={{ opacity: 0, y: 10 }}
                            animate={{ opacity: 1, y: 0 }}
                            key={idx}
                            className={`group ${msg.role === 'user' ? '' : ''}`}
                        >
                            {msg.role === 'user' ? (
                                <h3 className="text-lg font-semibold text-slate-900 mb-4">{msg.content}</h3>
                            ) : (
                                <div className="bg-white">
                                    <div className="prose prose-slate max-w-none prose-headings:font-bold prose-h1:text-2xl prose-h2:text-xl prose-p:text-slate-700 prose-li:text-slate-700">
                                        <ReactMarkdown
                                            remarkPlugins={[remarkMath]}
                                            rehypePlugins={[rehypeKatex]}
                                        >
                                            {msg.content}
                                        </ReactMarkdown>
                                    </div>

                                    {/* Citation Chips */}
                                    {msg.citations && msg.citations.length > 0 && (
                                        <div className="flex flex-wrap gap-2 mt-6 pt-4 border-t border-slate-100">
                                            {msg.citations.map((cit, i) => (
                                                <button
                                                    key={i}
                                                    onClick={() => onCitationClick(cit)} // Prop to open right panel
                                                    className="flex items-center gap-1.5 px-3 py-1.5 rounded-full bg-slate-50 border border-slate-200 text-xs font-medium text-slate-600 hover:bg-blue-50 hover:border-blue-200 hover:text-blue-600 transition-colors"
                                                >
                                                    <span className="w-4 h-4 rounded-full bg-slate-200 text-slate-600 flex items-center justify-center text-[10px] font-bold">{i + 1}</span>
                                                    {cit.paper_id.slice(0, 20)}...
                                                </button>
                                            ))}
                                        </div>
                                    )}
                                </div>
                            )}
                        </motion.div>
                    ))}

                    {loading && (
                        <motion.div
                            initial={{ opacity: 0 }}
                            animate={{ opacity: 1 }}
                            className="flex items-center gap-2 text-slate-400 italic"
                        >
                            <Sparkles size={16} className="animate-spin" />
                            <span>Analyzing documents...</span>
                        </motion.div>
                    )}
                </div>
            </div>

            {/* Input Area - Fixed at bottom center like a floating bar */}
            <div className="absolute bottom-6 left-0 right-0 px-4">
                <div className="max-w-3xl mx-auto">
                    <div className="bg-white p-2 rounded-2xl shadow-xl border border-slate-200 flex items-center gap-2">
                        <input
                            type="text"
                            className="flex-1 bg-transparent border-0 px-4 py-3 text-slate-700 placeholder-slate-400 focus:ring-0 text-base"
                            placeholder="Ask a question about your sources..."
                            value={input}
                            onChange={(e) => setInput(e.target.value)}
                            onKeyDown={(e) => e.key === 'Enter' && handleSend()}
                        />
                        <button
                            onClick={handleSend}
                            disabled={loading || !input.trim()}
                            className="p-3 rounded-xl bg-slate-900 text-white hover:bg-slate-700 disabled:opacity-50 transition-all"
                        >
                            <Send size={18} />
                        </button>
                    </div>
                    <p className="text-center text-xs text-slate-400 mt-3">
                        AI can make mistakes. Check important info.
                    </p>
                </div>
            </div>
        </div>
    );
};

export default Notebook;