src/endpoints/secure-generate.js

import express from 'express';
import fetch from 'node-fetch';
import fs from 'node:fs';
import path from 'node:path';
import multer from 'multer';
import { forwardFetchResponse } from '../util.js';

export const router = express.Router();

const HIDDEN_PROMPT_PLACEHOLDER = /__HIDDEN_PROMPT:([a-zA-Z0-9_.-]+)__/g;

const SCOPES = ['context', 'instruct', 'sysprompt', 'main', 'jailbreak', 'quickreply'];

// ---------- Admin key ----------
const ADMIN_KEY = process.env.ADMIN_KEY || 'admin123';

function isAdmin(req) {
    const key = req.query.admin_key || req.headers['x-admin-key'] || '';
    return key === ADMIN_KEY;
}

// ---------- Data helpers ----------

function getDataRoot() {
    return globalThis.DATA_ROOT || '';
}

function readJsonFile(filePath) {
    try {
        if (fs.existsSync(filePath)) {
            return JSON.parse(fs.readFileSync(filePath, 'utf8'));
        }
    } catch (err) {
        console.error('Error reading hidden prompts:', err);
    }
    return {};
}

/**
 * Directory for uploaded presets (full SillyTavern preset JSONs).
 */
function getPresetsDir() {
    return path.join(getDataRoot(), 'hidden_prompt_presets');
}

/**
 * Read a preset file by name. Returns the parsed JSON or null.
 * @param {string} name - Preset name (without .json extension)
 */
function readPreset(name) {
    const safeName = name.replace(/[^a-zA-Z0-9\u4e00-\u9fff_.-]/g, '_');
    const filePath = path.join(getPresetsDir(), `${safeName}.json`);
    try {
        if (fs.existsSync(filePath)) {
            return JSON.parse(fs.readFileSync(filePath, 'utf8'));
        }
    } catch (err) {
        console.error(`Error reading preset ${name}:`, err);
    }
    return null;
}

/**
 * List all available preset names.
 * @returns {string[]}
 */
function listPresets() {
    const dir = getPresetsDir();
    try {
        if (!fs.existsSync(dir)) return [];
        return fs.readdirSync(dir)
            .filter(f => f.endsWith('.json'))
            .map(f => f.replace(/\.json$/, ''));
    } catch (err) {
        console.error('Error listing presets:', err);
        return [];
    }
}

/**
 * Get all preset prompts as a flat map: "preset__identifier" -> { name, prompt, role }.
 * Used for injection during generation.
 */
function getAllPresetPromptsFlat() {
    const flat = {};
    for (const presetName of listPresets()) {
        const preset = readPreset(presetName);
        if (!preset || !Array.isArray(preset.prompts)) continue;
        const slug = presetName.replace(/\s+/g, '_');
        for (const p of preset.prompts) {
            if (!p.content || p.marker) continue;
            const qualifiedId = `preset__${slug}__${p.identifier}`;
            flat[qualifiedId] = {
                name: p.name || p.identifier,
                prompt: p.content,
                role: p.role || 'system',
            };
        }
    }
    return flat;
}

// ---------- Legacy hidden prompt helpers (kept for backward compat) ----------

function getHiddenPromptsForScope(scope) {
    const dataRoot = getDataRoot();
    const scopeFile = path.join(dataRoot, `hidden_prompts_${scope}.json`);
    const raw = readJsonFile(scopeFile);
    const result = {};
    for (const [id, data] of Object.entries(raw)) {
        result[`${scope}__${id}`] = data;
    }
    return result;
}

function getHiddenPrompts() {
    const dataRoot = getDataRoot();
    const merged = {};

    // Legacy: hidden_prompts.json (no prefix)
    const legacyFile = path.join(dataRoot, 'hidden_prompts.json');
    Object.assign(merged, readJsonFile(legacyFile));

    // Scope-specific files: scope__id
    for (const scope of SCOPES) {
        const scopeFile = path.join(dataRoot, `hidden_prompts_${scope}.json`);
        const raw = readJsonFile(scopeFile);
        for (const [id, data] of Object.entries(raw)) {
            merged[`${scope}__${id}`] = data;
        }
    }

    // Preset-based prompts: preset__presetSlug__identifier
    Object.assign(merged, getAllPresetPromptsFlat());

    return merged;
}

function replaceHiddenPromptPlaceholders(obj, hiddenPrompts) {
    if (typeof obj === 'string') {
        return obj.replace(HIDDEN_PROMPT_PLACEHOLDER, (_, id) => {
            const hp = hiddenPrompts[id];
            return hp ? hp.prompt : `[Hidden prompt ${id} not found]`;
        });
    }
    if (Array.isArray(obj)) {
        for (let i = 0; i < obj.length; i++) {
            obj[i] = replaceHiddenPromptPlaceholders(obj[i], hiddenPrompts);
        }
        return obj;
    }
    if (obj && typeof obj === 'object') {
        for (const key of Object.keys(obj)) {
            obj[key] = replaceHiddenPromptPlaceholders(obj[key], hiddenPrompts);
        }
        return obj;
    }
    return obj;
}

// ===================== ROUTES =====================

// ---------- Main proxy endpoint ----------
router.post('/', async (req, res) => {
    const { target_url, hidden_prompt_id, hidden_prompt_ids, has_hidden_placeholders, ...llm_params } = req.body;

    if (!target_url) {
        return res.status(400).send('Missing target_url');
    }

    const hiddenPrompts = getHiddenPrompts();

    // Replace __HIDDEN_PROMPT:id__ placeholders
    replaceHiddenPromptPlaceholders(llm_params, hiddenPrompts);

    // Inject hidden prompts by ID
    const idsToInject = Array.isArray(hidden_prompt_ids) && hidden_prompt_ids.length
        ? hidden_prompt_ids.filter(id => id)
        : (hidden_prompt_id ? [hidden_prompt_id] : []);

    if (idsToInject.length && llm_params.messages) {
        for (let i = idsToInject.length - 1; i >= 0; i--) {
            const id = idsToInject[i];
            const hiddenPrompt = hiddenPrompts[id];
            if (hiddenPrompt) {
                console.log(`Injecting hidden prompt: ${id}`);
                llm_params.messages.unshift({
                    role: hiddenPrompt.role || 'system',
                    content: hiddenPrompt.prompt,
                });
            }
        }
        console.log('>>> Final messages array sent to model:', JSON.stringify(llm_params.messages, null, 2));
    }

    try {
        const urlToFetch = target_url.startsWith('http')
            ? target_url
            : `${req.protocol}://${req.get('host') || '127.0.0.1'}${target_url.startsWith('/') ? target_url : '/' + target_url}`;

        const headers = { ...req.headers };
        delete headers.host;
        delete headers['content-length'];
        const isSameOrigin = urlToFetch.startsWith(`${req.protocol}://${req.get('host') || '127.0.0.1'}`);
        if (!isSameOrigin) {
            delete headers['x-csrf-token'];
            delete headers.cookie;
        }

        const response = await fetch(urlToFetch, {
            method: 'POST',
            headers: /** @type {any} */(headers),
            body: JSON.stringify(llm_params),
        });

        forwardFetchResponse(response, res);
    } catch (error) {
        console.error('Error in secure-generate proxy:', error);
        res.status(500).send('Error in secure-generate proxy: ' + error.message);
    }
});

// ---------- Legacy list endpoint ----------
router.get('/list', (req, res) => {
    const scope = /** @type {string|undefined} */(req.query.scope);
    let hiddenPrompts;
    if (scope && SCOPES.includes(scope)) {
        hiddenPrompts = getHiddenPromptsForScope(scope);
    } else {
        const dataRoot = getDataRoot();
        hiddenPrompts = readJsonFile(path.join(dataRoot, 'hidden_prompts.json'));
    }
    const list = Object.entries(hiddenPrompts).map(([id, data]) => ({
        id,
        label: /** @type {any} */(data).name || id,
    }));
    res.json(list);
});

// ===================== PRESET ENDPOINTS =====================

/**
 * GET /presets — list available preset names.
 */
router.get('/presets', (_req, res) => {
    res.json(listPresets());
});

/**
 * GET /presets/:name/prompts — flat list of prompts for a preset.
 * Regular users: name, identifier, role, description only (no content).
 * Admin users (admin_key): includes full content.
 */
router.get('/presets/:name/prompts', (req, res) => {
    const name = req.params.name;
    const preset = readPreset(name);
    if (!preset) {
        return res.status(404).json({ error: `Preset "${name}" not found` });
    }

    const prompts = preset.prompts || [];
    const admin = isAdmin(req);
    const presetSlug = name.replace(/\s+/g, '_');

    const result = prompts.map(p => {
        const base = {
            identifier: p.identifier,
            qualifiedId: `preset__${presetSlug}__${p.identifier}`,
            name: p.name || p.identifier,
            role: p.role || 'system',
            description: p.description || '',
            system_prompt: !!p.system_prompt,
            enabled: p.enabled !== false,
            marker: !!p.marker,
            injection_position: p.injection_position,
            injection_depth: p.injection_depth,
        };
        if (admin) {
            base.content = p.content || '';
        }
        return base;
    });

    res.json({ presetName: name, admin, prompts: result });
});

/**
 * GET /presets/:name/import-data — (legacy, kept for compatibility)
 */
router.get('/presets/:name/import-data', (req, res) => {
    const name = req.params.name;
    const preset = readPreset(name);
    if (!preset) {
        return res.status(404).json({ error: `Preset "${name}" not found` });
    }
    const admin = isAdmin(req);
    const result = JSON.parse(JSON.stringify(preset));
    if (!admin) {
        (result.prompts || []).forEach(p => { p.content = '[Hidden Content]'; });
    }
    res.json(result);
});

/**
 * GET /presets/:name/full — returns the FULL SillyTavern preset JSON.
 * This is the exact format the native preset loader expects (temperature, prompts, prompt_order, etc.).
 * Non-admin: content is replaced with '[Hidden Content]'.
 * Admin: full content.
 */
router.get('/presets/:name/full', (req, res) => {
    const name = req.params.name;
    const preset = readPreset(name);
    if (!preset) {
        return res.status(404).json({ error: `Preset "${name}" not found` });
    }

    const admin = isAdmin(req);
    // Deep copy to avoid mutating the cached preset
    const result = JSON.parse(JSON.stringify(preset));

    if (!admin) {
        (result.prompts || []).forEach(p => {
            p.content = '[Hidden Content]';
        });
    }

    res.json(result);
});

// ---------- Multer for preset uploads ----------
const presetUpload = multer({
    dest: path.join(getPresetsDir(), '_uploads'),
    fileFilter: (_req, file, cb) => {
        if (file.mimetype !== 'application/json' && !file.originalname.endsWith('.json')) {
            return cb(new Error('Only JSON files are allowed'));
        }
        cb(null, true);
    },
    limits: { fileSize: 10 * 1024 * 1024 },
});

/**
 * POST /presets/upload — admin-only: upload a SillyTavern preset file.
 * Form fields: file (the JSON), preset_name (display name).
 */
router.post('/presets/upload', presetUpload.single('file'), (req, res) => {
    if (!isAdmin(req)) {
        if (req.file && fs.existsSync(req.file.path)) fs.unlinkSync(req.file.path);
        return res.status(403).json({ error: 'Admin access required' });
    }
    if (!req.file) {
        return res.status(400).json({ error: 'No file uploaded' });
    }

    const presetName = (req.body?.preset_name || '').trim();
    if (!presetName) {
        fs.unlinkSync(req.file.path);
        return res.status(400).json({ error: 'Missing preset_name field' });
    }

    try {
        const data = JSON.parse(fs.readFileSync(req.file.path, 'utf8'));
        if (!Array.isArray(data.prompts) || data.prompts.length === 0) {
            fs.unlinkSync(req.file.path);
            return res.status(400).json({ error: 'Invalid preset: no prompts array found' });
        }

        // Save with the given name
        const safeName = presetName.replace(/[^a-zA-Z0-9\u4e00-\u9fff_.-]/g, '_');
        const dir = getPresetsDir();
        if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });
        const outputPath = path.join(dir, `${safeName}.json`);
        fs.renameSync(req.file.path, outputPath);

        console.log(`Preset uploaded: ${presetName} (${data.prompts.length} prompts)`);
        res.json({ success: true, name: presetName, promptCount: data.prompts.length });
    } catch (err) {
        if (req.file && fs.existsSync(req.file.path)) fs.unlinkSync(req.file.path);
        res.status(400).json({ error: 'Invalid JSON: ' + err.message });
    }
});