/*
    vibecoded :sob:
*/

import { Pool } from "pg";

export class KVDB {
  /**
   * @param {Pool | object} poolOrConfig - A pg.Pool instance or pg connection config object
   * @param {string} [table="main_kv"] - The table/context name (defaults to "main_kv")
   */
  constructor(poolOrConfig, table = "main_kv") {
    this.pool = poolOrConfig instanceof Pool ? poolOrConfig : new Pool(poolOrConfig);

    // Sanitize table name to prevent SQL injection in DDL
    this.tableName = table.replace(/[^a-zA-Z0-9_]/g, "");
  }

  /**
   * Initialize the database table. Call this once after instantiating the class.
   * @returns {Promise<void>}
   */
  async init() {
    await this.pool.query(`
      CREATE TABLE IF NOT EXISTS "${this.tableName}" (
        key TEXT PRIMARY KEY,
        value JSONB
      )
    `);
  }

  /**
   * Switch database table/context sharing the same connection pool
   * @param {string} newTableName 
   * @returns {Promise<KVDB>}
   */
  async table(newTableName) {
    const newDb = new KVDB(this.pool, newTableName);
    await newDb.init();
    return newDb;
  }

  /**
   * Set a key or nested dot-notation property.
   * Examples:
   *   await db.set("user", { name: "Alex" })
   *   await db.set("user.age", 30)
   * @param {string} keyPath 
   * @param {any} value 
   * @returns {Promise<void>}
   */
  async set(keyPath, value) {
    const parts = keyPath.split(".");
    const rootKey = parts[0];
    const jsonValue = JSON.stringify(value);

    if (parts.length === 1) {
      await this.pool.query(
        `INSERT INTO "${this.tableName}" (key, value) VALUES ($1, $2::jsonb)
         ON CONFLICT(key) DO UPDATE SET value = EXCLUDED.value`,
        [rootKey, jsonValue]
      );
      return;
    }

    const pathArr = parts.slice(1);

    await this.pool.query(
      `INSERT INTO "${this.tableName}" (key, value) VALUES ($1, '{}'::jsonb)
       ON CONFLICT(key) DO NOTHING`,
      [rootKey]
    );

    await this.pool.query(
      `UPDATE "${this.tableName}" 
       SET value = jsonb_set(value, $1::text[], $2::jsonb, true) 
       WHERE key = $3`,
      [pathArr, jsonValue, rootKey]
    );
  }

  /**
   * Get a key/nested property, OR find an element inside an array using a predicate function.
   * @param {string} keyPath 
   * @param {Function} [predicate] - Optional matching function for array items
   * @returns {Promise<any>}
   */
  async get(keyPath, predicate) {
    const parts = keyPath.split(".");
    const rootKey = parts[0];
    let data = null;

    if (parts.length === 1) {
      const { rows } = await this.pool.query(
        `SELECT value FROM "${this.tableName}" WHERE key = $1`,
        [rootKey]
      );
      // pg parses JSONB automatically
      data = rows.length > 0 ? rows[0].value : null; 
    } else {
      const pathArr = parts.slice(1);
      const { rows } = await this.pool.query(
        `SELECT value #> $1::text[] AS result FROM "${this.tableName}" WHERE key = $2`,
        [pathArr, rootKey]
      );
      data = (rows.length > 0 && rows[0].result !== undefined) ? rows[0].result : null;
    }

    if (typeof predicate === "function") {
      if (!Array.isArray(data)) return null;
      return data.find(predicate) ?? null;
    }

    return data;
  }

  /**
   * Check if a key/property exists, OR check if an array contains a specific item / predicate match.
   * @param {string} keyPath 
   * @param {any | Function} [itemOrPredicate] 
   * @returns {Promise<boolean>}
   */
  async has(keyPath, itemOrPredicate) {
    if (itemOrPredicate !== undefined) {
      const arr = await this.get(keyPath);
      if (!Array.isArray(arr)) return false;

      if (typeof itemOrPredicate === "function") {
        return arr.some(itemOrPredicate);
      }

      if (typeof itemOrPredicate === "object" && itemOrPredicate !== null) {
        const targetString = JSON.stringify(itemOrPredicate);
        return arr.some((el) => JSON.stringify(el) === targetString);
      }

      return arr.includes(itemOrPredicate);
    }

    const parts = keyPath.split(".");
    const rootKey = parts[0];

    if (parts.length === 1) {
      const { rows } = await this.pool.query(
        `SELECT 1 FROM "${this.tableName}" WHERE key = $1`,
        [rootKey]
      );
      return rows.length > 0;
    }

    const pathArr = parts.slice(1);
    const { rows } = await this.pool.query(
      `SELECT 1 FROM "${this.tableName}" 
       WHERE key = $1 AND value #> $2::text[] IS NOT NULL`,
      [rootKey, pathArr]
    );

    return rows.length > 0;
  }

  /**
   * Delete a key, property, OR remove item(s) from an array using a predicate or value.
   * @param {string} keyPath 
   * @param {any | Function} [itemOrPredicate] 
   * @returns {Promise<void>}
   */
  async delete(keyPath, itemOrPredicate) {
    if (itemOrPredicate !== undefined) {
      const arr = await this.get(keyPath);
      if (!Array.isArray(arr)) return;

      let filtered;
      if (typeof itemOrPredicate === "function") {
        filtered = arr.filter((item) => !itemOrPredicate(item));
      } else if (typeof itemOrPredicate === "object" && itemOrPredicate !== null) {
        const targetString = JSON.stringify(itemOrPredicate);
        filtered = arr.filter((item) => JSON.stringify(item) !== targetString);
      } else {
        filtered = arr.filter((item) => item !== itemOrPredicate);
      }

      await this.set(keyPath, filtered);
      return;
    }

    const parts = keyPath.split(".");
    const rootKey = parts[0];

    if (parts.length === 1) {
      await this.pool.query(`DELETE FROM "${this.tableName}" WHERE key = $1`, [rootKey]);
      return;
    }

    const pathArr = parts.slice(1);
    await this.pool.query(
      `UPDATE "${this.tableName}" 
       SET value = value #- $1::text[] 
       WHERE key = $2`,
      [pathArr, rootKey]
    );
  }

  /**
   * Appends an element to an array at the target key or nested path.
   * @param {string} keyPath 
   * @param {any} value 
   * @returns {Promise<void>}
   */
  async push(keyPath, value) {
    const parts = keyPath.split(".");
    const rootKey = parts[0];
    const wrappedValue = JSON.stringify([value]);

    if (parts.length === 1) {
      await this.pool.query(
        `INSERT INTO "${this.tableName}" (key, value) VALUES ($1, $2::jsonb)
         ON CONFLICT(key) DO UPDATE 
         SET value = CASE 
           WHEN jsonb_typeof("${this.tableName}".value) = 'array' 
           THEN "${this.tableName}".value || $2::jsonb
           ELSE $2::jsonb
         END`,
        [rootKey, wrappedValue]
      );
    } else {
      const pathArr = parts.slice(1);

      await this.pool.query(
        `INSERT INTO "${this.tableName}" (key, value) VALUES ($1, '{}'::jsonb)
         ON CONFLICT(key) DO NOTHING`,
        [rootKey]
      );

      await this.pool.query(
        `UPDATE "${this.tableName}" 
         SET value = jsonb_set(
           value, 
           $1::text[], 
           COALESCE(value #> $1::text[], '[]'::jsonb) || $2::jsonb, 
           true
         ) 
         WHERE key = $3`,
        [pathArr, wrappedValue, rootKey]
      );
    }
  }

  /**
   * Updates matching item(s) in an array using a predicate and an update callback/object.
   * @param {string} keyPath 
   * @param {Function} predicate 
   * @param {Object | Function} updateFnOrObject 
   * @returns {Promise<boolean>}
   */
  async update(keyPath, predicate, updateFnOrObject) {
    const arr = await this.get(keyPath);

    if (!Array.isArray(arr)) return false;

    let updatedAny = false;

    const newArr = arr.map((item) => {
      if (!predicate(item)) return item;

      updatedAny = true;

      if (typeof updateFnOrObject === "function") {
        return updateFnOrObject(item);
      }

      if (typeof item === "object" && item !== null && typeof updateFnOrObject === "object") {
        return { ...item, ...updateFnOrObject };
      }

      return updateFnOrObject;
    });

    if (updatedAny) {
      await this.set(keyPath, newArr);
    }

    return updatedAny;
  }

  /**
   * Renames a key or nested property path without changing its value.
   * @param {string} oldKeyPath 
   * @param {string} newKeyPath 
   * @returns {Promise<boolean>} True if target key existed and was renamed
   */
  async rename(oldKeyPath, newKeyPath) {
    if (!(await this.has(oldKeyPath))) return false;

    const oldParts = oldKeyPath.split(".");
    let targetNewPath = newKeyPath;

    if (oldParts.length > 1 && !newKeyPath.startsWith(oldParts.slice(0, -1).join(".") + ".")) {
      const parentPath = oldParts.slice(0, -1).join(".");
      targetNewPath = `${parentPath}.${newKeyPath}`;
    }

    const value = await this.get(oldKeyPath);

    const oldRoot = oldParts[0];
    const newParts = targetNewPath.split(".");
    const newRoot = newParts[0];

    if (oldParts.length === 1 && newParts.length === 1) {
      await this.pool.query(
        `UPDATE "${this.tableName}" SET key = $1 WHERE key = $2`,
        [newRoot, oldRoot]
      );
    } else {
      await this.set(targetNewPath, value);
      await this.delete(oldKeyPath);
    }

    return true;
  }
}