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

const CONTINUE_PROMPT = "Continue the story forward.\nUse `edit_text` tool in append mode to add new text to the story.\nWait for the approval after adding.";

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 ChatSidebar = () => {
    const appState = useAppState();
    const { 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());

    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: ChatMessage[]) => {
        if (!currentStory || !connection || !model) return;

        for (const message of newMessages) {
            dispatch({
                type: 'ADD_CHAT_MESSAGE',
                storyId: currentStory.id,
                message,
            });
        }

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

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

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

        try {
            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;
                }
                const reasoningContent = delta?.reasoning_content;
                if (reasoningContent) {
                    accumulatedReasoning += reasoningContent;
                }
                if (content || reasoningContent) {
                    dispatch({
                        type: 'ADD_CHAT_MESSAGE',
                        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',
                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',
                        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 handleSendMessage = useCallback(async () => {
        if (!currentStory || !input.trim() || !connection || !model || isLoading) 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]);

    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) return;
        dispatch({
            type: 'CLEAR_CHAT',
            storyId: currentStory.id,
        });
    };

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

    return (
        <Sidebar side="right" open={chatOpen} onToggle={() => dispatch({ type: 'SET_CHAT_OPEN', open: !chatOpen })} class={sidebarStyles.mobileOverlay}>
            <div class={styles.chat}>
                {!currentStory ? (
                    <div class={styles.placeholder}>
                        Select a story to start chatting
                    </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) => (
                            <div key={message.id} class={styles.message} data-role={message.role}>
                                <RoleHeader message={message} chatMessages={currentStory.chatMessages} />

                                {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}>
                                <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>}
                                <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={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 || !input.trim()}
                                >
                                    Send
                                </button>
                                <button
                                    class={styles.continueButton}
                                    onClick={handleContinue}
                                    disabled={isDisabled}
                                    title="Continue"
                                >
                                    <ChevronsRight size={14} />
                                </button>
                            </div>
                        )}
                    </div>
                )}
            </div>
        </Sidebar>
    );
};
