frontend/components/Chat/ChatLayout.tsx

"use client";

import { useState, useEffect } from "react";
import Sidebar from "./Sidebar";
import MessageBubble from "./MessageBubble";
import ChatInput from "./ChatInput";
import { Menu } from "lucide-react";

interface Message {
    role: "user" | "assistant";
    content: string;
    sources?: any[];
    intent?: string;
    thought_process?: string;
}

interface Session {
    id: string;
    title: string;
    date: string;
}

export default function ChatLayout() {
    const [isSidebarOpen, setIsSidebarOpen] = useState(true);
    const [query, setQuery] = useState("");
    const [attachedFile, setAttachedFile] = useState<File | null>(null);
    const [messages, setMessages] = useState<Message[]>([]);
    const [isLoading, setIsLoading] = useState(false);
    const [sessions, setSessions] = useState<Session[]>([]);
    const [currentSessionId, setCurrentSessionId] = useState<string | null>(null);

    // Load history from local storage on mount
    useEffect(() => {
        const saved = localStorage.getItem("chat_sessions");
        if (saved) {
            setSessions(JSON.parse(saved));
        }
    }, []);

    const handleNewChat = () => {
        setMessages([]);
        setCurrentSessionId(null);
    };

    const handleSearch = async () => {
        if (!query.trim()) return;

        // Add User Message
        const userMsg: Message = { role: "user", content: query };
        setMessages((prev) => [...prev, userMsg]);
        setQuery("");
        setIsLoading(true);

        try {
            let finalQuery = userMsg.content;

            // Upload File if exists
            if (attachedFile) {
                const formData = new FormData();
                formData.append("file", attachedFile);

                // Optimistic UI update: Show file as uploaded
                setMessages((prev) => [...prev, { role: "assistant", content: `📂 Analyzing ${attachedFile.name}...`, isLoading: true }]);

                const uploadRes = await fetch(`${process.env.NEXT_PUBLIC_API_BASE || "http://localhost:8000"}/api/v1/upload`, {
                    method: "POST",
                    body: formData,
                });
                const uploadData = await uploadRes.json();

                // Remove the loading message logic would be complex, simplfying:
                // Append file context to query
                finalQuery = `[Context from uploaded file ${attachedFile.name}]:\n${uploadData.content}\n\nUser Question: ${userMsg.content}`;
                setAttachedFile(null); // Clear file
            }

            const response = await fetch(`${process.env.NEXT_PUBLIC_API_BASE || "http://localhost:8000"}/api/v1/query`, {
                method: "POST",
                headers: { "Content-Type": "application/json" },
                body: JSON.stringify({ query: finalQuery }),
            });

            const data = await response.json();

            // Add Assistant Message
            const aiMsg: Message = {
                role: "assistant",
                content: data.answer,
                sources: data.sources,
                intent: data.intent,
                thought_process: data.thought_process,
            };
            setMessages((prev) => [...prev, aiMsg]);

            // Save Session (Simple Logic)
            if (!currentSessionId) {
                const newId = Date.now().toString();
                const newSession = {
                    id: newId,
                    title: userMsg.content.slice(0, 30) + "...",
                    date: new Date().toLocaleDateString(),
                };
                setSessions((prev) => [newSession, ...prev]);
                setCurrentSessionId(newId);
                localStorage.setItem("chat_sessions", JSON.stringify([newSession, ...sessions]));
            }

        } catch (error) {
            console.error("Error:", error);
            setMessages((prev) => [
                ...prev,
                { role: "assistant", content: "Sorry, something went wrong. Please try again." },
            ]);
        } finally {
            setIsLoading(false);
        }
    };

    const handleChallenge = async (msgIndex: number) => {
        const targetMsg = messages[msgIndex];
        if (!targetMsg || targetMsg.role !== "assistant") return;

        // Find the preceding user message for context (simple heuristic: index - 1)
        const userQuery = messages[msgIndex - 1]?.content || "Unknown context";
        const sourcesText = targetMsg.sources
            ? targetMsg.sources.map(s => `Title: ${s.title}\nContent: ${s.snippet}`).join("\n\n")
            : "No sources.";

        setIsLoading(true);

        try {
            const response = await fetch(`${process.env.NEXT_PUBLIC_API_BASE || "http://localhost:8000"}/api/v1/challenge`, {
                method: "POST",
                headers: { "Content-Type": "application/json" },
                body: JSON.stringify({
                    original_query: userQuery,
                    original_answer: targetMsg.content,
                    sources_text: sourcesText
                }),
            });

            const data = await response.json();

            // Add Critique Message
            const critiqueMsg: Message = {
                role: "assistant",
                content: data.answer,
                sources: [],
                intent: "CRITIQUE",
                thought_process: "Devil's Advocate Mode: Analyzing potential flaws in the previous answer."
            };
            setMessages((prev) => [...prev, critiqueMsg]);

        } catch (error) {
            console.error("Challenge Error:", error);
        } finally {
            setIsLoading(false);
        }
    };

    return (
        <div className="flex h-screen bg-white text-gray-900 font-sans overflow-hidden">
            {/* Sidebar */}
            <Sidebar
                isOpen={isSidebarOpen}
                sessions={sessions}
                currentSessionId={currentSessionId}
                onNewChat={handleNewChat}
                onSelectSession={(id) => console.log("Select session:", id)} // Placeholder for loading specific session
                onDeleteSession={(id, e) => {
                    e.stopPropagation();
                    const newSessions = sessions.filter(s => s.id !== id);
                    setSessions(newSessions);
                    localStorage.setItem("chat_sessions", JSON.stringify(newSessions));
                    if (currentSessionId === id) handleNewChat();
                }}
            />

            {/* Main Content */}
            <div className="flex-1 flex flex-col h-full relative">
                {/* Header / Mobile Toggle */}
                <div className="absolute top-4 left-4 z-20">
                    <button
                        onClick={() => setIsSidebarOpen(!isSidebarOpen)}
                        className="p-2 text-gray-400 hover:text-gray-900 rounded-lg hover:bg-gray-100 transition-colors"
                    >
                        <Menu size={20} />
                    </button>
                </div>

                {/* Chat Area */}
                <div className="flex-1 overflow-y-auto scroll-smooth">
                    {messages.length === 0 ? (
                        <div className="h-full flex flex-col items-center justify-center p-8 text-center opacity-50">
                            <div className="w-16 h-16 bg-gray-100 rounded-2xl flex items-center justify-center mb-6">
                                <span className="text-3xl">✨</span>
                            </div>
                            <h2 className="text-2xl font-semibold mb-2 text-gray-900">Trust-First Copilot</h2>
                            <p className="max-w-md text-gray-500">
                                Ask anything. I check reliable sources before answering.
                            </p>
                        </div>
                    ) : (
                        <div className="pb-32">
                            {messages.map((msg, idx) => (
                                <MessageBubble
                                    key={idx}
                                    role={msg.role}
                                    content={msg.content}
                                    sources={msg.sources}
                                    intent={msg.intent}
                                    thought_process={msg.thought_process}
                                    onChallenge={msg.role === "assistant" ? () => handleChallenge(idx) : undefined}
                                />
                            ))}
                            {isLoading && (
                                <MessageBubble role="assistant" content="" isLoading={true} />
                            )}
                        </div>
                    )}
                </div>

                {/* Input Area */}
                <div className="flex-shrink-0 bg-white pt-10">
                    <ChatInput
                        value={query}
                        onChange={setQuery}
                        onSubmit={handleSearch}
                        isLoading={isLoading}
                        onFileSelect={setAttachedFile}
                        attachedFile={attachedFile}
                    />
                </div>
            </div>
        </div>
    );
}