import { extension_settings, renderExtensionTemplateAsync } from '../../extensions.js';
import {
    activateSendButtons,
    deactivateSendButtons,
    eventSource,
    event_types,
    extension_prompt_roles,
    extension_prompt_types,
    generateQuietPrompt,
    isGenerating,
    saveSettingsDebounced,
    setExtensionPrompt,
    setSendButtonState,
    streamingProcessor,
} from '../../../script.js';

export { init };

const MODULE = 'keepalive';

// Cap the keepalive completion to a single token. The full prompt prefix is
// still processed by the backend (which is what refreshes the server-side
// cache), but we never pay for a full completion.
const KEEPALIVE_RESPONSE_LENGTH = 1;

// Extension-prompt key used to inject the keepalive message as a USER message at depth 0.
const KEEPALIVE_INJECT_ID = 'keepalive_ping';

// Original default message; used to migrate untouched installs to the current default wording.
const LEGACY_DEFAULT_MESSAGE = 'Keep the cache warm.';

// Idle-tracking key used when no connection profile is active (manual API panel).
const NONE_KEY = '__none__';

// When the active profile is already past its timeout (e.g. you just switched back
// to a profile that went cold), wait this short grace before pinging so an imminent
// user message can pre-empt the keepalive.
const EXPIRED_GRACE_MS = 3000;

// If the idle timer fires more than this many ms later than it was scheduled to, assume the
// machine slept / the tab was throttled — far more wall-clock time passed than the interval,
// so the cache almost certainly lapsed. Skip the ping and go dormant rather than guess.
const MAX_TIMER_DRIFT_MS = 5000;

const defaultSettings = {
    enabled: false,
    timeoutSeconds: 295,
    message: '[Note: just reply with an empty response to keep the session alive]',
    // Per-profile enable map: { [profileId]: boolean }. Unlisted profiles default to enabled.
    profiles: {},
};

let idleTimer = null;
let keepaliveActive = false;
// Whether keepalive is "armed" for the CURRENT chat. It stays dormant until the user sends a new
// message in the open chat — on a freshly loaded or switched-to chat (even with history already
// present) we have no way to know whether its prompt is still cached server-side, so pinging would
// be pointless. Reset on page reload and whenever the chat changes.
let armed = false;
// Wall-clock time (Date.now() + delay) the pending idle timer is expected to fire, or null when
// none is scheduled. Used to detect a timer that fired far too late (sleep/throttle).
let scheduledFireTime = null;
// Per-profile "time since last message", keyed by connection profile id (or NONE_KEY).
// In-memory only: on reload every profile is treated as freshly active.
const lastActivity = {};

/**
 * @returns {string} The active connection profile id, or NONE_KEY for a manual connection.
 */
function getActiveProfileKey() {
    return extension_settings.connectionManager?.selectedProfile || NONE_KEY;
}

/**
 * Whether keepalive should run for the CURRENTLY ACTIVE connection. Requires the global
 * master switch; for a real profile it also requires that profile's per-profile toggle
 * (default on). A manual connection (no profile selected) is governed by the master switch alone.
 * @returns {boolean}
 */
function isKeepaliveEnabledForActive() {
    const s = extension_settings[MODULE];
    if (!s.enabled) {
        return false;
    }
    const profileId = extension_settings.connectionManager?.selectedProfile;
    if (!profileId) {
        return true;
    }
    const perProfile = s.profiles?.[profileId];
    return perProfile === undefined ? true : !!perProfile;
}

/**
 * Clears the pending idle timer, if any.
 */
function clearIdleTimer() {
    if (idleTimer) {
        clearTimeout(idleTimer);
        idleTimer = null;
    }
}

/**
 * (Re)schedules the keepalive for the ACTIVE profile based on how long it has been since
 * that profile's last message. Does nothing (and leaves the timer cleared) until keepalive
 * has been armed for the current chat, or when it is disabled for the active connection.
 * Called on activity AND whenever the active profile changes, so each profile keeps its own
 * independent idle clock.
 */
function scheduleIdleTimer() {
    clearIdleTimer();
    if (!armed || !isKeepaliveEnabledForActive()) {
        scheduledFireTime = null;
        return;
    }
    const timeoutMs = Math.max(60, Number(extension_settings[MODULE].timeoutSeconds) || 295) * 1000;
    const last = lastActivity[getActiveProfileKey()] ?? Date.now();
    const remaining = timeoutMs - (Date.now() - last);
    const delay = remaining > 0 ? remaining : EXPIRED_GRACE_MS;
    scheduledFireTime = Date.now() + delay;
    idleTimer = setTimeout(fireKeepalive, delay);
}

/**
 * Marks the active profile as just-used (resets its idle clock) and reschedules its keepalive.
 */
function markActivity() {
    lastActivity[getActiveProfileKey()] = Date.now();
    scheduleIdleTimer();
}

/**
 * The user sent a new message in the current chat — the ONLY thing that arms keepalive.
 * Background/quiet generations, swipes, regenerations, and loading chat history deliberately do
 * NOT arm it: none of those prove the user is actively continuing THIS chat from a known-cached
 * prompt (GENERATION_STARTED in particular fires for every quiet/background generation too).
 */
function onMessageSent() {
    armed = true;
    markActivity();
}

/**
 * Opening or switching chats disarms keepalive: a freshly loaded chat is not proof its prompt is
 * still cached server-side, so we wait for a new sent message in that chat before running again.
 */
function onChatChanged() {
    armed = false;
    scheduleIdleTimer();
}

/**
 * Fires a background ("quiet") generation against the active model to keep the
 * server-side prompt cache warm. The full prompt prefix (current chat context +
 * the keepalive message) is sent — which is what refreshes the cache — while the
 * completion is capped to a single token (KEEPALIVE_RESPONSE_LENGTH) so we never
 * pay for a real response.
 *
 * Only the CURRENTLY ACTIVE profile is ever pinged — keepalive never switches the
 * connection in the background. Each profile's idle clock is tracked separately, so
 * switching profiles re-targets the timer at whichever profile is now active.
 *
 * Quiet generations are forced non-streaming by the core (Generate excludes
 * type === 'quiet' from the StreamingProcessor) and are never saved to chat, so the
 * output is simply discarded. generateQuietPrompt's responseLength applies a *global*
 * temporary response-length override for the duration of the request, so we hold the
 * same generation lock the blocking summary path uses — setSendButtonState(true) flips
 * `is_send_press` (which every user send entry point checks and bails on) and
 * deactivateSendButtons() reflects the busy state — and release both in `finally`.
 */
async function fireKeepalive() {
    if (!armed || !isKeepaliveEnabledForActive()) {
        return;
    }

    // If the timer fired much later than scheduled (machine sleep, background-tab throttling,
    // etc.), far more wall-clock time has elapsed than the idle interval, so the cache has
    // almost certainly expired. Don't ping on a stale assumption — go dormant until the next
    // sent message re-arms keepalive, exactly like a freshly opened chat.
    if (scheduledFireTime !== null && (Date.now() - scheduledFireTime) > MAX_TIMER_DRIFT_MS) {
        console.debug('[Keepalive] Idle timer fired', Date.now() - scheduledFireTime, 'ms late; skipping and disarming until next activity.');
        armed = false;
        clearIdleTimer();
        scheduledFireTime = null;
        return;
    }

    // Never start while any generation is in progress, or while a previous
    // keepalive is still running. Reschedule and bail.
    if (keepaliveActive || isGenerating() || streamingProcessor) {
        scheduleIdleTimer();
        return;
    }

    keepaliveActive = true;
    setSendButtonState(true);
    deactivateSendButtons();

    try {
        // Deliver the keepalive as a USER message at depth 0 (the quiet-prompt path forces a
        // system role — openai.js hardcodes role:'system' for it), using the same depth+role
        // injection Author's Note uses. The quiet prompt itself stays empty so the full chat
        // context is still assembled (which is what keeps the cached prefix warm).
        setExtensionPrompt(KEEPALIVE_INJECT_ID, extension_settings[MODULE].message, extension_prompt_types.IN_CHAT, 0, false, extension_prompt_roles.USER);
        await generateQuietPrompt({ quietPrompt: '', responseLength: KEEPALIVE_RESPONSE_LENGTH });
        // The ping kept this profile warm; treat it as fresh activity so the next
        // keepalive for it is a full timeout away.
        lastActivity[getActiveProfileKey()] = Date.now();
    } catch (error) {
        // Best-effort: any failure here is harmless and intentionally swallowed.
        console.debug('[Keepalive] Background generation ended:', error);
    } finally {
        setExtensionPrompt(KEEPALIVE_INJECT_ID, '', extension_prompt_types.IN_CHAT, 0, false, extension_prompt_roles.USER);
        setSendButtonState(false);
        activateSendButtons();
        keepaliveActive = false;
        scheduleIdleTimer();
    }
}

/**
 * Renders the per-profile enable checkboxes from the current Connection Manager profiles.
 */
function renderProfileToggles() {
    const container = $('#keepalive_profiles_list');
    if (!container.length) {
        return;
    }
    container.empty();

    const profiles = extension_settings.connectionManager?.profiles;
    if (!Array.isArray(profiles) || profiles.length === 0) {
        container.append(
            $('<small>')
                .attr('data-i18n', 'ext_keepalive_no_profiles')
                .text('No connection profiles found. Keepalive uses the active model, governed by the master toggle above.'),
        );
        return;
    }

    const perProfile = extension_settings[MODULE].profiles || {};
    const sorted = profiles.slice().sort((a, b) => String(a.name).localeCompare(String(b.name)));
    for (const profile of sorted) {
        const enabled = perProfile[profile.id] === undefined ? true : !!perProfile[profile.id];
        const label = $('<label>').addClass('checkbox_label');
        const input = $('<input>')
            .attr('type', 'checkbox')
            .addClass('keepalive_profile_toggle')
            .attr('data-profile-id', profile.id)
            .prop('checked', enabled);
        const span = $('<span>').text(profile.name);
        label.append(input, span);
        container.append(label);
    }
}

/**
 * Loads settings into the global store and hydrates the UI controls.
 */
function loadSettings() {
    if (extension_settings[MODULE] === undefined) {
        extension_settings[MODULE] = {};
    }

    for (const key of Object.keys(defaultSettings)) {
        if (extension_settings[MODULE][key] === undefined) {
            extension_settings[MODULE][key] = structuredClone(defaultSettings[key]);
        }
    }

    // Guard against a corrupted/legacy value for the per-profile map.
    if (typeof extension_settings[MODULE].profiles !== 'object' || extension_settings[MODULE].profiles === null) {
        extension_settings[MODULE].profiles = {};
    }

    // Migrate the original default message to the current default so installs that never
    // customized it pick up the improved wording (custom messages are left untouched).
    if (extension_settings[MODULE].message === LEGACY_DEFAULT_MESSAGE) {
        extension_settings[MODULE].message = defaultSettings.message;
        saveSettingsDebounced();
    }

    $('#keepalive_enabled').prop('checked', extension_settings[MODULE].enabled);
    $('#keepalive_timeout').val(extension_settings[MODULE].timeoutSeconds);
    $('#keepalive_message').val(extension_settings[MODULE].message);
    renderProfileToggles();
}

/**
 * Wires up the settings UI control handlers.
 */
function setupListeners() {
    $('#keepalive_enabled').on('change', function () {
        extension_settings[MODULE].enabled = !!$(this).prop('checked');
        saveSettingsDebounced();
        scheduleIdleTimer();
    });

    $('#keepalive_timeout').on('input', function () {
        extension_settings[MODULE].timeoutSeconds = Number($(this).val()) || 295;
        saveSettingsDebounced();
        scheduleIdleTimer();
    });

    $('#keepalive_message').on('input', function () {
        extension_settings[MODULE].message = String($(this).val());
        saveSettingsDebounced();
    });

    // Per-profile toggles are rendered dynamically; use a delegated handler.
    $('#keepalive_profiles_list').on('change', '.keepalive_profile_toggle', function () {
        const profileId = $(this).attr('data-profile-id');
        if (!profileId) {
            return;
        }
        extension_settings[MODULE].profiles[profileId] = !!$(this).prop('checked');
        saveSettingsDebounced();
        // If the toggled profile is the one currently active, reschedule immediately.
        if (profileId === extension_settings.connectionManager?.selectedProfile) {
            scheduleIdleTimer();
        }
    });
}

async function init() {
    const settingsHtml = await renderExtensionTemplateAsync(MODULE, 'settings');
    $('#extensions_settings2').append(settingsHtml);

    loadSettings();
    setupListeners();

    // Only the user sending a new message arms keepalive (for the current chat).
    eventSource.on(event_types.MESSAGE_SENT, onMessageSent);

    // These reset the active profile's idle clock (so generation time isn't counted as idle) but
    // never arm keepalive on their own — they also fire for background/quiet generations and while
    // a chat is loading. They only re-arm the timer if a sent message already armed it.
    const activityEvents = [
        event_types.MESSAGE_RECEIVED,
        event_types.GENERATION_STARTED,
        event_types.GENERATION_ENDED,
    ];
    for (const event of activityEvents) {
        eventSource.on(event, markActivity);
    }

    // Opening/switching a chat disarms keepalive — it must be re-armed by a new sent message.
    eventSource.on(event_types.CHAT_CHANGED, onChatChanged);

    // When the active profile changes, re-target the idle timer at the new profile using
    // ITS own last-activity time (do NOT mark activity — switching is not using it).
    eventSource.on(event_types.CONNECTION_PROFILE_LOADED, scheduleIdleTimer);

    // Keep the per-profile toggle list in sync with Connection Manager.
    eventSource.on(event_types.CONNECTION_PROFILE_CREATED, renderProfileToggles);
    eventSource.on(event_types.CONNECTION_PROFILE_UPDATED, renderProfileToggles);
    eventSource.on(event_types.CONNECTION_PROFILE_DELETED, renderProfileToggles);

    scheduleIdleTimer();
}