/*
This file is part of the Notesnook project (https://notesnook.com/)

Copyright (C) 2023 Streetwriters (Private) Limited

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

import { match } from "fuzzyjs";
import Database from "./index.js";
import { Item, Note, TrashItem } from "../types.js";
import { DatabaseSchema, RawDatabaseSchema } from "../database/index.js";
import { AnyColumnWithTable, Kysely, sql } from "@streetwriters/kysely";
import { FilteredSelector } from "../database/sql-collection.js";
import { VirtualizedGrouping } from "../utils/virtualized-grouping.js";
import { logger } from "../logger.js";
import { rebuildSearchIndex } from "../database/fts.js";
import { transformQuery } from "../utils/query-transformer.js";

type SearchResults<T> = {
  sorted: (limit?: number) => Promise<VirtualizedGrouping<T>>;
  items: (limit?: number) => Promise<T[]>;
  ids: () => Promise<string[]>;
};

type FuzzySearchField<T> = {
  weight?: number;
  name: keyof T;
  column: AnyColumnWithTable<DatabaseSchema, keyof DatabaseSchema>;
};
export default class Lookup {
  constructor(private readonly db: Database) {}

  notes(query: string, notes?: FilteredSelector<Note>): SearchResults<Note> {
    return this.toSearchResults(async (limit) => {
      const db = this.db.sql() as unknown as Kysely<RawDatabaseSchema>;
      const excludedIds = this.db.trash.cache.notes;

      if (query.length <= 2) {
        const results = await db
          .selectFrom((eb) =>
            eb
              .selectFrom("notes")
              .$if(!!notes, (eb) =>
                eb.where("id", "in", notes!.filter.select("id"))
              )
              .$if(excludedIds.length > 0, (eb) =>
                eb.where("id", "not in", excludedIds)
              )
              .where("title", "like", `%${query}%`)
              .select(["id"])
              .unionAll((eb) =>
                eb
                  .selectFrom("content")
                  .$if(!!notes, (eb) =>
                    eb.where("id", "in", notes!.filter.select("id"))
                  )
                  .$if(excludedIds.length > 0, (eb) =>
                    eb.where("id", "not in", excludedIds)
                  )
                  .where("locked", "!=", true)
                  .where("data", "like", `%${query}%`)
                  .select(["noteId as id"])
                  .$castTo<{ id: string; dateCreated: number }>()
              )
              .as("results")
          )
          .select(["results.id"])
          .groupBy("results.id")
          .$if(!!limit, (eb) => eb.limit(limit!))
          // filter out ids that have no note against them
          .where(
            "results.id",
            "in",
            (notes || this.db.notes.all).filter.select("id")
          )
          .execute()
          .catch((e) => {
            logger.error(e, `Error while searching`, { query });
            return [];
          });
        return results.map((r) => r.id);
      }

      query = transformQuery(query);
      const results = await db
        .selectFrom((eb) =>
          eb
            .selectFrom("notes_fts")
            .$if(!!notes, (eb) =>
              eb.where("id", "in", notes!.filter.select("id"))
            )
            .$if(excludedIds.length > 0, (eb) =>
              eb.where("id", "not in", excludedIds)
            )
            .where("title", "match", query)
            .select(["id", sql<number>`rank * 10`.as("rank")])
            .unionAll((eb) =>
              eb
                .selectFrom("content_fts")
                .$if(!!notes, (eb) =>
                  eb.where("id", "in", notes!.filter.select("id"))
                )
                .$if(excludedIds.length > 0, (eb) =>
                  eb.where("id", "not in", excludedIds)
                )
                .where("data", "match", query)
                .select(["noteId as id", "rank"])
                .$castTo<{ id: string; rank: number }>()
            )
            .as("results")
        )
        .select(["results.id"])
        .groupBy("results.id")
        .orderBy(sql`SUM(results.rank)`, "desc")
        .$if(!!limit, (eb) => eb.limit(limit!))

        // filter out ids that have no note against them
        .where(
          "results.id",
          "in",
          (notes || this.db.notes.all).filter.select("id")
        )
        .execute()
        .catch((e) => {
          logger.error(e, `Error while searching`, { query });
          return [];
        });
      return results.map((r) => r.id);
    }, notes || this.db.notes.all);
  }

  notebooks(query: string) {
    return this.search(this.db.notebooks.all, query, [
      { name: "id", column: "notebooks.id", weight: -100 },
      { name: "title", column: "notebooks.title", weight: 10 },
      { name: "description", column: "notebooks.description" }
    ]);
  }

  tags(query: string) {
    return this.search(this.db.tags.all, query, [
      { name: "id", column: "tags.id", weight: -100 },
      { name: "title", column: "tags.title" }
    ]);
  }

  reminders(query: string) {
    return this.search(this.db.reminders.all, query, [
      { name: "id", column: "reminders.id", weight: -100 },
      { name: "title", column: "reminders.title", weight: 10 },
      { name: "description", column: "reminders.description" }
    ]);
  }

  trash(query: string): SearchResults<TrashItem> {
    return {
      sorted: async (limit?: number) => {
        const { ids, items } = await this.filterTrash(query, limit);
        return new VirtualizedGrouping<TrashItem>(
          ids.length,
          this.db.options.batchSize,
          () => Promise.resolve(ids),
          async (start, end) => {
            return {
              ids: ids.slice(start, end),
              items: items.slice(start, end)
            };
          }
        );
      },
      items: async (limit?: number) => {
        const { items } = await this.filterTrash(query, limit);
        return items;
      },
      ids: () => this.filterTrash(query).then(({ ids }) => ids)
    };
  }

  attachments(query: string) {
    return this.search(this.db.attachments.all, query, [
      { name: "id", column: "attachments.id", weight: -100 },
      { name: "filename", column: "attachments.filename", weight: 5 },
      { name: "mimeType", column: "attachments.mimeType" },
      { name: "hash", column: "attachments.hash" }
    ]);
  }

  private search<T extends Item>(
    selector: FilteredSelector<T>,
    query: string,
    fields: FuzzySearchField<T>[]
  ) {
    return this.toSearchResults(
      (limit) => this.filter(selector, query, fields, limit),
      selector
    );
  }

  private async filter<T extends Item>(
    selector: FilteredSelector<T>,
    query: string,
    fields: FuzzySearchField<T>[],
    limit?: number
  ) {
    const results: Map<string, number> = new Map();
    const columns = fields.map((f) => f.column);
    for await (const item of selector.fields(columns)) {
      if (limit && results.size >= limit) break;

      for (const field of fields) {
        const result = match(query, `${item[field.name]}`);
        if (result.match) {
          const oldScore = results.get(item.id) || 0;
          results.set(item.id, oldScore + result.score * (field.weight || 1));
        }
      }
    }
    selector.fields([]);

    return Array.from(results.entries())
      .sort((a, b) => a[1] - b[1])
      .map((a) => a[0]);
  }

  private toSearchResults<T extends Item>(
    ids: (limit?: number) => Promise<string[]>,
    selector: FilteredSelector<T>
  ): SearchResults<T> {
    return {
      sorted: async (limit?: number) =>
        this.toVirtualizedGrouping(await ids(limit), selector),
      items: async (limit?: number) => this.toItems(await ids(limit), selector),
      ids
    };
  }

  private async filterTrash(query: string, limit?: number) {
    const items = await this.db.trash.all();

    const results: Map<string, { rank: number; item: TrashItem }> = new Map();
    for (const item of items) {
      if (limit && results.size >= limit) break;

      const result = match(query, item.title);
      if (result.match) {
        results.set(item.id, { rank: result.score, item });
      }
    }

    const sorted = Array.from(results.entries()).sort(
      (a, b) => a[1].rank - b[1].rank
    );
    return {
      ids: sorted.map((a) => a[0]),
      items: sorted.map((a) => a[1].item)
    };
  }

  private toVirtualizedGrouping<T extends Item>(
    ids: string[],
    selector: FilteredSelector<T>
  ) {
    return new VirtualizedGrouping<T>(
      ids.length,
      this.db.options.batchSize,
      () => Promise.resolve(ids),
      async (start, end) => {
        const items = await selector.records(ids);
        return {
          ids: ids.slice(start, end),
          items: Object.values(items).slice(start, end)
        };
      }
    );
  }

  private toItems<T extends Item>(
    ids: string[],
    selector: FilteredSelector<T>
  ) {
    if (!ids.length) return [];
    return selector.items(ids);
  }

  async rebuild() {
    const db = this.db.sql() as unknown as Kysely<RawDatabaseSchema>;
    await rebuildSearchIndex(db);
  }
}
