ai-translate/src/services/dbService.js

import initSqlJs from 'sql.js';
import localforage from 'localforage';
import { v4 as uuidv4 } from 'uuid';

const DB_NAME = 'ai_translate_db_v1';
const STORE_KEY = 'sqlite_binary';

let db = null;
let SQL = null;

export const dbService = {
    init: async () => {
        if (db) return db;

        try {
            // Load WASM
            SQL = await initSqlJs({
                // Locate the WASM file in public folder
                locateFile: file => `/${file}`
            });

            // Try to load saved binary from IndexedDB
            const savedData = await localforage.getItem(STORE_KEY);

            if (savedData) {
                // Load existing DB
                db = new SQL.Database(new Uint8Array(savedData));
                console.log("[DB] Loaded existing database.");
            } else {
                // Create new DB
                db = new SQL.Database();
                console.log("[DB] Created new database.");
                dbService._createTables();
                await dbService.save(); // Initial save
            }

            return db;
        } catch (e) {
            console.error("[DB] Initialization Failed:", e);
            throw e;
        }
    },

    _createTables: () => {
        // Stories Table
        db.run(`
            CREATE TABLE IF NOT EXISTS stories (
                id TEXT PRIMARY KEY,
                title TEXT,
                author TEXT,
                summary TEXT,
                cover_image TEXT, -- Base64 or URL
                created_at INTEGER,
                last_accessed INTEGER,
                metadata TEXT -- JSON: { genres, languages, personality, world... }
            );
        `);

        // Files Table
        db.run(`
            CREATE TABLE IF NOT EXISTS files (
                id TEXT PRIMARY KEY, -- Use existing file IDs if possible, or new UUIDs
                story_id TEXT,
                name TEXT,
                content TEXT,
                translated_content TEXT,
                status TEXT,
                size INTEGER,
                last_modified INTEGER,
                FOREIGN KEY(story_id) REFERENCES stories(id) ON DELETE CASCADE
            );
        `);

        // Dictionary Table (Global or per story? Let's make it simple first: Per Story stored in metadata or separate?)
        // For now, Dictionary is just text in App.jsx. We can store it in metadata if needed.
    },

    save: async () => {
        if (!db) return;
        try {
            const data = db.export();
            await localforage.setItem(STORE_KEY, data);
            // console.log("[DB] Saved to IndexedDB");
        } catch (e) {
            console.error("[DB] Save Failed:", e);
        }
    },

    // --- Story Operations ---

    createStory: async (title = "Truyện Mới") => {
        if (!db) {
            console.error("[DB] Database is null!");
            throw new Error("Database not initialized");
        }
        console.log("[DB] Executing createStory for:", title);
        const id = uuidv4();
        const now = Date.now();
        const metadata = JSON.stringify({
            genres: [], languages: ['Tiếng Trung'], mcPersonality: [], worldSetting: [], contextNotes: '', additionalDictionary: '', customPrompt: ''
        });

        try {
            db.run(`INSERT INTO stories (id, title, created_at, last_accessed, metadata) VALUES (?, ?, ?, ?, ?)`,
                [id, title, now, now, metadata]);
            await dbService.save();
            console.log("[DB] Insert Successful, ID:", id);
            return id;
        } catch (e) {
            console.error("[DB] Insert Failed:", e);
            throw e;
        }
    },

    getStories: () => {
        if (!db) return [];
        const result = db.exec("SELECT id, title, author, cover_image, last_accessed FROM stories ORDER BY created_at DESC");
        if (result.length === 0) return [];

        // Transform result [columns, values] to objects
        const columns = result[0].columns;
        return result[0].values.map(row => {
            const obj = {};
            columns.forEach((col, i) => obj[col] = row[i]);
            return obj;
        });
    },

    getStory: (id) => {
        if (!db) return null;
        const result = db.exec("SELECT * FROM stories WHERE id = ?", [id]);
        if (result.length === 0) return null;

        const columns = result[0].columns;
        const row = result[0].values[0];
        const story = {};
        columns.forEach((col, i) => story[col] = row[i]);

        // Parse metadata
        try {
            if (story.metadata) Object.assign(story, JSON.parse(story.metadata));
        } catch (e) {
            console.error("Failed to parse metadata", e);
        }
        return story;
    },

    updateStory: async (id, data) => {
        if (!db) return;
        const now = Date.now();

        // Separate columns vs metadata
        const directCols = ['title', 'author', 'summary', 'cover_image'];
        const metadataKeys = ['genres', 'languages', 'mcPersonality', 'worldSetting', 'contextNotes', 'image_prompt', 'additionalDictionary', 'customPrompt'];

        // 1. Get current metadata
        const current = dbService.getStory(id);
        if (!current) return;

        // 2. Prepare Updates
        const updates = [];
        const values = [];

        directCols.forEach(col => {
            if (data[col] !== undefined) {
                updates.push(`${col} = ?`);
                values.push(data[col]);
            }
        });

        // 3. Update Metadata
        const newMetadata = {};
        metadataKeys.forEach(k => {
            if (data[k] !== undefined) newMetadata[k] = data[k];
            else if (current[k] !== undefined) newMetadata[k] = current[k];
        });

        updates.push(`metadata = ?`);
        values.push(JSON.stringify(newMetadata));

        updates.push(`last_accessed = ?`);
        values.push(now);

        // Where ID
        values.push(id);

        const query = `UPDATE stories SET ${updates.join(', ')} WHERE id = ?`;
        db.run(query, values);
        await dbService.save();
    },

    deleteStory: async (id) => {
        db.run("DELETE FROM stories WHERE id = ?", [id]);
        // Files cascade delete handled by schema? sql.js usually doesn't enforce FK by default unless enabled
        // Manually delete files to be safe
        db.run("DELETE FROM files WHERE story_id = ?", [id]);
        await dbService.save();
    },

    // --- File Operations ---

    getFiles: (storyId) => {
        const res = db.exec("SELECT * FROM files WHERE story_id = ?", [storyId]);
        if (res.length === 0) return [];

        const columns = res[0].columns;
        return res[0].values.map(row => {
            const f = {};
            columns.forEach((col, i) => {
                // Map snake_case to camelCase
                if (col === 'translated_content') f['translatedContent'] = row[i];
                else if (col === 'last_modified') f['lastModified'] = row[i];
                else f[col] = row[i];
            });
            return f;
        });
    },

    saveFilesBatch: async (storyId, files) => {
        // Use transaction
        db.run("BEGIN TRANSACTION");
        const stmt = db.prepare("INSERT OR REPLACE INTO files (id, story_id, name, content, translated_content, status, size, last_modified) VALUES (?, ?, ?, ?, ?, ?, ?, ?)");

        files.forEach(f => {
            stmt.run([f.id, storyId, f.name, f.content, f.translatedContent || null, f.status || 'IDLE', f.content.length, Date.now()]);
        });

        stmt.free();
        db.run("COMMIT");
        await dbService.save();
    },

    deleteFile: async (id) => {
        db.run("DELETE FROM files WHERE id = ?", [id]);
        await dbService.save();
    },

    deleteAllFiles: async (storyId) => {
        db.run("DELETE FROM files WHERE story_id = ?", [storyId]);
        await dbService.save();
    }
};