frontend/src/components/ChatPanel.tsx

"use client";

import { useState, useRef, useEffect, useCallback } from "react";
import { Send, Map as MapIcon, Database, Sparkles, BarChart3, Brain } from "lucide-react";
import { cn } from "@/lib/utils";
import dynamic from "next/dynamic";
import ReactMarkdown from "react-markdown";
import remarkGfm from "remark-gfm";

// Dynamic import for chart renderer to avoid SSR issues
const ChartRenderer = dynamic(
    () => import("./charts/ChartRenderer"),
    { ssr: false, loading: () => <div className="h-64 bg-slate-100 animate-pulse rounded-xl" /> }
);

// ============================================================================
// ============================================================================
// Types
// ============================================================================

import type { MapLayer } from "./MapViewer";

interface ChartData {
    type: 'bar' | 'line' | 'pie' | 'donut';
    title?: string;
    data: Array<{ [key: string]: any }>;
    xKey?: string;
    yKey?: string;
    lines?: Array<{ key: string; color?: string; name?: string }>;
}

interface Message {
    role: "user" | "assistant";
    content: string;
    sql?: string;
    citations?: string[];
    intent?: string;
    chart?: ChartData;
    rawData?: Array<any>;
    thoughts?: string;
    status?: string; // e.g., "Thinking...", "Writing SQL...", etc.
}

interface ChatPanelProps {
    onMapUpdate?: (geojson: any) => void;
    layers?: MapLayer[];
}

// ============================================================================
// Loading Indicator Component
// ============================================================================

function LoadingDots() {
    return (
        <div className="flex gap-1">
            <span className="w-1.5 h-1.5 bg-indigo-400 rounded-full animate-bounce" style={{ animationDelay: "0ms" }} />
            <span className="w-1.5 h-1.5 bg-indigo-400 rounded-full animate-bounce" style={{ animationDelay: "150ms" }} />
            <span className="w-1.5 h-1.5 bg-indigo-400 rounded-full animate-bounce" style={{ animationDelay: "300ms" }} />
        </div>
    );
}

// ============================================================================
// Message Bubble Component
// ============================================================================

interface MessageBubbleProps {
    message: Message;
}

function MessageBubble({ message: m }: MessageBubbleProps) {
    const isUser = m.role === "user";
    const isLoading = !!m.status && !m.content;
    const hasContent = !!m.content;
    const hasThoughts = !!m.thoughts;

    // Auto-collapse reasoning when generation is done
    const [isThinkingOpen, setIsThinkingOpen] = useState(false);

    useEffect(() => {
        if (m.status === "Thinking..." || m.status === "Reasoning...") {
            setIsThinkingOpen(true);
        } else if (!m.status && hasThoughts) {
            // Collapse when done
            setIsThinkingOpen(false);
        }
    }, [m.status, hasThoughts]);

    return (
        <div className={cn(
            "flex flex-col",
            isUser ? "ml-auto items-end max-w-[90%]" : "mr-auto items-start w-full"
        )}>
            {/* Main Message Bubble */}
            <div className={cn(
                "px-4 py-3 rounded-2xl text-sm leading-relaxed shadow-sm",
                isUser
                    ? "bg-gradient-to-br from-indigo-600 to-indigo-700 text-white rounded-br-md max-w-full"
                    : "bg-white text-slate-700 rounded-bl-md border border-slate-100 max-w-[95%]"
            )}>
                {/* Status Indicator (inside bubble when loading) */}
                {isLoading && (
                    <div className="flex items-center gap-2 text-slate-400">
                        <LoadingDots />
                        <span className="text-sm italic">{m.status}</span>
                    </div>
                )}

                {/* Content with Markdown */}
                {hasContent && (
                    <div className={cn(
                        "prose prose-sm max-w-none break-words",
                        isUser ? "prose-invert text-white" : "prose-slate"
                    )}>
                        <ReactMarkdown
                            remarkPlugins={[remarkGfm]}
                            components={{
                                pre: ({ node, ...props }) => (
                                    <div className="overflow-auto w-full my-2 bg-black/10 rounded-lg p-2" {...props as any} />
                                ),
                                code: ({ node, ...props }) => (
                                    <code className="bg-black/10 rounded px-1 py-0.5" {...props as any} />
                                )
                            }}
                        >
                            {m.content}
                        </ReactMarkdown>
                    </div>
                )}

                {/* Inline status when content is present but still loading more */}
                {hasContent && m.status && (
                    <div className="mt-2 flex items-center gap-2 text-xs text-slate-400 italic border-t border-slate-100 pt-2">
                        <LoadingDots />
                        <span>{m.status}</span>
                    </div>
                )}
            </div>

            {/* Thought Process (Collapsible) - Auto-collapses when done */}
            {hasThoughts && (
                <div className="mt-2 w-full max-w-[95%] group">
                    <button
                        onClick={() => setIsThinkingOpen(!isThinkingOpen)}
                        className="text-xs text-indigo-500 cursor-pointer hover:text-indigo-700 flex items-center gap-1 font-medium select-none list-none outline-none"
                    >
                        <div className="flex items-center gap-1.5 px-2 py-1 bg-indigo-50 rounded-full border border-indigo-100">
                            <Brain className="w-3 h-3 text-indigo-500" />
                            <span>AI Reasoning</span>
                            <span className="ml-1 opacity-50">{isThinkingOpen ? "▼" : "▶"}</span>
                        </div>
                    </button>

                    {isThinkingOpen && (
                        <div className="mt-2 p-3 bg-slate-50 rounded-xl border border-slate-100 text-xs text-slate-600 space-y-2 animate-in slide-in-from-top-2 duration-200">
                            <div className="flex gap-2">
                                <div className="w-0.5 bg-indigo-200 rounded-full shrink-0" />
                                <div className="leading-relaxed prose prose-xs max-w-none prose-slate whitespace-pre-wrap break-words w-full">
                                    <ReactMarkdown remarkPlugins={[remarkGfm]}>{m.thoughts}</ReactMarkdown>
                                </div>
                            </div>
                        </div>
                    )}
                </div>
            )}

            {/* Chart Visualization */}
            {m.chart && (
                <div className="mt-3 w-full max-w-md">
                    <ChartRenderer chart={m.chart} />
                </div>
            )}

            {/* Data Citations */}
            {m.citations && m.citations.length > 0 && (
                <div className="mt-2 flex flex-wrap gap-1.5">
                    {m.citations.map((citation, idx) => (
                        <span
                            key={idx}
                            className="inline-flex items-center gap-1 text-xs px-2 py-1 bg-blue-50 text-blue-700 rounded-full border border-blue-100"
                        >
                            <Database className="w-3 h-3" />
                            {citation}
                        </span>
                    ))}
                </div>
            )}

            {/* SQL Query and Raw Data (collapsible) */}
            {(m.sql || m.rawData) && (
                <details className="mt-2 w-full max-w-[95%]">
                    <summary className="text-xs text-slate-500 cursor-pointer hover:text-slate-700 flex items-center gap-1">
                        <BarChart3 className="w-3 h-3" />
                        View SQL & Raw Data
                    </summary>
                    <div className="mt-2 space-y-3">
                        {m.sql && (
                            <div className="text-xs font-mono bg-slate-800 text-green-400 p-3 rounded-lg overflow-x-auto whitespace-pre-wrap">
                                {m.sql}
                            </div>
                        )}

                        {m.rawData && m.rawData.length > 0 && m.rawData[0] && (
                            <div className="overflow-x-auto border rounded-lg max-h-60">
                                <table className="min-w-full divide-y divide-slate-200 text-xs">
                                    <thead className="bg-slate-50 sticky top-0">
                                        <tr>
                                            {Object.keys(m.rawData[0] || {}).map((key) => (
                                                <th key={key} className="px-3 py-2 text-left font-medium text-slate-500 uppercase tracking-wider">
                                                    {key}
                                                </th>
                                            ))}
                                        </tr>
                                    </thead>
                                    <tbody className="bg-white divide-y divide-slate-200">
                                        {m.rawData.map((row, idx) => (
                                            <tr key={idx} className="hover:bg-slate-50">
                                                {Object.values(row).map((val: any, vIdx) => (
                                                    <td key={vIdx} className="px-3 py-2 whitespace-nowrap text-slate-600">
                                                        {typeof val === 'object' ? JSON.stringify(val) : String(val)}
                                                    </td>
                                                ))}
                                            </tr>
                                        ))}
                                    </tbody>
                                </table>
                            </div>
                        )}
                    </div>
                </details>
            )}
        </div>
    );
}

// ============================================================================
// Main ChatPanel Component
// ============================================================================

export default function ChatPanel({ onMapUpdate, layers = [] }: ChatPanelProps) {
    const [messages, setMessages] = useState<Message[]>([
        {
            role: "assistant",
            content: "Welcome! I'm GeoQuery, your Territorial Intelligence Assistant for Panama."
        }
    ]);
    const [input, setInput] = useState("");
    const [loading, setLoading] = useState(false);
    const messagesEndRef = useRef<HTMLDivElement>(null);

    // Slash command state
    const [showSuggestions, setShowSuggestions] = useState(false);
    const [suggestionQuery, setSuggestionQuery] = useState("");
    const [activeSuggestionIndex, setActiveSuggestionIndex] = useState(0);
    const inputRef = useRef<HTMLDivElement>(null);
    const [cursorPosition, setCursorPosition] = useState<{ top: number, left: number } | null>(null);
    const slashContext = useRef<{ node: Node, index: number } | null>(null);

    // Auto-scroll to bottom when new messages arrive
    useEffect(() => {
        messagesEndRef.current?.scrollIntoView({ behavior: "smooth" });
    }, [messages]);

    // ========================================================================
    // Message Updater - Properly handles immutable state updates
    // ========================================================================
    const updateLastMessage = useCallback((updater: (msg: Message) => Message) => {
        setMessages(prev => {
            const newMsgs = [...prev];
            const lastMsgIndex = newMsgs.length - 1;
            // Create a new message object to avoid mutation
            newMsgs[lastMsgIndex] = updater({ ...newMsgs[lastMsgIndex] });
            return newMsgs;
        });
    }, []);

    // ========================================================================
    // Slash Command Logic & Rich Text Input
    // ========================================================================

    // Filter layers based on query
    const filteredLayers = layers.filter(layer =>
        layer.name.toLowerCase().includes(suggestionQuery.toLowerCase())
    );

    const handleInput = () => {
        if (!inputRef.current) return;

        const text = inputRef.current.innerText;
        setInput(text); // Keep simple text state for disabled check

        // Detect slash command using Selection API
        const selection = window.getSelection();
        if (!selection || selection.rangeCount === 0) {
            setShowSuggestions(false);
            return;
        }

        const range = selection.getRangeAt(0);
        const node = range.startContainer;

        // Only trigger if we are in a text node
        if (node.nodeType === Node.TEXT_NODE && node.textContent) {
            const textBeforeCursor = node.textContent.slice(0, range.startOffset);
            const lastSlashIndex = textBeforeCursor.lastIndexOf("/");

            if (lastSlashIndex !== -1) {
                // Check if contains spaces (simple heuristic to stop looking too far back)
                const query = textBeforeCursor.slice(lastSlashIndex + 1);
                if (!query.includes(" ")) {
                    setShowSuggestions(true);
                    setSuggestionQuery(query);
                    setActiveSuggestionIndex(0);

                    // Save context for replacement
                    slashContext.current = { node, index: lastSlashIndex };

                    // Get coordinates for popup
                    const rect = range.getBoundingClientRect();
                    const inputRect = inputRef.current.getBoundingClientRect();
                    setCursorPosition({
                        top: rect.top - inputRect.top,
                        left: rect.left - inputRect.left
                    });
                    return;
                }
            }
        }
        setShowSuggestions(false);
    };

    const handleSuggestionSelect = (layer: MapLayer) => {
        if (!inputRef.current || !slashContext.current) return;

        // Use stored context
        const { node, index } = slashContext.current;

        // Validate node is still in DOM (basic check)
        if (!document.contains(node)) return;

        const selection = window.getSelection();
        const range = document.createRange();

        try {
            // Calculate end index: slash index + 1 (for '/') + query length
            // But query length might have changed since last render?
            // Better: Replace from slash index to the CURRENT cursor (if we still have focus)
            // OR simpler: Replace from slash index to (slash index + 1 + suggestionQuery.length)

            // We'll trust suggestionQuery state as it drives the filtering
            const endOffset = index + 1 + suggestionQuery.length;

            // Safety check limits
            const safeEndOffset = Math.min(endOffset, (node.textContent?.length || 0));

            range.setStart(node, index);
            range.setEnd(node, safeEndOffset);
            range.deleteContents();

            // Create the chip
            const chip = document.createElement("span");
            chip.contentEditable = "false";
            // Fixed height 20px (h-5), smaller text, constrained width
            chip.className = "inline-flex items-center gap-1 px-1.5 h-5 mx-1 rounded-full text-[11px] font-medium select-none align-middle transition-transform hover:scale-105 cursor-default max-w-[160px] truncate border";

            // Style: 15% opacity background, matching border
            chip.style.backgroundColor = `${layer.style.color}26`; // 26 = ~15% opacity
            chip.style.color = layer.style.color;
            chip.style.borderColor = `${layer.style.color}40`; // 25% opacity border

            chip.dataset.layerId = layer.id;
            chip.dataset.layerName = layer.name;

            // Dot icon
            chip.innerHTML = `
                <span style="background-color: ${layer.style.color}" class="w-1.5 h-1.5 rounded-full inline-block shrink-0"></span>
                <span class="truncate">${layer.name}</span>
            `;

            // Insert chip
            range.insertNode(chip);

            // Add a space after
            const space = document.createTextNode("\u00A0");
            range.setStartAfter(chip);
            range.setEndAfter(chip);
            range.insertNode(space);

            // Restore Selection to end
            range.setStartAfter(space);
            range.setEndAfter(space);

            if (selection) {
                selection.removeAllRanges();
                selection.addRange(range);
            }

            setShowSuggestions(false);
            setInput(inputRef.current.innerText); // update state

            // Force focus back
            inputRef.current.focus();

        } catch (e) {
            console.error("Error inserting chip:", e);
        }
    };

    const handleKeyDown = (e: React.KeyboardEvent<HTMLDivElement>) => {
        if (showSuggestions && filteredLayers.length > 0) {
            if (e.key === "ArrowDown") {
                e.preventDefault();
                setActiveSuggestionIndex(prev => (prev + 1) % filteredLayers.length);
            } else if (e.key === "ArrowUp") {
                e.preventDefault();
                setActiveSuggestionIndex(prev => (prev - 1 + filteredLayers.length) % filteredLayers.length);
            } else if (e.key === "Enter" || e.key === "Tab") {
                e.preventDefault();
                handleSuggestionSelect(filteredLayers[activeSuggestionIndex]);
            } else if (e.key === "Escape") {
                setShowSuggestions(false);
            }
            return;
        }

        if (e.key === "Enter" && !e.shiftKey) {
            e.preventDefault();
            sendMessage();
        }
    };

    // ========================================================================
    // Send Message Handler
    // ========================================================================
    const sendMessage = useCallback(async () => {
        if (!inputRef.current || (!inputRef.current.innerText.trim() && !input.trim()) || loading) return;

        // Construct message with IDs from DOM
        let fullMessage = "";
        const nodes = inputRef.current.childNodes;
        let displayMessage = ""; // For user chat bubble (clean text)

        nodes.forEach(node => {
            if (node.nodeType === Node.TEXT_NODE) {
                fullMessage += node.textContent;
                displayMessage += node.textContent;
            } else if (node.nodeType === Node.ELEMENT_NODE) {
                const el = node as HTMLElement;
                if (el.dataset.layerId) {
                    // It's a chip
                    fullMessage += ` Layer ${el.dataset.layerName} (ID: ${el.dataset.layerId}) `;
                    displayMessage += ` ${el.dataset.layerName} `;
                } else {
                    fullMessage += el.innerText;
                    displayMessage += el.innerText;
                }
            }
        });

        // Fallback if empty (shouldn't happen with check above)
        if (!fullMessage.trim()) return;

        setMessages(prev => [...prev, { role: "user", content: displayMessage }]);

        // Clear input
        if (inputRef.current) inputRef.current.innerHTML = "";
        setInput("");

        setLoading(true);
        setShowSuggestions(false);

        // Add placeholder assistant message with initial status
        setMessages(prev => [...prev, {
            role: "assistant",
            content: "",
            thoughts: "",
            status: "Thinking..."
        }]);

        try {
            const history = messages.map(m => ({ role: m.role, content: m.content }));
            const apiUrl = process.env.NEXT_PUBLIC_API_URL || "http://localhost:8000/api/v1";
            const response = await fetch(`${apiUrl}/chat/stream`, {
                method: "POST",
                headers: { "Content-Type": "application/json" },
                body: JSON.stringify({ message: fullMessage, history }) // Send ID-enriched message
            });

            const reader = response.body?.getReader();
            const decoder = new TextDecoder();

            if (!reader) {
                updateLastMessage(msg => ({ ...msg, content: "Error: No response stream", status: undefined }));
                return;
            }

            let buffer = "";
            let currentEventType: string | null = null;

            while (true) {
                const { done, value } = await reader.read();
                if (done) break;

                const chunk = decoder.decode(value, { stream: true });
                buffer += chunk;

                const lines = buffer.split('\n');
                buffer = lines.pop() || ""; // Keep partial line in buffer

                for (const line of lines) {
                    if (line.startsWith("event:")) {
                        currentEventType = line.slice(6).trim();
                    } else if (line.startsWith("data:")) {
                        const dataStr = line.slice(5).trim();
                        if (!dataStr) continue;

                        try {
                            const data = JSON.parse(dataStr);

                            // Handle map updates OUTSIDE state updater
                            if (currentEventType === "result" && data.geojson && onMapUpdate) {
                                onMapUpdate(data.geojson);
                            }

                            // Update message based on event type
                            if (currentEventType === "chunk") {
                                if (data.type === "text") {
                                    updateLastMessage(msg => ({
                                        ...msg,
                                        content: msg.content + data.content
                                    }));
                                } else if (data.type === "thought") {
                                    updateLastMessage(msg => ({
                                        ...msg,
                                        thoughts: (msg.thoughts || "") + data.content,
                                        status: "Reasoning..."
                                    }));
                                }
                            } else if (currentEventType === "result") {
                                updateLastMessage(msg => ({
                                    ...msg,
                                    content: data.response || msg.content,
                                    sql: data.sql_query,
                                    chart: data.chart_data,
                                    rawData: data.raw_data,
                                    citations: data.data_citations,
                                    status: undefined // Clear status on completion
                                }));
                            } else if (currentEventType === "status") {
                                updateLastMessage(msg => ({
                                    ...msg,
                                    status: data.status
                                }));
                            }
                        } catch (e) {
                            console.error("Error parsing SSE JSON:", e);
                        }
                    }
                }
            }
        } catch (err) {
            console.error("Stream error:", err);
            setMessages(prev => [...prev, {
                role: "assistant",
                content: "Error connecting to server. Please try again."
            }]);
        } finally {
            setLoading(false);
            // Ensure status is cleared after stream ends
            updateLastMessage(msg => ({ ...msg, status: undefined }));
        }
    }, [input, loading, messages, onMapUpdate, updateLastMessage, layers]);

    // ========================================================================
    // Render
    // ========================================================================
    return (
        <div className="flex flex-col h-full bg-gradient-to-b from-slate-50 to-white border-r shadow-lg relative">
            {/* Header */}
            <div className="p-4 border-b flex items-center justify-between bg-white/80 backdrop-blur-sm">
                <h1 className="font-bold text-lg flex items-center gap-2">
                    <div className="p-1.5 bg-indigo-600 rounded-lg">
                        <MapIcon className="w-4 h-4 text-white" />
                    </div>
                    <span className="bg-gradient-to-r from-indigo-600 to-purple-600 bg-clip-text text-transparent">
                        GeoQuery
                    </span>
                </h1>
                <div className="flex items-center gap-1.5 text-xs font-medium px-2 py-1 bg-emerald-50 text-emerald-700 rounded-full border border-emerald-100">
                    <Sparkles className="w-3 h-3" />
                    Gemini 3 Flash
                </div>
            </div>

            {/* Messages */}
            <div className="flex-1 overflow-y-auto p-4 space-y-4">
                {messages.map((m, i) => (
                    <MessageBubble key={i} message={m} />
                ))}
                <div ref={messagesEndRef} />
            </div>

            {/* Input */}
            <div className="p-4 border-t bg-white/80 backdrop-blur-sm relative">
                {/* Suggestions Popup */}
                {showSuggestions && filteredLayers.length > 0 && (
                    <div
                        className="absolute bg-white rounded-xl shadow-xl border border-slate-200 overflow-hidden z-20 max-h-48 overflow-y-auto w-64 animate-in fade-in slide-in-from-bottom-2 duration-200"
                        style={{
                            bottom: "100%", // Always anchor to bottom of input area for now, simpler than dynamic text coords
                            left: cursorPosition ? Math.min(cursorPosition.left + 20, 300) : 20,
                            marginBottom: "10px"
                        }}
                    >
                        <div className="p-2 bg-slate-50 text-xs text-slate-500 border-b border-slate-100 font-medium flex items-center gap-1">
                            <Sparkles className="w-3 h-3 text-indigo-400" />
                            Reference Map Layer
                        </div>
                        {filteredLayers.map((layer, idx) => (
                            <button
                                key={layer.id}
                                onMouseDown={(e) => {
                                    e.preventDefault(); // Prevent losing focus from input
                                    handleSuggestionSelect(layer);
                                }}
                                className={cn(
                                    "w-full px-3 py-2 flex items-center gap-2 text-sm text-left hover:bg-slate-50 transition-colors",
                                    idx === activeSuggestionIndex ? "bg-indigo-50 text-indigo-700" : "text-slate-700"
                                )}
                            >
                                <div
                                    className="w-2.5 h-2.5 rounded-full shrink-0 ring-2 ring-white shadow-sm"
                                    style={{ backgroundColor: layer.style.color }}
                                />
                                <span className="truncate font-medium">{layer.name}</span>
                                {layer.pointMarker?.icon && <span className="text-base ml-auto opacity-70">{layer.pointMarker.icon}</span>}
                            </button>
                        ))}
                    </div>
                )}

                <div className="flex items-end gap-2 bg-white border border-slate-200 rounded-2xl px-4 py-3 shadow-sm focus-within:ring-2 focus-within:ring-indigo-100 focus-within:border-indigo-300 transition-all">
                    <div
                        ref={inputRef}
                        contentEditable
                        className="flex-1 outline-none text-sm text-gray-900 leading-relaxed min-h-[20px] max-h-32 overflow-y-auto whitespace-pre-wrap break-words empty:before:content-[attr(data-placeholder)] empty:before:text-slate-400"
                        data-placeholder="Ask about Panama's data..."
                        onInput={handleInput}
                        onKeyDown={handleKeyDown}
                        // Prevent tab interactions escaping the editor improperly
                        suppressContentEditableWarning
                    />



                    {/* Send Button */}
                    <button
                        onClick={sendMessage}
                        disabled={loading || (!input.trim() && !inputRef.current?.innerText.trim())}
                        className={cn(
                            "p-2 rounded-xl transition-all shrink-0 pb-2",
                            loading || (!input.trim() && !inputRef.current?.innerText.trim())
                                ? "bg-slate-100 text-slate-400 cursor-not-allowed"
                                : "bg-indigo-600 text-white hover:bg-indigo-700 shadow-sm hover:shadow"
                        )}
                    >
                        <Send className="w-4 h-4" />
                    </button>
                </div>
                <p className="text-xs text-slate-400 mt-2 text-center">
                    GeoQuery can query datasets and visualize results as maps and charts.
                    <br />
                    Use '/' to reference layers.
                </p>
            </div>
        </div>
    );
}