const sqlite3 = require("sqlite3").verbose();
const path = require("path");
const { app } = require("electron");

const MODEL_TABLE_CREATION = `
CREATE TABLE IF NOT EXISTS models (
  id TEXT PRIMARY KEY,
  name TEXT NOT NULL,
  short_description TEXT NOT NULL,
  avatar_url TEXT,
  long_description TEXT NOT NULL,
  author TEXT NOT NULL,
  version TEXT NOT NULL,
  model_url TEXT NOT NULL,
  nsfw INTEGER NOT NULL,
  tags TEXT NOT NULL,
  default_greeting TEXT NOT NULL,
  type TEXT NOT NULL,
  created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
  updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
);`;

const MODEL_VERSION_TABLE_CREATION = `
CREATE TABLE IF NOT EXISTS model_versions (
  id TEXT PRIMARY KEY,
  name TEXT NOT NULL,
  quant_method TEXT NOT NULL,
  bits INTEGER NOT NULL,
  size INTEGER NOT NULL,
  max_ram_required INTEGER NOT NULL,
  usecase TEXT NOT NULL,
  download_link TEXT NOT NULL,
  model_id TEXT NOT NULL,
  start_download_at INTEGER DEFAULT -1,
  finish_download_at INTEGER DEFAULT -1,
  created_at DATETIME DEFAULT CURRENT_TIMESTAMP
);`;

const MODEL_TABLE_INSERTION = `
INSERT OR IGNORE INTO models (
  id,
  name,
  short_description,
  avatar_url,
  long_description,
  author,
  version,
  model_url,
  nsfw,
  tags,
  default_greeting,
  type
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?)`;

const MODEL_VERSION_TABLE_INSERTION = `
INSERT INTO model_versions (
  id,
  name,
  quant_method,
  bits,
  size,
  max_ram_required,
  usecase,
  download_link,
  model_id,
  start_download_at
) VALUES (?,?,?,?,?,?,?,?,?,?)`;

const getDbPath = () => {
  return path.join(app.getPath("userData"), "jan.db");
};

function init() {
  const db = new sqlite3.Database(getDbPath());
  console.debug(`Database located at ${getDbPath()}`);

  db.serialize(() => {
    db.run(MODEL_TABLE_CREATION);
    db.run(MODEL_VERSION_TABLE_CREATION);
    db.run(
      "CREATE TABLE IF NOT EXISTS conversations ( id INTEGER PRIMARY KEY, name TEXT, model_id TEXT, image TEXT, message TEXT, created_at DATETIME DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP);"
    );
    db.run(
      "CREATE TABLE IF NOT EXISTS messages ( id INTEGER PRIMARY KEY, name TEXT, conversation_id INTEGER, user TEXT, message TEXT, created_at DATETIME DEFAULT CURRENT_TIMESTAMP, updated_at DATETIME DEFAULT CURRENT_TIMESTAMP);"
    );
  });

  db.close();
}

/**
 * Store a model in the database when user start downloading it
 *
 * @param params: { model, modelVersion }
 */
function storeModel(params: any) {
  return new Promise((res) => {
    const db = new sqlite3.Database(getDbPath());
    console.debug("Inserting", JSON.stringify(params));

    const model = params.model;
    const modelTags = model.tags.join(",");
    const modelVersion = params.modelVersion;

    db.serialize(() => {
      const stmt = db.prepare(MODEL_TABLE_INSERTION);
      stmt.run(
        model.id,
        model.name,
        model.shortDescription,
        model.avatarUrl,
        model.longDescription,
        model.author,
        model.version,
        model.modelUrl,
        model.nsfw,
        modelTags,
        model.greeting,
        model.type,
        function (err: any) {
          if (err) {
            // Handle the insertion error here
            console.error(err.message);
            res(undefined);
            return;
          }
          // @ts-ignoreF
          const id = this.lastID;
          res(id);
          return;
        }
      );

      // insert modelVersion to MODEL_VERSION_TABLE_INSERTION
      const stmt2 = db.prepare(MODEL_VERSION_TABLE_INSERTION);
      stmt2.run(
        modelVersion.id,
        modelVersion.name,
        modelVersion.quantMethod,
        modelVersion.bits,
        modelVersion.size,
        modelVersion.maxRamRequired,
        modelVersion.usecase,
        modelVersion.downloadLink,
        model.id,
        modelVersion.startDownloadAt,
        function (err: any) {
          if (err) {
            // Handle the insertion error here
            console.error(err.message);
            res(undefined);
            return;
          }
          // @ts-ignoreF
          const id = this.lastID;
          res(id);
          return;
        }
      );
      stmt.finalize();
      stmt2.finalize();
    });

    db.close();
  });
}

/**
 * Update the finished download time of a model
 *
 * @param model Product
 */
function updateFinishedDownloadAt(modelVersionId: string) {
  return new Promise((res) => {
    const db = new sqlite3.Database(getDbPath());
    const time = Date.now();
    console.debug(
      `Updating finished downloaded model version ${modelVersionId}`
    );
    const stmt = `UPDATE model_versions SET finish_download_at = ? WHERE id = ?`;
    db.run(stmt, [time, modelVersionId], (err: any) => {
      if (err) {
        console.log(err);
        res(undefined);
      } else {
        console.log("Updated 1 row");
        res("Updated");
      }
    });

    db.close();
  });
}

/**
 * Get all unfinished models from the database
 */
function getUnfinishedDownloadModels() {
  return new Promise((res) => {
    const db = new sqlite3.Database(getDbPath());

    const query = `SELECT * FROM model_versions WHERE finish_download_at = -1 ORDER BY start_download_at DESC`;
    db.all(query, (err: Error, row: any) => {
      if (row) {
        res(row);
      } else {
        res([]);
      }
    });
    db.close();
  });
}

async function getFinishedDownloadModels() {
  const db = new sqlite3.Database(getDbPath());

  const query = `SELECT * FROM model_versions WHERE finish_download_at != -1 ORDER BY finish_download_at DESC`;
  const modelVersions: any = await new Promise((resolve, reject) => {
    db.all(query, (err: Error, rows: any[]) => {
      if (err) {
        reject(err);
      } else {
        resolve(rows);
      }
    });
  });

  const models = await Promise.all(
    modelVersions.map(async (modelVersion) => {
      const modelQuery = `SELECT * FROM models WHERE id = ?`;
      return new Promise((resolve, reject) => {
        db.get(modelQuery, [modelVersion.model_id], (err: Error, row: any) => {
          if (err) {
            reject(err);
          } else {
            resolve(row);
          }
        });
      });
    })
  );

  const downloadedModels = [];
  modelVersions.forEach((modelVersion: any) => {
    const model = models.find((m: any) => m.id === modelVersion.model_id);

    if (!model) {
      return;
    }

    const assistantModel = {
      id: modelVersion.id,
      name: modelVersion.name,
      quantMethod: modelVersion.quant_method,
      bits: modelVersion.bits,
      size: modelVersion.size,
      maxRamRequired: modelVersion.max_ram_required,
      usecase: modelVersion.usecase,
      downloadLink: modelVersion.download_link,
      startDownloadAt: modelVersion.start_download_at,
      finishDownloadAt: modelVersion.finish_download_at,
      productId: model.id,
      productName: model.name,
      shortDescription: model.short_description,
      longDescription: model.long_description,
      avatarUrl: model.avatar_url,
      author: model.author,
      version: model.version,
      modelUrl: model.model_url,
      nsfw: model.nsfw === 0 ? false : true,
      greeting: model.default_greeting,
      type: model.type,
      createdAt: new Date(model.created_at).getTime(),
      updatedAt: new Date(model.updated_at ?? "").getTime(),
      status: "",
      releaseDate: -1,
      tags: model.tags.split(","),
    };
    downloadedModels.push(assistantModel);
  });

  db.close();

  return downloadedModels;
}

function deleteDownloadModel(modelId: string) {
  return new Promise((res) => {
    const db = new sqlite3.Database(getDbPath());
    console.debug(`Deleting ${modelId}`);
    db.serialize(() => {
      const stmt = db.prepare("DELETE FROM model_versions WHERE id = ?");
      stmt.run(modelId);
      stmt.finalize();
      res(modelId);
    });

    db.close();
  });
}

async function fetchModelVersion(db: any, versionId: string) {
  return new Promise((resolve, reject) => {
    db.get(
      "SELECT * FROM model_versions WHERE id = ?",
      [versionId],
      (err, row) => {
        if (err) {
          reject(err);
        } else {
          if (row) {
            const product = {
              id: row.id,
              slug: row.slug,
              name: row.name,
              description: row.description,
              avatarUrl: row.avatar_url,
              longDescription: row.long_description,
              technicalDescription: row.technical_description,
              author: row.author,
              version: row.version,
              modelUrl: row.model_url,
              nsfw: row.nsfw,
              greeting: row.greeting,
              type: row.type,
              inputs: row.inputs,
              outputs: row.outputs,
              createdAt: new Date(row.created_at),
              updatedAt: new Date(row.updated_at),
              fileName: row.file_name,
              downloadUrl: row.download_url,
            };
            resolve(product);
          } else {
            resolve(undefined);
          }
        }
      }
    );
  });
}

async function fetchModel(db: any, modelId: string) {
  return new Promise((resolve, reject) => {
    db.get("SELECT * FROM models WHERE id = ?", [modelId], (err, row) => {
      if (err) {
        reject(err);
      } else {
        resolve(row);
      }
    });
  });
}

const getModelById = async (versionId: string): Promise<any | undefined> => {
  const db = new sqlite3.Database(getDbPath());
  const modelVersion: any | undefined = await fetchModelVersion(db, versionId);
  if (!modelVersion) return undefined;
  const model: any | undefined = await fetchModel(db, modelVersion.model_id);
  if (!model) return undefined;

  const assistantModel = {
    id: modelVersion.id,
    name: modelVersion.name,
    quantMethod: modelVersion.quant_method,
    bits: modelVersion.bits,
    size: modelVersion.size,
    maxRamRequired: modelVersion.max_ram_required,
    usecase: modelVersion.usecase,
    downloadLink: modelVersion.download_link,
    startDownloadAt: modelVersion.start_download_at,
    finishDownloadAt: modelVersion.finish_download_at,
    productId: model.id,
    productName: model.name,
    shortDescription: model.short_description,
    longDescription: model.long_description,
    avatarUrl: model.avatar_url,
    author: model.author,
    version: model.version,
    modelUrl: model.model_url,
    nsfw: model.nsfw === 0 ? false : true,
    greeting: model.default_greeting,
    type: model.type,
    createdAt: new Date(model.created_at).getTime(),
    updatedAt: new Date(model.updated_at ?? "").getTime(),
    status: "",
    releaseDate: -1,
    tags: model.tags.split(","),
  };

  return assistantModel;
};

function getConversations() {
  return new Promise((res) => {
    const db = new sqlite3.Database(getDbPath());

    db.all(
      "SELECT * FROM conversations ORDER BY updated_at DESC",
      (err: any, row: any) => {
        res(row);
      }
    );
    db.close();
  });
}

function storeConversation(conversation: any): Promise<number | undefined> {
  return new Promise((res) => {
    const db = new sqlite3.Database(getDbPath());

    db.serialize(() => {
      const stmt = db.prepare(
        "INSERT INTO conversations (name, model_id, image, message) VALUES (?, ?, ?, ?)"
      );
      stmt.run(
        conversation.name,
        conversation.model_id,
        conversation.image,
        conversation.message,
        function (err: any) {
          if (err) {
            // Handle the insertion error here
            console.error(err.message);
            res(undefined);
            return;
          }
          // @ts-ignoreF
          const id = this.lastID;
          res(id);
          return;
        }
      );
      stmt.finalize();
    });

    db.close();
  });
}

function storeMessage(message: any): Promise<number | undefined> {
  return new Promise((res) => {
    const db = new sqlite3.Database(getDbPath());

    db.serialize(() => {
      const stmt = db.prepare(
        "INSERT INTO messages (name, conversation_id, user, message) VALUES (?, ?, ?, ?)"
      );
      stmt.run(
        message.name,
        message.conversation_id,
        message.user,
        message.message,
        function (err: any) {
          if (err) {
            // Handle the insertion error here
            console.error(err.message);
            res(undefined);
            return;
          }
          //@ts-ignore
          const id = this.lastID;
          res(id);
          return;
        }
      );
      stmt.finalize();
    });

    db.close();
  });
}

function updateMessage(message: any): Promise<number | undefined> {
  return new Promise((res) => {
    const db = new sqlite3.Database(getDbPath());

    db.serialize(() => {
      const stmt = db.prepare(
        "UPDATE messages SET message = ?, updated_at = ? WHERE id = ?"
      );
      stmt.run(message.message, message.updated_at, message.id);
      stmt.finalize();
      res(message.id);
    });

    db.close();
  });
}

function deleteConversation(id: any) {
  return new Promise((res) => {
    const db = new sqlite3.Database(getDbPath());

    db.serialize(() => {
      const deleteConv = db.prepare("DELETE FROM conversations WHERE id = ?");
      deleteConv.run(id);
      deleteConv.finalize();
      const deleteMessages = db.prepare(
        "DELETE FROM messages WHERE conversation_id = ?"
      );
      deleteMessages.run(id);
      deleteMessages.finalize();
      res(id);
    });

    db.close();
  });
}

function getConversationMessages(conversation_id: any) {
  return new Promise((res) => {
    const db = new sqlite3.Database(getDbPath());

    const query = `SELECT * FROM messages WHERE conversation_id = ${conversation_id} ORDER BY id DESC`;
    db.all(query, (err: Error, row: any) => {
      res(row);
    });
    db.close();
  });
}

module.exports = {
  init,
  getConversations,
  deleteConversation,
  storeConversation,
  storeMessage,
  updateMessage,
  getConversationMessages,
  storeModel,
  updateFinishedDownloadAt,
  getUnfinishedDownloadModels,
  getFinishedDownloadModels,
  deleteDownloadModel,
  getModelById,
};