import type { DatabaseSync } from "node:sqlite";
import { getDb, withTransaction } from "./db";

export interface Entity {
  name: string;
  entityType: string;
  observations: string[];
}

export interface Relation {
  from: string;
  to: string;
  relationType: string;
}

export interface KnowledgeGraph {
  entities: Entity[];
  relations: Relation[];
}

interface EntityRow {
  id: number;
  name: string;
  entity_type: string;
}

function escapeFtsQuery(query: string): string {
  const trimmed = query.trim();
  if (trimmed === "") {
    return "";
  }
  return `"${trimmed.replace(/"/g, '""')}"`;
}

export class KnowledgeGraphManager {
  // Opened lazily: LM Studio re-runs the tools provider on every config edit
  // (each keystroke in the Database Path field), and opening eagerly would
  // mkdir and create a DB file for every intermediate path typed.
  constructor(private readonly databasePath: string) {}

  private get db(): DatabaseSync {
    return getDb(this.databasePath);
  }

  async createEntities(entities: Entity[]): Promise<Entity[]> {
    const findByName = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const insertEntity = this.db.prepare("INSERT INTO entities (name, entity_type) VALUES (?, ?)");
    const insertObservation = this.db.prepare(
      "INSERT INTO observations (entity_id, content) VALUES (?, ?)"
    );

    const created: Entity[] = [];
    const seenInBatch = new Set<string>();

    withTransaction(this.db, () => {
      for (const entity of entities) {
        if (seenInBatch.has(entity.name)) continue;
        if (findByName.get(entity.name)) continue;

        seenInBatch.add(entity.name);
        const result = insertEntity.run(entity.name, entity.entityType);
        const entityId = result.lastInsertRowid as number;
        for (const observation of entity.observations) {
          insertObservation.run(entityId, observation);
        }
        created.push(entity);
      }
    });

    return created;
  }

  async createRelations(relations: Relation[]): Promise<Relation[]> {
    const findId = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const findExisting = this.db.prepare(
      "SELECT id FROM relations WHERE from_id = ? AND to_id = ? AND relation_type = ?"
    );
    const insertRelation = this.db.prepare(
      "INSERT INTO relations (from_id, to_id, relation_type) VALUES (?, ?, ?)"
    );

    // Validate every relation before writing anything, so a batch with one
    // bad reference creates nothing (matches the original JSONL manager).
    for (const relation of relations) {
      if (!findId.get(relation.from)) {
        throw new Error(`Entity with name ${relation.from} not found`);
      }
      if (!findId.get(relation.to)) {
        throw new Error(`Entity with name ${relation.to} not found`);
      }
    }

    const created: Relation[] = [];
    const seenInBatch = new Set<string>();

    withTransaction(this.db, () => {
      for (const relation of relations) {
        const key = JSON.stringify([relation.from, relation.to, relation.relationType]);
        if (seenInBatch.has(key)) continue;

        const fromRow = findId.get(relation.from) as { id: number };
        const toRow = findId.get(relation.to) as { id: number };
        if (findExisting.get(fromRow.id, toRow.id, relation.relationType)) continue;

        seenInBatch.add(key);
        insertRelation.run(fromRow.id, toRow.id, relation.relationType);
        created.push(relation);
      }
    });

    return created;
  }

  async addObservations(
    observations: { entityName: string; contents: string[] }[]
  ): Promise<{ entityName: string; addedObservations: string[] }[]> {
    const findId = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const findExistingObservation = this.db.prepare(
      "SELECT 1 FROM observations WHERE entity_id = ? AND content = ?"
    );
    const insertObservation = this.db.prepare(
      "INSERT INTO observations (entity_id, content) VALUES (?, ?)"
    );

    for (const o of observations) {
      if (!findId.get(o.entityName)) {
        throw new Error(`Entity with name ${o.entityName} not found`);
      }
    }

    const results: { entityName: string; addedObservations: string[] }[] = [];

    withTransaction(this.db, () => {
      for (const o of observations) {
        const row = findId.get(o.entityName) as { id: number };
        const added: string[] = [];
        for (const content of o.contents) {
          if (findExistingObservation.get(row.id, content)) continue;
          insertObservation.run(row.id, content);
          added.push(content);
        }
        results.push({ entityName: o.entityName, addedObservations: added });
      }
    });

    return results;
  }

  async deleteEntities(entityNames: string[]): Promise<{ deleted: string[]; notFound: string[] }> {
    const findId = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const deleteEntity = this.db.prepare("DELETE FROM entities WHERE id = ?");

    const deleted: string[] = [];
    const notFound: string[] = [];

    withTransaction(this.db, () => {
      for (const name of entityNames) {
        const row = findId.get(name) as { id: number } | undefined;
        if (!row) {
          notFound.push(name);
          continue;
        }
        deleteEntity.run(row.id); // ON DELETE CASCADE removes its observations and relations
        deleted.push(name);
      }
    });

    return { deleted, notFound };
  }

  async deleteObservations(
    deletions: { entityName: string; observations: string[] }[]
  ): Promise<{ deletedCount: number; missingEntities: string[] }> {
    const findId = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const deleteObservation = this.db.prepare(
      "DELETE FROM observations WHERE entity_id = ? AND content = ?"
    );

    let deletedCount = 0;
    const missingEntities: string[] = [];

    withTransaction(this.db, () => {
      for (const d of deletions) {
        const row = findId.get(d.entityName) as { id: number } | undefined;
        if (!row) {
          missingEntities.push(d.entityName);
          continue;
        }
        for (const content of d.observations) {
          const result = deleteObservation.run(row.id, content);
          deletedCount += Number(result.changes);
        }
      }
    });

    return { deletedCount, missingEntities };
  }

  async deleteRelations(relations: Relation[]): Promise<{ deletedCount: number }> {
    const findId = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const deleteRelation = this.db.prepare(
      "DELETE FROM relations WHERE from_id = ? AND to_id = ? AND relation_type = ?"
    );

    let deletedCount = 0;

    withTransaction(this.db, () => {
      for (const r of relations) {
        const fromRow = findId.get(r.from) as { id: number } | undefined;
        const toRow = findId.get(r.to) as { id: number } | undefined;
        if (!fromRow || !toRow) continue;
        const result = deleteRelation.run(fromRow.id, toRow.id, r.relationType);
        deletedCount += Number(result.changes);
      }
    });

    return { deletedCount };
  }

  async readGraph(): Promise<KnowledgeGraph> {
    const entityRows = this.db
      .prepare("SELECT id, name, entity_type FROM entities")
      .all() as unknown as EntityRow[];
    return this.buildGraphFromEntityRows(entityRows);
  }

  async searchNodes(query: string): Promise<KnowledgeGraph> {
    const likePattern = `%${query}%`;

    const nameOrTypeMatches = this.db
      .prepare(
        `SELECT id, name, entity_type FROM entities
         WHERE name LIKE ? OR entity_type LIKE ?`
      )
      .all(likePattern, likePattern) as unknown as EntityRow[];

    // Substring match against observation content (mirrors the original JSONL
    // manager's `.includes()` search), so short/partial queries like "Ac" still
    // match an observation containing "Acme Corp".
    const observationLikeMatches = this.db
      .prepare(
        `SELECT DISTINCT e.id, e.name, e.entity_type
         FROM observations o
         JOIN entities e ON e.id = o.entity_id
         WHERE o.content LIKE ?`
      )
      .all(likePattern) as unknown as EntityRow[];

    // FTS5 match as an additional (non-substring, tokenized) signal, kept for
    // relevance-style matching across word boundaries. Falls back to nothing
    // for a query that can't be parsed as an FTS expression even after
    // escaping (defensive; escapeFtsQuery already quotes the whole query).
    const ftsQuery = escapeFtsQuery(query);
    let observationFtsMatches: EntityRow[] = [];
    if (ftsQuery) {
      try {
        observationFtsMatches = this.db
          .prepare(
            `SELECT DISTINCT e.id, e.name, e.entity_type
             FROM observations_fts f
             JOIN observations o ON o.id = f.rowid
             JOIN entities e ON e.id = o.entity_id
             WHERE observations_fts MATCH ?`
          )
          .all(ftsQuery) as unknown as EntityRow[];
      } catch {
        observationFtsMatches = [];
      }
    }

    const byId = new Map<number, EntityRow>();
    for (const row of [...nameOrTypeMatches, ...observationLikeMatches, ...observationFtsMatches]) {
      byId.set(row.id, row);
    }

    return this.buildGraphFromEntityRows([...byId.values()]);
  }

  async openNodes(names: string[]): Promise<KnowledgeGraph> {
    if (names.length === 0) {
      return { entities: [], relations: [] };
    }

    const placeholders = names.map(() => "?").join(", ");
    const entityRows = this.db
      .prepare(`SELECT id, name, entity_type FROM entities WHERE name IN (${placeholders})`)
      .all(...names) as unknown as EntityRow[];

    return this.buildGraphFromEntityRows(entityRows);
  }

  private buildGraphFromEntityRows(entityRows: EntityRow[]): KnowledgeGraph {
    if (entityRows.length === 0) {
      return { entities: [], relations: [] };
    }

    const ids = entityRows.map(row => row.id);
    const placeholders = ids.map(() => "?").join(", ");

    const observationRows = this.db
      .prepare(`SELECT entity_id, content FROM observations WHERE entity_id IN (${placeholders})`)
      .all(...ids) as { entity_id: number; content: string }[];

    const observationsByEntity = new Map<number, string[]>();
    for (const row of observationRows) {
      const list = observationsByEntity.get(row.entity_id) ?? [];
      list.push(row.content);
      observationsByEntity.set(row.entity_id, list);
    }

    const entities: Entity[] = entityRows.map(row => ({
      name: row.name,
      entityType: row.entity_type,
      observations: observationsByEntity.get(row.id) ?? [],
    }));

    const relationRows = this.db
      .prepare(
        `SELECT e1.name AS from_name, e2.name AS to_name, r.relation_type
         FROM relations r
         JOIN entities e1 ON r.from_id = e1.id
         JOIN entities e2 ON r.to_id = e2.id
         WHERE r.from_id IN (${placeholders}) OR r.to_id IN (${placeholders})`
      )
      .all(...ids, ...ids) as { from_name: string; to_name: string; relation_type: string }[];

    const relations: Relation[] = relationRows.map(row => ({
      from: row.from_name,
      to: row.to_name,
      relationType: row.relation_type,
    }));

    return { entities, relations };
  }
}

import type { DatabaseSync } from "node:sqlite";
import { getDb, withTransaction } from "./db";

export interface Entity {
  name: string;
  entityType: string;
  observations: string[];
}

export interface Relation {
  from: string;
  to: string;
  relationType: string;
}

export interface KnowledgeGraph {
  entities: Entity[];
  relations: Relation[];
}

interface EntityRow {
  id: number;
  name: string;
  entity_type: string;
}

function escapeFtsQuery(query: string): string {
  const trimmed = query.trim();
  if (trimmed === "") {
    return "";
  }
  return `"${trimmed.replace(/"/g, '""')}"`;
}

export class KnowledgeGraphManager {
  // Opened lazily: LM Studio re-runs the tools provider on every config edit
  // (each keystroke in the Database Path field), and opening eagerly would
  // mkdir and create a DB file for every intermediate path typed.
  constructor(private readonly databasePath: string) {}

  private get db(): DatabaseSync {
    return getDb(this.databasePath);
  }

  async createEntities(entities: Entity[]): Promise<Entity[]> {
    const findByName = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const insertEntity = this.db.prepare("INSERT INTO entities (name, entity_type) VALUES (?, ?)");
    const insertObservation = this.db.prepare(
      "INSERT INTO observations (entity_id, content) VALUES (?, ?)"
    );

    const created: Entity[] = [];
    const seenInBatch = new Set<string>();

    withTransaction(this.db, () => {
      for (const entity of entities) {
        if (seenInBatch.has(entity.name)) continue;
        if (findByName.get(entity.name)) continue;

        seenInBatch.add(entity.name);
        const result = insertEntity.run(entity.name, entity.entityType);
        const entityId = result.lastInsertRowid as number;
        for (const observation of entity.observations) {
          insertObservation.run(entityId, observation);
        }
        created.push(entity);
      }
    });

    return created;
  }

  async createRelations(relations: Relation[]): Promise<Relation[]> {
    const findId = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const findExisting = this.db.prepare(
      "SELECT id FROM relations WHERE from_id = ? AND to_id = ? AND relation_type = ?"
    );
    const insertRelation = this.db.prepare(
      "INSERT INTO relations (from_id, to_id, relation_type) VALUES (?, ?, ?)"
    );

    // Validate every relation before writing anything, so a batch with one
    // bad reference creates nothing (matches the original JSONL manager).
    for (const relation of relations) {
      if (!findId.get(relation.from)) {
        throw new Error(`Entity with name ${relation.from} not found`);
      }
      if (!findId.get(relation.to)) {
        throw new Error(`Entity with name ${relation.to} not found`);
      }
    }

    const created: Relation[] = [];
    const seenInBatch = new Set<string>();

    withTransaction(this.db, () => {
      for (const relation of relations) {
        const key = JSON.stringify([relation.from, relation.to, relation.relationType]);
        if (seenInBatch.has(key)) continue;

        const fromRow = findId.get(relation.from) as { id: number };
        const toRow = findId.get(relation.to) as { id: number };
        if (findExisting.get(fromRow.id, toRow.id, relation.relationType)) continue;

        seenInBatch.add(key);
        insertRelation.run(fromRow.id, toRow.id, relation.relationType);
        created.push(relation);
      }
    });

    return created;
  }

  async addObservations(
    observations: { entityName: string; contents: string[] }[]
  ): Promise<{ entityName: string; addedObservations: string[] }[]> {
    const findId = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const findExistingObservation = this.db.prepare(
      "SELECT 1 FROM observations WHERE entity_id = ? AND content = ?"
    );
    const insertObservation = this.db.prepare(
      "INSERT INTO observations (entity_id, content) VALUES (?, ?)"
    );

    for (const o of observations) {
      if (!findId.get(o.entityName)) {
        throw new Error(`Entity with name ${o.entityName} not found`);
      }
    }

    const results: { entityName: string; addedObservations: string[] }[] = [];

    withTransaction(this.db, () => {
      for (const o of observations) {
        const row = findId.get(o.entityName) as { id: number };
        const added: string[] = [];
        for (const content of o.contents) {
          if (findExistingObservation.get(row.id, content)) continue;
          insertObservation.run(row.id, content);
          added.push(content);
        }
        results.push({ entityName: o.entityName, addedObservations: added });
      }
    });

    return results;
  }

  async deleteEntities(entityNames: string[]): Promise<{ deleted: string[]; notFound: string[] }> {
    const findId = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const deleteEntity = this.db.prepare("DELETE FROM entities WHERE id = ?");

    const deleted: string[] = [];
    const notFound: string[] = [];

    withTransaction(this.db, () => {
      for (const name of entityNames) {
        const row = findId.get(name) as { id: number } | undefined;
        if (!row) {
          notFound.push(name);
          continue;
        }
        deleteEntity.run(row.id); // ON DELETE CASCADE removes its observations and relations
        deleted.push(name);
      }
    });

    return { deleted, notFound };
  }

  async deleteObservations(
    deletions: { entityName: string; observations: string[] }[]
  ): Promise<{ deletedCount: number; missingEntities: string[] }> {
    const findId = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const deleteObservation = this.db.prepare(
      "DELETE FROM observations WHERE entity_id = ? AND content = ?"
    );

    let deletedCount = 0;
    const missingEntities: string[] = [];

    withTransaction(this.db, () => {
      for (const d of deletions) {
        const row = findId.get(d.entityName) as { id: number } | undefined;
        if (!row) {
          missingEntities.push(d.entityName);
          continue;
        }
        for (const content of d.observations) {
          const result = deleteObservation.run(row.id, content);
          deletedCount += Number(result.changes);
        }
      }
    });

    return { deletedCount, missingEntities };
  }

  async deleteRelations(relations: Relation[]): Promise<{ deletedCount: number }> {
    const findId = this.db.prepare("SELECT id FROM entities WHERE name = ?");
    const deleteRelation = this.db.prepare(
      "DELETE FROM relations WHERE from_id = ? AND to_id = ? AND relation_type = ?"
    );

    let deletedCount = 0;

    withTransaction(this.db, () => {
      for (const r of relations) {
        const fromRow = findId.get(r.from) as { id: number } | undefined;
        const toRow = findId.get(r.to) as { id: number } | undefined;
        if (!fromRow || !toRow) continue;
        const result = deleteRelation.run(fromRow.id, toRow.id, r.relationType);
        deletedCount += Number(result.changes);
      }
    });

    return { deletedCount };
  }

  async readGraph(): Promise<KnowledgeGraph> {
    const entityRows = this.db
      .prepare("SELECT id, name, entity_type FROM entities")
      .all() as unknown as EntityRow[];
    return this.buildGraphFromEntityRows(entityRows);
  }

  async searchNodes(query: string): Promise<KnowledgeGraph> {
    const likePattern = `%${query}%`;

    const nameOrTypeMatches = this.db
      .prepare(
        `SELECT id, name, entity_type FROM entities
         WHERE name LIKE ? OR entity_type LIKE ?`
      )
      .all(likePattern, likePattern) as unknown as EntityRow[];

    // Substring match against observation content (mirrors the original JSONL
    // manager's `.includes()` search), so short/partial queries like "Ac" still
    // match an observation containing "Acme Corp".
    const observationLikeMatches = this.db
      .prepare(
        `SELECT DISTINCT e.id, e.name, e.entity_type
         FROM observations o
         JOIN entities e ON e.id = o.entity_id
         WHERE o.content LIKE ?`
      )
      .all(likePattern) as unknown as EntityRow[];

    // FTS5 match as an additional (non-substring, tokenized) signal, kept for
    // relevance-style matching across word boundaries. Falls back to nothing
    // for a query that can't be parsed as an FTS expression even after
    // escaping (defensive; escapeFtsQuery already quotes the whole query).
    const ftsQuery = escapeFtsQuery(query);
    let observationFtsMatches: EntityRow[] = [];
    if (ftsQuery) {
      try {
        observationFtsMatches = this.db
          .prepare(
            `SELECT DISTINCT e.id, e.name, e.entity_type
             FROM observations_fts f
             JOIN observations o ON o.id = f.rowid
             JOIN entities e ON e.id = o.entity_id
             WHERE observations_fts MATCH ?`
          )
          .all(ftsQuery) as unknown as EntityRow[];
      } catch {
        observationFtsMatches = [];
      }
    }

    const byId = new Map<number, EntityRow>();
    for (const row of [...nameOrTypeMatches, ...observationLikeMatches, ...observationFtsMatches]) {
      byId.set(row.id, row);
    }

    return this.buildGraphFromEntityRows([...byId.values()]);
  }

  async openNodes(names: string[]): Promise<KnowledgeGraph> {
    if (names.length === 0) {
      return { entities: [], relations: [] };
    }

    const placeholders = names.map(() => "?").join(", ");
    const entityRows = this.db
      .prepare(`SELECT id, name, entity_type FROM entities WHERE name IN (${placeholders})`)
      .all(...names) as unknown as EntityRow[];

    return this.buildGraphFromEntityRows(entityRows);
  }

  private buildGraphFromEntityRows(entityRows: EntityRow[]): KnowledgeGraph {
    if (entityRows.length === 0) {
      return { entities: [], relations: [] };
    }

    const ids = entityRows.map(row => row.id);
    const placeholders = ids.map(() => "?").join(", ");

    const observationRows = this.db
      .prepare(`SELECT entity_id, content FROM observations WHERE entity_id IN (${placeholders})`)
      .all(...ids) as { entity_id: number; content: string }[];

    const observationsByEntity = new Map<number, string[]>();
    for (const row of observationRows) {
      const list = observationsByEntity.get(row.entity_id) ?? [];
      list.push(row.content);
      observationsByEntity.set(row.entity_id, list);
    }

    const entities: Entity[] = entityRows.map(row => ({
      name: row.name,
      entityType: row.entity_type,
      observations: observationsByEntity.get(row.id) ?? [],
    }));

    const relationRows = this.db
      .prepare(
        `SELECT e1.name AS from_name, e2.name AS to_name, r.relation_type
         FROM relations r
         JOIN entities e1 ON r.from_id = e1.id
         JOIN entities e2 ON r.to_id = e2.id
         WHERE r.from_id IN (${placeholders}) OR r.to_id IN (${placeholders})`
      )
      .all(...ids, ...ids) as { from_name: string; to_name: string; relation_type: string }[];

    const relations: Relation[] = relationRows.map(row => ({
      from: row.from_name,
      to: row.to_name,
      relationType: row.relation_type,
    }));

    return { entities, relations };
  }
}