import { ContentEditable } from "@common/components/ContentEditable";
import { highlight } from "@common/highlight";
import { useInputState } from "@common/hooks/useInputState";
import clsx from "clsx";
import { Check, ChevronsRight, Edit2, RefreshCw, Sparkles, Trash2, X } from "lucide-preact";
import { useCallback, useEffect, useMemo, useRef, useState } from "preact/hooks";
import styles from '../assets/chat-sidebar.module.css';
import sidebarStyles from '../assets/sidebar.module.css';
import { useAppState, type ChatMessage } from "../contexts/state";
import LLM from "../utils/llm";
import Prompt from "../utils/prompt";
import { Tools } from "../utils/tools";
import { useChapterSummarization } from "../utils/useChapterSummarization";
import { Sidebar } from "./sidebar";

const CONTINUE_PROMPT = "Continue the story naturally.\nUse `edit_text` tool in append mode to add new text to the story.\nWait for the approval after adding.\nNote: added text could be cropped due to limit, do not make any attempts to add it back.";

interface RoleHeaderProps {
    message: ChatMessage;
    chatMessages: ChatMessage[];
}

const RoleHeader = ({ message, chatMessages }: RoleHeaderProps) => {
    const toolName = useMemo(() => {
        if (message.role !== 'tool') return;
        for (const m of chatMessages.toReversed()) {
            if (m.role !== 'assistant') continue;
            const toolCall = m.tool_calls?.find(tc => tc.id === message.tool_call_id);
            if (toolCall) return toolCall.function.name;
        }
    }, [message, chatMessages]);

    return (
        <div class={styles.role}>
            {message.role}
            {toolName && (
                <span class={styles.toolBadge}>
                    {toolName}
                </span>
            )}
        </div>
    );
};

export const ChatPanel = () => {
    const appState = useAppState();
    const { currentWorld, currentStory, dispatch, connection, model, enableThinking, chatOpen } = appState;
    const { summarizeAll, isSummarizing } = useChapterSummarization();
    const [input, setInput] = useInputState('');
    const [isLoading, setIsLoading] = useState(false);
    const [tokenCount, setTokenCount] = useState<{ taken: number; total: number } | null>(null);

    const [error, setError] = useState<string | null>(null);
    const messagesRef = useRef<HTMLDivElement>(null);
    const abortControllerRef = useRef<AbortController>(new AbortController());

    // Edit state
    const [editingMessageId, setEditingMessageId] = useState<string | null>(null);
    const [editingContent, setEditingContent] = useInputState('');

    const appStateRef = useRef(appState);
    appStateRef.current = appState;

    const lastMessage = currentStory?.chatMessages.at(-1);

    useEffect(() => {
        if (messagesRef.current) {
            messagesRef.current.scrollTo({
                top: messagesRef.current.scrollHeight,
                behavior: 'smooth',
            });
        }
    }, [
        currentStory?.chatMessages.length,
        lastMessage?.content.length,
        lastMessage?.role === 'assistant' && lastMessage?.reasoning_content?.length,
    ]);

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

    useEffect(() => {
        return () => {
            abortControllerRef.current?.abort();
        };
    }, []);

    useEffect(() => {
        if (!currentStory || !connection || !model) {
            setTokenCount(null);
            return;
        }

        const countTokens = async () => {
            try {
                const messages: ChatMessage[] = [];

                if (input.trim()) {
                    messages.push({
                        id: crypto.randomUUID(),
                        role: 'user',
                        content: input.trim(),
                    });
                }

                const chatRequest = Prompt.compilePrompt(appStateRef.current, messages);
                const countRequest: LLM.CountTokensRequest = {
                    model: model.id,
                    input: chatRequest?.messages ?? [],
                    tools: chatRequest?.tools,
                    enable_thinking: chatRequest?.enable_thinking,
                };

                const response = await LLM.countTokens(connection, countRequest);

                setTokenCount({
                    taken: response.input_tokens,
                    total: model.max_context ?? response.input_tokens,
                });
            } catch {
                setTokenCount(null);
            }
        };

        const timeoutId = setTimeout(countTokens, 300);
        return () => clearTimeout(timeoutId);
    }, [currentStory, connection, model, input, currentStory?.chatMessages.length]);

    const sendMessage = useCallback(async (
        newMessages: Iterable<ChatMessage>,
        excludedMessageIds: string[] = [],
    ) => {
        if (!currentStory || !currentWorld || !connection || !model) return;

        for (const message of newMessages) {
            if (excludedMessageIds.includes(message.id)) continue;
            dispatch({
                type: 'ADD_CHAT_MESSAGE',
                worldId: currentWorld.id,
                storyId: currentStory.id,
                message,
            });
        }

        const assistantMessageId = crypto.randomUUID();
        dispatch({
            type: 'ADD_CHAT_MESSAGE',
            worldId: currentWorld.id,
            storyId: currentStory.id,
            message: {
                id: assistantMessageId,
                role: 'assistant',
                content: '',
                reasoning_content: 'Generating...',
            },
        });

        const request = Prompt.compilePrompt(appStateRef.current, newMessages, excludedMessageIds);

        if (!request) {
            setError('Failed to compile prompt');
            setIsLoading(false);
            return;
        }

        try {
            const charName = currentWorld.title ?? 'Assistant';
            const prefix = `${charName}: `;
            let accumulatedContent = '';
            let accumulatedReasoning = '';
            let tool_calls: LLM.ToolCall[] | undefined;

            for await (const chunk of LLM.generateStream(connection, request)) {
                if (abortControllerRef.current?.signal.aborted) {
                    break;
                }
                const delta = chunk.choices[0]?.delta;

                if (delta?.tool_calls) {
                    tool_calls = delta.tool_calls;
                }

                const content = delta?.content;
                if (content) {
                    accumulatedContent += content;
                    if (currentWorld?.chatOnly && accumulatedContent.trimStart().startsWith(prefix)) {
                        accumulatedContent = accumulatedContent.trimStart().slice(prefix.length).trimStart();
                    }
                }
                const reasoningContent = delta?.reasoning_content;
                if (reasoningContent) {
                    accumulatedReasoning += reasoningContent;
                }
                if (content || reasoningContent) {
                    dispatch({
                        type: 'ADD_CHAT_MESSAGE',
                        worldId: currentWorld.id,
                        storyId: currentStory.id,
                        message: {
                            id: assistantMessageId,
                            role: 'assistant',
                            content: accumulatedContent,
                            reasoning_content: accumulatedReasoning,
                            tool_calls,
                        },
                    });
                }
            }
            const assistantMessage: ChatMessage = {
                id: assistantMessageId,
                role: 'assistant',
                content: accumulatedContent,
                reasoning_content: accumulatedReasoning,
                tool_calls,
            };
            dispatch({
                type: 'ADD_CHAT_MESSAGE',
                worldId: currentWorld.id,
                storyId: currentStory.id,
                message: assistantMessage,
            });

            if (tool_calls) {
                const toolMessages: ChatMessage[] = [];
                for (const tool of tool_calls) {
                    if (abortControllerRef.current?.signal.aborted) {
                        break;
                    }
                    const content = await Tools.executeTool(appStateRef.current, tool);
                    const message: ChatMessage = {
                        id: crypto.randomUUID(),
                        role: 'tool',
                        content,
                        tool_call_id: tool.id,
                    };
                    dispatch({
                        type: 'ADD_CHAT_MESSAGE',
                        worldId: currentWorld.id,
                        storyId: currentStory.id,
                        message,
                    });
                    toolMessages.push(message);
                }

                if (!abortControllerRef.current?.signal.aborted) {
                    return sendMessage([...newMessages, assistantMessage, ...toolMessages]);
                }
            }
        } catch (err) {
            const errorMessage = err instanceof Error
                ? err.message
                : 'Failed to generate response';

            setError(errorMessage);
        }
    }, [currentStory, connection, model]);

    const handleRegenerate = useCallback(async () => {
        if (!currentStory || !connection || !model || isLoading) return;

        // Only regenerate if last message is assistant
        const lastMessage = currentStory.chatMessages.at(-1);
        const isAssistant = lastMessage?.role === 'assistant';

        setIsLoading(true);
        setError(null);
        abortControllerRef.current = new AbortController();

        const excludedMessages: string[] = [];
        try {
            if (isAssistant) {
                // Delete the last assistant message and regenerate
                dispatch({
                    type: 'DELETE_CHAT_MESSAGES_FROM',
                    worldId: currentWorld!.id,
                    storyId: currentStory.id,
                    messageId: lastMessage.id,
                });
                excludedMessages.push(lastMessage.id);
            }
            await sendMessage([], excludedMessages);
        } finally {
            setIsLoading(false);
        }
    }, [currentStory, currentWorld, connection, model, isLoading, sendMessage]);

    const handleSendMessage = useCallback(async () => {
        if (!currentStory || !connection || !model || isLoading) return;

        if (!input.trim()) {
            if (lastMessage?.role === 'user') {
                handleRegenerate();
            }
            return;
        }

        setInput('');
        setIsLoading(true);
        setError(null);
        abortControllerRef.current = new AbortController();

        try {
            await sendMessage([{
                id: crypto.randomUUID(),
                role: 'user' as const,
                content: input.trim(),
            }]);
        } finally {
            setIsLoading(false);
        }
    }, [currentStory, input, connection, model, isLoading, sendMessage, handleRegenerate]);

    const handleContinue = useCallback(async () => {
        if (!currentStory || !connection || !model || isLoading) return;

        setInput('');
        setIsLoading(true);
        setError(null);
        abortControllerRef.current = new AbortController();

        try {
            await sendMessage([{
                id: crypto.randomUUID(),
                role: 'user' as const,
                content: (CONTINUE_PROMPT + '\n\n' + input).trim(),
            }]);
        } finally {
            setIsLoading(false);
        }
    }, [currentStory, input, connection, model, isLoading, sendMessage]);

    const handleStopGeneration = useCallback(() => {
        abortControllerRef.current?.abort();
        setIsLoading(false);
    }, []);

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

    const handleClear = () => {
        if (!currentStory || !currentWorld) return;
        dispatch({
            type: 'CLEAR_CHAT',
            worldId: currentWorld.id,
            storyId: currentStory.id,
        });
    };

    const handleDeleteMessage = useCallback((messageId: string) => {
        if (!currentStory || !currentWorld) return;
        dispatch({
            type: 'DELETE_CHAT_MESSAGES_FROM',
            worldId: currentWorld.id,
            storyId: currentStory.id,
            messageId,
        });
    }, [currentStory, currentWorld, dispatch]);

    const handleStartEdit = useCallback((message: ChatMessage) => {
        setEditingMessageId(message.id);
        setEditingContent(message.content);
    }, [setEditingContent]);

    const handleSaveEdit = useCallback(() => {
        if (!currentStory || !currentWorld || !editingMessageId) return;

        dispatch({
            type: 'EDIT_CHAT_MESSAGE',
            worldId: currentWorld.id,
            storyId: currentStory.id,
            messageId: editingMessageId,
            content: editingContent.trim(),
        });

        setEditingMessageId(null);
        setEditingContent('');
    }, [currentStory, currentWorld, editingMessageId, editingContent, dispatch, setEditingContent]);

    const handleCancelEdit = useCallback(() => {
        setEditingMessageId(null);
        setEditingContent('');
    }, [setEditingContent]);

    const isDisabled = !currentStory || !connection || !model || isLoading;

    return (
        <div class={styles.chat}>
            {!currentStory ? (
                <div class={styles.placeholder}>
                    Select a chat to start
                </div>
            ) : !connection || !model ? (
                <div class={styles.placeholder}>
                    {!connection ? 'Connect to an LLM server' : 'Select a model'} to start chatting
                </div>
            ) : currentStory.chatMessages.length === 0 ? (
                <div class={styles.placeholder}>
                    No messages yet
                </div>
            ) : (
                <div class={styles.messages} ref={messagesRef}>
                    {currentStory.chatMessages.map((message) => {
                        const isEditing = editingMessageId === message.id;
                        const canEdit = message.role === 'user' || message.role === 'assistant';

                        return (
                            <div key={message.id} class={styles.message} data-role={message.role}>
                                <div class={styles.messageHeader}>
                                    <RoleHeader message={message} chatMessages={currentStory.chatMessages} />

                                    {!isLoading && canEdit && (
                                        <div class={styles.messageActions}>
                                            <button
                                                class={styles.iconButton}
                                                onClick={() => handleStartEdit(message)}
                                                title="Edit message"
                                            >
                                                <Edit2 size={12} />
                                            </button>
                                            <button
                                                class={styles.iconButton}
                                                onClick={() => handleDeleteMessage(message.id)}
                                                title="Delete this and all following messages"
                                            >
                                                <Trash2 size={12} />
                                            </button>
                                        </div>
                                    )}
                                </div>

                                {isEditing ? (
                                    <div class={styles.editContainer}>
                                        <ContentEditable
                                            autoLines
                                            class={styles.editTextarea}
                                            value={highlight(editingContent)}
                                            onInput={setEditingContent}
                                            onKeyDown={(e) => {
                                                if (e.key === 'Enter' && e.ctrlKey) {
                                                    e.preventDefault();
                                                    handleSaveEdit();
                                                }
                                                if (e.key === 'Escape') {
                                                    e.preventDefault();
                                                    handleCancelEdit();
                                                }
                                            }}
                                        />
                                        <div class={styles.editActions}>
                                            <button
                                                class={clsx(styles.iconButton, styles.saveButton)}
                                                onClick={handleSaveEdit}
                                                disabled={!editingContent.trim()}
                                                title="Save (Ctrl+Enter)"
                                            >
                                                <Check size={14} />
                                            </button>
                                            <button
                                                class={clsx(styles.iconButton, styles.cancelButton)}
                                                onClick={handleCancelEdit}
                                                title="Cancel (Esc)"
                                            >
                                                <X size={14} />
                                            </button>
                                        </div>
                                    </div>
                                ) : (
                                    <>
                                        {message.role === 'assistant' && message.reasoning_content && (
                                            <div class={styles.reasoningContent}>
                                                {message.reasoning_content}
                                            </div>
                                        )}

                                        <div
                                            class={styles.content}
                                            dangerouslySetInnerHTML={{ __html: highlight(message.content, false).trim() }}
                                        />

                                        {message.role === 'assistant' && message.tool_calls && (
                                            <div class={styles.toolCalls}>
                                                {message.tool_calls.map((tool) => (
                                                    <span key={tool.id} class={styles.toolBadge}>
                                                        {tool.function.name}
                                                    </span>
                                                ))}
                                            </div>
                                        )}
                                    </>
                                )}
                            </div>
                        );
                    })}
                    {error && (
                        <div class={clsx(styles.message, styles.errorMessage)} data-role="assistant">
                            <div class={styles.role}>error</div>
                            <div class={styles.errorText}>{error}</div>
                        </div>
                    )}
                    <button class={styles.clearButton} onClick={handleClear}>
                        Clear chat
                    </button>
                </div>
            )}
            {currentStory && (
                <div class={styles.inputContainer}>
                    <div class={styles.optionsRow}>
                        <label class={styles.toggleContainer}>
                            {!currentWorld?.chatOnly && (<>
                                <input
                                    type="checkbox"
                                    checked={enableThinking}
                                    onChange={(e) => dispatch({
                                        type: 'SET_ENABLE_THINKING',
                                        enable: (e.target as HTMLInputElement).checked,
                                    })}
                                    disabled={isDisabled}
                                />
                                <span>Enable thinking</span>
                            </>)}
                        </label>
                        <div class={styles.tokenCounter}>
                            {tokenCount && <span>{tokenCount.taken} / {tokenCount.total} tokens</span>}

                            {!currentWorld?.chatOnly && (<button
                                class={styles.summarizeButton}
                                onClick={summarizeAll}
                                disabled={isSummarizing || !currentStory || !connection || !model}
                                title={isSummarizing ? 'Summarizing...' : 'Summarize'}>
                                <Sparkles size={14} />
                            </button>
                            )}
                        </div>
                    </div>
                    <ContentEditable
                        autoLines
                        class={styles.input}
                        value={highlight(input)}
                        onInput={setInput}
                        onKeyDown={handleKeyDown}
                        placeholder={
                            isLoading
                                ? 'Generating...'
                                : isDisabled
                                    ? 'Connect to an LLM server to chat'
                                    : 'Type a message...'}
                        disabled={isDisabled}
                    />
                    {isLoading ? (
                        <button
                            class={styles.stopButton}
                            onClick={handleStopGeneration}
                        >
                            Stop
                        </button>
                    ) : (
                        <div class={styles.buttonRow}>
                            <button
                                class={styles.sendButton}
                                onClick={handleSendMessage}
                                disabled={isDisabled}
                            >
                                Send
                            </button>
                            {!currentWorld?.chatOnly && (
                                <button
                                    class={styles.continueButton}
                                    onClick={handleContinue}
                                    disabled={isDisabled}
                                    title="Continue"
                                >
                                    <ChevronsRight size={14} />
                                </button>
                            )}
                            <button
                                class={styles.regenerateButton}
                                onClick={handleRegenerate}
                                disabled={isDisabled}
                                title="Regenerate last response"
                            >
                                <RefreshCw size={14} />
                            </button>
                        </div>
                    )}
                </div>
            )}
        </div>
    );
};

export const ChatSidebar = () => {
    const { currentWorld, chatOpen, dispatch } = useAppState();

    // In chat-only worlds, chat is a full editor tab — no sidebar needed
    if (currentWorld?.chatOnly) return null;

    return (
        <Sidebar
            side="right"
            open={chatOpen}
            onToggle={() => dispatch({
                type: 'SET_CHAT_OPEN',
                open: !chatOpen,
            })}
            class={sidebarStyles.mobileOverlay}
        >
            <ChatPanel />
        </Sidebar>
    );
};