import { ContentEditable } from "@common/components/ContentEditable";
import { highlight } from "@common/highlight";
import { useInputState } from "@common/hooks/useInputState";
import { useMediaQuery } from "@common/hooks/useMediaQuery";
import clsx from "clsx";
import { Check, ChevronsRight, Edit2, GitFork, RefreshCw, SendHorizonal, Sparkles, Square, Trash2, X } from "lucide-preact";
import { useCallback, useEffect, useMemo, useRef, useState } from "preact/hooks";
import { MOBILE_BREAKPOINT } from '../assets/breakpoints';
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";


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

const RoleHeader = ({ message, chatMessages }: RoleHeaderProps) => {
    const { currentWorld, userName } = useAppState();

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

    let displayName: string = message.role;
    let roleLabel = null;

    if (message.role === 'tool' && toolName) {
        displayName = toolName;
        roleLabel = message.role;
    } else if (currentWorld?.chatOnly) {
        if (message.role === 'user' && userName) {
            displayName = userName;
            roleLabel = message.role;
        } else if (message.role === 'assistant' && currentWorld.title) {
            displayName = currentWorld.title;
            roleLabel = message.role;
        }
    }

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

export const ChatPanel = ({ visible }: { visible: boolean }) => {
    const appState = useAppState();
    const { currentWorld, currentStory, dispatch, connection, model, enableThinking, chatOpen, continuePrompt } = 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, visible]);

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

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

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

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

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

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

                setTokenCount({
                    taken: response.input_tokens,
                    total: model.context_length ?? 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 (config: Prompt.CompileConfig = {}) => {
        if (!currentStory || !currentWorld || !connection || !model) return;

        const { newMessages = [], excludedMessageIds = [] } = config;
        const excludedSet = new Set(excludedMessageIds);

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

        const continuedMessage = config.continueLast ? currentStory.chatMessages.at(-1) : null;
        const targetMessageId = continuedMessage?.id ?? crypto.randomUUID();
        const targetRole = continuedMessage?.role ?? 'assistant';

        if (!continuedMessage) {
            dispatch({
                type: 'ADD_CHAT_MESSAGE',
                worldId: currentWorld.id,
                storyId: currentStory.id,
                message: {
                    id: targetMessageId,
                    role: 'assistant',
                    content: '',
                    reasoning_content: 'Generating...',
                },
            });
        }

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

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

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

            for await (const chunk of LLM.generateStream(connection, request)) {
                if (abortControllerRef.current?.signal.aborted) {
                    break;
                }
                if ('error' in chunk) {
                    throw new Error(chunk.error);
                }
                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: targetMessageId,
                            role: targetRole,
                            content: accumulatedContent,
                            reasoning_content: accumulatedReasoning,
                            tool_calls,
                        },
                    });
                }
            }
            const finalMessage: ChatMessage = {
                id: targetMessageId,
                role: targetRole,
                content: accumulatedContent,
                reasoning_content: accumulatedReasoning,
                tool_calls,
            };
            dispatch({
                type: 'ADD_CHAT_MESSAGE',
                worldId: currentWorld.id,
                storyId: currentStory.id,
                message: finalMessage,
            });

            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: [
                            ...newMessages,
                            finalMessage,
                            ...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 excludedMessageIds = new Set<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,
                });
                excludedMessageIds.add(lastMessage.id);
            }
            await sendMessage({ excludedMessageIds });
        } 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({
                newMessages: [{
                    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 {
            if (currentWorld?.chatOnly) {
                await sendMessage({ continueLast: true });
            } else {
                await sendMessage({
                    newMessages: [{
                        id: crypto.randomUUID(),
                        role: 'user' as const,
                        content: (continuePrompt + '\n\n' + input).trim(),
                    }]
                });
            }
        } finally {
            setIsLoading(false);
        }
    }, [currentStory, currentWorld, 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 handleDeleteMessage = useCallback((messageId: string) => {
        if (!currentStory || !currentWorld) return;
        if (!confirm('Delete this and all following messages?')) return;
        dispatch({
            type: 'DELETE_CHAT_MESSAGES_FROM',
            worldId: currentWorld.id,
            storyId: currentStory.id,
            messageId,
        });
    }, [currentStory, currentWorld, dispatch]);

    const handleForkChat = useCallback((messageId: string) => {
        if (!currentStory || !currentWorld) return;
        dispatch({
            type: 'DUPLICATE_STORY',
            worldId: currentWorld.id,
            id: currentStory.id,
            upToMessageId: 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]);

    if (!visible) return null;

    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}>
                                            {isEditing ? (
                                                <>
                                                    <button
                                                        class={clsx(styles.iconButton, styles.saveButton)}
                                                        onClick={handleSaveEdit}
                                                        disabled={!editingContent.trim()}
                                                        title="Save (Ctrl+Enter)"
                                                    >
                                                        <Check size={12} />
                                                    </button>
                                                    <button
                                                        class={clsx(styles.iconButton, styles.cancelButton)}
                                                        onClick={handleCancelEdit}
                                                        title="Cancel (Esc)"
                                                    >
                                                        <X size={12} />
                                                    </button>
                                                </>
                                            ) : (
                                                <>
                                                    {currentWorld?.chatOnly && (
                                                        <button
                                                            class={styles.iconButton}
                                                            onClick={() => handleForkChat(message.id)}
                                                            title="Fork chat up to this message"
                                                        >
                                                            <GitFork size={12} />
                                                        </button>
                                                    )}
                                                    <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
                                            highlight
                                            class={styles.editTextarea}
                                            value={editingContent}
                                            onInput={setEditingContent}
                                            onKeyDown={(e) => {
                                                if (e.key === 'Enter' && e.ctrlKey) {
                                                    e.preventDefault();
                                                    handleSaveEdit();
                                                }
                                                if (e.key === 'Escape') {
                                                    e.preventDefault();
                                                    handleCancelEdit();
                                                }
                                            }}
                                        />
                                    </div>
                                ) : (
                                    <>
                                        {message.role === 'assistant' && message.reasoning_content?.trim() && (
                                            <div class={styles.reasoningContent}>
                                                {message.reasoning_content.trim()}
                                            </div>
                                        )}

                                        <div
                                            class={styles.content}
                                            dangerouslySetInnerHTML={{ __html: highlight(Prompt.substituteVars(appState, 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>
                    )}
                </div>
            )}
            {currentStory && (
                <div class={styles.inputContainer}>
                    <div class={styles.inputWrapper}>
                        <ContentEditable
                            autoLines
                            highlight
                            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}
                        />
                        <div class={styles.inputFooter}>
                            <div class={styles.inputFooterLeft}>
                                {tokenCount && (
                                    <span class={styles.tokenCounter}>
                                        {tokenCount.taken} / {tokenCount.total}
                                    </span>
                                )}
                                {!currentWorld?.chatOnly && (
                                    <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>Thinking</span>
                                    </label>
                                )}
                            </div>
                            <div class={styles.inputFooterRight}>
                                {isLoading ? (
                                    <button class={styles.stopButton} onClick={handleStopGeneration} title="Stop">
                                        <Square size={14} />
                                    </button>
                                ) : (
                                    <>
                                        {!currentWorld?.chatOnly && (
                                            <button
                                                class={styles.actionButton}
                                                onClick={summarizeAll}
                                                disabled={isSummarizing || !currentStory || !connection || !model}
                                                title={isSummarizing ? 'Summarizing...' : 'Summarize'}
                                            >
                                                <Sparkles size={14} />
                                            </button>
                                        )}
                                        <button
                                            class={styles.actionButton}
                                            onClick={handleContinue}
                                            disabled={isDisabled}
                                            title="Continue"
                                        >
                                            <ChevronsRight size={14} />
                                        </button>
                                        <button
                                            class={styles.actionButton}
                                            onClick={handleRegenerate}
                                            disabled={isDisabled}
                                            title="Regenerate last response"
                                        >
                                            <RefreshCw size={14} />
                                        </button>
                                        <button
                                            class={styles.sendButton}
                                            onClick={handleSendMessage}
                                            disabled={isDisabled}
                                            title="Send"
                                        >
                                            <SendHorizonal size={14} />
                                        </button>
                                    </>
                                )}
                            </div>
                        </div>
                    </div>
                </div>
            )}
        </div>
    );
};

export const ChatSidebar = () => {
    const { currentWorld, chatOpen, dispatch } = useAppState();
    const isMobile = useMediaQuery(`(max-width: ${MOBILE_BREAKPOINT}px)`);

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

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