src/components/chat.tsx

"use client";

import { useState, useRef, useEffect } from "react";
import { Globe, Send, Loader2, Plus, ChevronDown, ChevronRight } from "lucide-react";
import { Markdown } from "@/components/markdown";
import { cn } from "@/lib/utils";

interface ReasoningBlock {
    content: string;
    duration?: number;
}

interface ToolInvocation {
    toolName: string;
    args: Record<string, unknown>;
    result?: string;
    status: "pending" | "running" | "complete" | "error";
}

// Each step can be a reasoning block or a tool invocation
type MessageStep =
    | { type: "reasoning"; block: ReasoningBlock }
    | { type: "tool"; invocation: ToolInvocation };

interface Message {
    id: string;
    role: "user" | "assistant";
    content: string;
    steps: MessageStep[];
}

function ThinkingBlock({
    content,
    duration,
    isStreaming
}: {
    content: string;
    duration?: number;
    isStreaming?: boolean;
}) {
    const [isOpen, setIsOpen] = useState(isStreaming);

    useEffect(() => {
        if (isStreaming) setIsOpen(true);
    }, [isStreaming]);

    const formatDuration = (s?: number) => {
        if (!s) return "";
        if (s < 1) return " · <1s";
        return ` · ${Math.round(s)}s`;
    };

    return (
        <div className="thinking-block">
            <button
                className="thinking-header w-full"
                onClick={() => setIsOpen(!isOpen)}
            >
                {isOpen ? <ChevronDown className="w-4 h-4" /> : <ChevronRight className="w-4 h-4" />}
                <span>
                    {isStreaming ? "Thinking..." : `Thought process${formatDuration(duration)}`}
                </span>
                {isStreaming && (
                    <span className="ml-auto flex gap-1">
                        <span className="w-1.5 h-1.5 rounded-full bg-current dot-1" />
                        <span className="w-1.5 h-1.5 rounded-full bg-current dot-2" />
                        <span className="w-1.5 h-1.5 rounded-full bg-current dot-3" />
                    </span>
                )}
            </button>
            {isOpen && (
                <div className="thinking-content">
                    <Markdown>{content}</Markdown>
                </div>
            )}
        </div>
    );
}

function ToolBlock({
    toolName,
    args,
    result,
    status
}: {
    toolName: string;
    args: Record<string, unknown>;
    result?: string;
    status: string;
}) {
    const [isOpen, setIsOpen] = useState(false);
    const isRunning = status === "running" || status === "pending";
    const query = (args.query as string) || JSON.stringify(args);

    return (
        <div className="tool-block">
            <button
                className="tool-header w-full"
                onClick={() => setIsOpen(!isOpen)}
            >
                {isRunning ? (
                    <Loader2 className="w-4 h-4 animate-spin text-blue-500" />
                ) : (
                    <Globe className="w-4 h-4 text-green-500" />
                )}
                <span className="font-medium">
                    {toolName === "web_search" ? "Web Search" : toolName}
                </span>
                <span className="text-muted-foreground truncate flex-1 text-left ml-1">
                    {query}
                </span>
                {isOpen ? <ChevronDown className="w-4 h-4 shrink-0" /> : <ChevronRight className="w-4 h-4 shrink-0" />}
            </button>
            {isOpen && result && (
                <div className="tool-content">{result}</div>
            )}
        </div>
    );
}

export function Chat() {
    const [messages, setMessages] = useState<Message[]>([]);
    const [input, setInput] = useState("");
    const [isLoading, setIsLoading] = useState(false);
    const [searchEnabled, setSearchEnabled] = useState(false);
    const [streamingMessage, setStreamingMessage] = useState<Message | null>(null);
    const messagesEndRef = useRef<HTMLDivElement>(null);
    const textareaRef = useRef<HTMLTextAreaElement>(null);
    const messagesRef = useRef<Message[]>(messages);
    const searchEnabledRef = useRef(searchEnabled);

    const makeId = () => {
        try {
            return globalThis.crypto?.randomUUID?.() ?? `${Date.now()}-${Math.random().toString(16).slice(2)}`;
        } catch {
            return `${Date.now()}-${Math.random().toString(16).slice(2)}`;
        }
    };

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

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

    useEffect(() => {
        messagesRef.current = messages;
    }, [messages]);

    useEffect(() => {
        searchEnabledRef.current = searchEnabled;
    }, [searchEnabled]);

    useEffect(() => {
        const textarea = textareaRef.current;
        if (textarea) {
            textarea.style.height = "auto";
            textarea.style.height = Math.min(textarea.scrollHeight, 160) + "px";
        }
    }, [input]);

    const handleSubmit = async (e: React.FormEvent) => {
        e.preventDefault();
        if (!input.trim() || isLoading) return;

        const userMessage: Message = {
            id: makeId(),
            role: "user",
            content: input.trim(),
            steps: [],
        };

        setMessages((prev) => [...prev, userMessage]);
        setInput("");
        setIsLoading(true);

        const assistantMessage: Message = {
            id: makeId(),
            role: "assistant",
            content: "",
            steps: [],
        };
        setStreamingMessage(assistantMessage);

        try {
            const messagesSnapshot = messagesRef.current;
            const searchEnabledSnapshot = searchEnabledRef.current;
            const response = await fetch("/api/chat", {
                method: "POST",
                headers: { "Content-Type": "application/json" },
                body: JSON.stringify({
                    messages: [...messagesSnapshot, userMessage].map((m) => ({
                        role: m.role,
                        content: m.content,
                    })),
                    searchEnabled: searchEnabledSnapshot,
                }),
            });

            if (!response.ok) throw new Error("Failed to get response");

            const reader = response.body?.getReader();
            if (!reader) throw new Error("No reader available");

            const decoder = new TextDecoder();
            let buffer = "";
            let currentContent = "";
            let reasoningStart = 0;
            let inReasoning = false;
            let currentReasoningIdx = -1;

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

                buffer += decoder.decode(value, { stream: true });
                const lines = buffer.split("\n");
                buffer = lines.pop() || "";

                for (const line of lines) {
                    if (!line.startsWith("data: ")) continue;
                    const data = line.slice(6);
                    if (data === "[DONE]") continue;

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

                        if (parsed.type === "reasoning_start") {
                            // Finalize previous reasoning block if any
                            if (inReasoning && reasoningStart > 0) {
                                const duration = (Date.now() - reasoningStart) / 1000;
                                setStreamingMessage((prev) => {
                                    if (!prev || currentReasoningIdx < 0) return prev;
                                    const steps = [...prev.steps];
                                    const step = steps[currentReasoningIdx];
                                    if (step?.type === "reasoning") {
                                        steps[currentReasoningIdx] = {
                                            type: "reasoning",
                                            block: { ...step.block, duration },
                                        };
                                    }
                                    return { ...prev, steps };
                                });
                            }
                            // Reset for new reasoning block
                            inReasoning = false;
                            currentReasoningIdx = -1;
                        } else if (parsed.type === "reasoning") {
                            if (!inReasoning) {
                                inReasoning = true;
                                reasoningStart = Date.now();
                                // Add new reasoning step
                                setStreamingMessage((prev) => {
                                    if (!prev) return prev;
                                    const newStep: MessageStep = {
                                        type: "reasoning",
                                        block: { content: parsed.content },
                                    };
                                    currentReasoningIdx = prev.steps.length;
                                    return { ...prev, steps: [...prev.steps, newStep] };
                                });
                            } else {
                                // Append to current reasoning block
                                setStreamingMessage((prev) => {
                                    if (!prev || currentReasoningIdx < 0) return prev;
                                    const steps = [...prev.steps];
                                    const step = steps[currentReasoningIdx];
                                    if (step?.type === "reasoning") {
                                        steps[currentReasoningIdx] = {
                                            type: "reasoning",
                                            block: { ...step.block, content: step.block.content + parsed.content },
                                        };
                                    }
                                    return { ...prev, steps };
                                });
                            }
                        } else if (parsed.type === "content") {
                            // Finalize reasoning if we were in one
                            if (inReasoning && reasoningStart > 0) {
                                const duration = (Date.now() - reasoningStart) / 1000;
                                setStreamingMessage((prev) => {
                                    if (!prev || currentReasoningIdx < 0) return prev;
                                    const steps = [...prev.steps];
                                    const step = steps[currentReasoningIdx];
                                    if (step?.type === "reasoning") {
                                        steps[currentReasoningIdx] = {
                                            type: "reasoning",
                                            block: { ...step.block, duration },
                                        };
                                    }
                                    return { ...prev, steps };
                                });
                                inReasoning = false;
                            }
                            currentContent += parsed.content;
                            setStreamingMessage((prev) =>
                                prev ? { ...prev, content: currentContent } : prev
                            );
                        } else if (parsed.type === "tool_call") {
                            setStreamingMessage((prev) => {
                                if (!prev) return prev;
                                // Find existing tool step by name
                                const existingIdx = prev.steps.findIndex(
                                    (s) => s.type === "tool" && s.invocation.toolName === parsed.name
                                );
                                if (existingIdx >= 0) {
                                    const steps = [...prev.steps];
                                    const step = steps[existingIdx];
                                    if (step.type === "tool") {
                                        steps[existingIdx] = {
                                            type: "tool",
                                            invocation: {
                                                ...step.invocation,
                                                status: parsed.status,
                                                result: parsed.result,
                                            },
                                        };
                                    }
                                    return { ...prev, steps };
                                }
                                // Add new tool step
                                const newStep: MessageStep = {
                                    type: "tool",
                                    invocation: {
                                        toolName: parsed.name,
                                        args: parsed.args,
                                        status: parsed.status,
                                        result: parsed.result,
                                    },
                                };
                                return { ...prev, steps: [...prev.steps, newStep] };
                            });
                        }
                    } catch {
                        // Ignore parse errors
                    }
                }
            }

            // Finalize reasoning duration if stream ended while still in reasoning
            if (inReasoning && reasoningStart > 0) {
                const duration = (Date.now() - reasoningStart) / 1000;
                setStreamingMessage((prev) => {
                    if (!prev || currentReasoningIdx < 0) return prev;
                    const steps = [...prev.steps];
                    const step = steps[currentReasoningIdx];
                    if (step?.type === "reasoning") {
                        steps[currentReasoningIdx] = {
                            type: "reasoning",
                            block: { ...step.block, duration },
                        };
                    }
                    return { ...prev, steps };
                });
            }

            // Get the final message and add to messages list
            setStreamingMessage((prev) => {
                if (prev) {
                    // Use setTimeout to avoid nested state updates
                    setTimeout(() => {
                        setMessages((msgs) => [...msgs, prev]);
                    }, 0);
                }
                return null;
            });
        } catch (error) {
            console.error("Chat error:", error);
            setStreamingMessage(null);
        } finally {
            setIsLoading(false);
        }
    };

    const handleKeyDown = (e: React.KeyboardEvent) => {
        if (e.key === "Enter" && !e.shiftKey) {
            e.preventDefault();
            handleSubmit(e);
        }
    };

    const clearChat = () => {
        setMessages([]);
        setStreamingMessage(null);
    };

    const hasMessages = messages.length > 0 || streamingMessage;

    const renderMessage = (message: Message, isStreaming = false) => {
        const lastStep = message.steps[message.steps.length - 1];
        const isStreamingReasoning = isStreaming && lastStep?.type === "reasoning" && !lastStep.block.duration;
        const hasAnyContent = message.steps.length > 0 || message.content;

        return (
            <div className="space-y-3">
                {message.steps.map((step, idx) => {
                    if (step.type === "reasoning") {
                        const isLastReasoning = idx === message.steps.length - 1;
                        return (
                            <ThinkingBlock
                                key={idx}
                                content={step.block.content}
                                duration={step.block.duration}
                                isStreaming={isStreaming && isLastReasoning && !step.block.duration}
                            />
                        );
                    } else {
                        return (
                            <ToolBlock
                                key={idx}
                                toolName={step.invocation.toolName}
                                args={step.invocation.args}
                                result={step.invocation.result}
                                status={step.invocation.status}
                            />
                        );
                    }
                })}
                {message.content && (
                    <Markdown>{message.content}</Markdown>
                )}
                {isStreaming && !hasAnyContent && (
                    <div className="flex items-center gap-2 py-2 text-muted-foreground">
                        <span className="flex gap-1">
                            <span className="w-2 h-2 rounded-full bg-current dot-1" />
                            <span className="w-2 h-2 rounded-full bg-current dot-2" />
                            <span className="w-2 h-2 rounded-full bg-current dot-3" />
                        </span>
                    </div>
                )}
            </div>
        );
    };

    return (
        <div className="flex flex-col h-screen bg-background">
            {/* Header */}
            <header className="shrink-0 border-b border-[#27272a] bg-[#09090b]">
                <div className="max-w-3xl mx-auto px-4 sm:px-6 h-14 flex items-center justify-between">
                    <div className="flex items-center gap-2.5">
                        <img src="/teich.png" alt="Teich AI" className="w-8 h-8" />
                        <span className="font-semibold text-[#fafafa]">Qwen3-4B-Thinking-2507-Claude-4.5-Opus</span>
                    </div>
                    {hasMessages && (
                        <button
                            onClick={clearChat}
                            className="flex items-center gap-1.5 px-3 py-1.5 text-sm text-[#a1a1aa] hover:text-[#fafafa] hover:bg-[#27272a] rounded-md transition-colors"
                        >
                            <Plus className="w-4 h-4" />
                            New Chat
                        </button>
                    )}
                </div>
            </header>

            {/* Messages */}
            <main className="flex-1 overflow-y-auto">
                <div className="max-w-3xl mx-auto px-4 sm:px-6 py-6">
                    {!hasMessages && (
                        <div className="flex flex-col items-center justify-center min-h-[50vh] text-center animate-fade-in">
                            <h2 className="text-xl font-semibold mb-2 text-[#fafafa]">How can I help you?</h2>
                            <p className="text-[#a1a1aa] text-sm max-w-sm mb-8">
                                A reasoning model that thinks step by step. Enable web search for real-time information.
                            </p>
                            <div className="flex flex-wrap gap-2 justify-center">
                                {["Explain quantum computing", "Write a Python function", "What's in the news?"].map((prompt) => (
                                    <button
                                        key={prompt}
                                        onClick={() => setInput(prompt)}
                                        className="px-3 py-1.5 text-sm border border-[#27272a] text-[#a1a1aa] rounded-md hover:bg-[#18181b] hover:text-[#fafafa] hover:border-[#3f3f46] transition-colors"
                                    >
                                        {prompt}
                                    </button>
                                ))}
                            </div>
                        </div>
                    )}

                    <div className="space-y-6">
                        {messages.map((message) => (
                            <div key={message.id} className="animate-slide-in">
                                {message.role === "user" ? (
                                    <div className="flex justify-end">
                                        <div className="user-message px-4 py-2.5 rounded-2xl rounded-br-sm max-w-[85%]">
                                            <p className="text-sm whitespace-pre-wrap">{message.content}</p>
                                        </div>
                                    </div>
                                ) : (
                                    renderMessage(message)
                                )}
                            </div>
                        ))}

                        {streamingMessage && (
                            <div className="animate-fade-in">
                                {renderMessage(streamingMessage, true)}
                            </div>
                        )}
                    </div>

                    <div ref={messagesEndRef} className="h-4" />
                </div>
            </main>

            {/* Input */}
            <footer className="shrink-0 border-t border-[#27272a] bg-[#09090b]">
                <div className="max-w-3xl mx-auto px-4 sm:px-6 py-4">
                    <form onSubmit={handleSubmit}>
                        <div className="input-container flex items-end gap-2 p-2">
                            <button
                                type="button"
                                onClick={() => setSearchEnabled((v) => !v)}
                                className={cn(
                                    "shrink-0 w-9 h-9 rounded-lg flex items-center justify-center transition-colors",
                                    searchEnabled
                                        ? "bg-[#ff4c00] text-white"
                                        : "text-[#a1a1aa] hover:text-[#fafafa] hover:bg-[#27272a]"
                                )}
                                title={searchEnabled ? "Disable web search" : "Enable web search"}
                            >
                                <Globe className="w-4 h-4" />
                            </button>

                            <textarea
                                ref={textareaRef}
                                value={input}
                                onChange={(e) => setInput(e.target.value)}
                                onKeyDown={handleKeyDown}
                                placeholder="Send a message..."
                                rows={1}
                                disabled={isLoading}
                                className="flex-1 bg-transparent border-0 resize-none text-sm leading-6 text-[#fafafa] placeholder:text-[#71717a] focus:outline-none min-h-[36px] max-h-[160px] py-1.5"
                            />

                            <button
                                type="submit"
                                disabled={isLoading || !input.trim()}
                                className={cn(
                                    "shrink-0 w-9 h-9 rounded-lg flex items-center justify-center transition-colors",
                                    input.trim() && !isLoading
                                        ? "bg-[#ff4c00] text-white hover:bg-[#ff6a2a]"
                                        : "text-[#a1a1aa] cursor-not-allowed"
                                )}
                            >
                                {isLoading ? (
                                    <Loader2 className="w-4 h-4 animate-spin" />
                                ) : (
                                    <Send className="w-4 h-4" />
                                )}
                            </button>
                        </div>

                        <div className="flex items-center justify-center gap-3 mt-2 text-xs text-[#a1a1aa]">
                            {searchEnabled && (
                                <span className="flex items-center gap-1.5">
                                    <span className="w-1.5 h-1.5 bg-[#ff4c00] rounded-full" />
                                    Web search on
                                </span>
                            )}
                            <span>Enter to send · Shift+Enter for newline</span>
                        </div>

                        <div className="mt-2 text-center text-[11px] leading-4 text-[#71717a]">
                            Inputs and outputs are logged anonymously for debugging/usage analytics. We do not track who sent messages.
                        </div>
                    </form>
                </div>
            </footer>
        </div>
    );
}