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>
    );
}