import { type PluginContext, tool, type ToolsProvider } from "@lmstudio/sdk";
import { exec } from "child_process";
import { promisify } from "util";
import { copyFile, mkdtemp, rm } from "fs/promises";
import { existsSync } from "fs";
import { join } from "path";
import * as os from "os";
import { z } from "zod";
import Database from "better-sqlite3";

const execAsync = promisify(exec);
const ZOTERO_DB_PATH = "/home/arkantu/Zotero/zotero.sqlite";
const ZOTERO_STORAGE_PATH = "/home/arkantu/Zotero/storage";
const MAX_OUTPUT_CHARS = 3500;

function truncate(text: string, max: number = MAX_OUTPUT_CHARS): string {
  if (!text) return "";
  if (text.length <= max) return text;
  return text.substring(0, max) + "\n...[truncated output to protect context window]";
}

async function openZoteroDB(dbPath: string): Promise<any> {
  try {
    const db = new Database(dbPath, { readonly: true, fileMustExist: true });
    db.pragma("query_only = ON");
    return { db, tmpDir: null };
  } catch {
    const tmpDir = await mkdtemp(join(os.tmpdir(), "zotero-"));
    const tmpPath = join(tmpDir, "zotero.sqlite");
    await copyFile(dbPath, tmpPath);
    const db = new Database(tmpPath, { readonly: true, fileMustExist: true });
    db.pragma("query_only = ON");
    return { db, tmpDir };
  }
}

async function closeDB(db: any, tmpDir: string | null) {
  try { db.close(); } catch {}
  if (tmpDir) {
    await rm(tmpDir, { recursive: true, force: true }).catch(() => {});
  }
}

function resolveFieldIds(db: any): Record<string, number> {
  const rows = db.prepare(
    "SELECT fieldName, fieldID FROM fields WHERE fieldName IN ('title','date','DOI','abstractNote','url','ISBN','publisher')"
  ).all() as Array<{ fieldName: string; fieldID: number }>;
  const fids: Record<string, number> = {};
  for (const r of rows) fids[r.fieldName] = r.fieldID;
  return fids;
}

function resolveAuthorTypeId(db: any): number {
  const row = db.prepare(
    "SELECT creatorTypeID FROM creatorTypes WHERE creatorType = 'author'"
  ).get() as { creatorTypeID: number } | undefined;
  return row?.creatorTypeID ?? 1;
}

function buildMainQuery(fids: Record<string, number>, authorTypeId: number) {
  const titleId    = fids["title"]        ?? -1;
  const dateId     = fids["date"]         ?? -1;
  const doiId      = fids["DOI"]          ?? -1;
  const abstractId = fids["abstractNote"] ?? -1;

  return `
    SELECT
        i.itemID    AS id,
        i.key       AS key,
        tv.value    AS title,
        GROUP_CONCAT(c.lastName || ', ' || COALESCE(c.firstName,''), '; ') AS authors,
        dv.value    AS year,
        doiv.value  AS doi,
        av.value    AS abstract,
        att.path    AS pdf_path,
        atti.key    AS storage_key,
        COALESCE(GROUP_CONCAT(DISTINCT t.name), '') AS tags,
        COALESCE(GROUP_CONCAT(DISTINCT col.collectionName), '') AS collections
    FROM items i
    LEFT JOIN itemData    td   ON td.itemID   = i.itemID AND td.fieldID   = ${titleId}
    LEFT JOIN itemDataValues tv ON tv.valueID = td.valueID
    LEFT JOIN itemData    dd   ON dd.itemID   = i.itemID AND dd.fieldID   = ${dateId}
    LEFT JOIN itemDataValues dv ON dv.valueID = dd.valueID
    LEFT JOIN itemData    doid ON doid.itemID = i.itemID AND doid.fieldID = ${doiId}
    LEFT JOIN itemDataValues doiv ON doiv.valueID = doid.valueID
    LEFT JOIN itemData    ad   ON ad.itemID   = i.itemID AND ad.fieldID   = ${abstractId}
    LEFT JOIN itemDataValues av ON av.valueID = ad.valueID
    LEFT JOIN itemCreators ic  ON ic.itemID   = i.itemID AND ic.creatorTypeID = ${authorTypeId}
    LEFT JOIN creators c       ON c.creatorID = ic.creatorID
    LEFT JOIN (
        SELECT parentItemID, MIN(itemID) AS itemID, path
        FROM itemAttachments
        WHERE contentType = 'application/pdf' AND path LIKE 'storage:%'
        GROUP BY parentItemID
    ) att ON att.parentItemID = i.itemID
    LEFT JOIN items atti ON atti.itemID = att.itemID
    LEFT JOIN itemTags it2 ON it2.itemID = i.itemID
    LEFT JOIN tags t ON t.tagID = it2.tagID
    LEFT JOIN collectionItems ci ON ci.itemID = i.itemID
    LEFT JOIN collections col ON col.collectionID = ci.collectionID
    WHERE i.itemTypeID NOT IN (14, 26)
      AND tv.value IS NOT NULL
    GROUP BY i.itemID
  `;
}

// ── 1. Tool: zotero_search ──
const zotero_search = tool({
  name: "zotero_search",
  description: "Search papers in local Zotero library by keywords, title, author, abstract, tags or DOI.",
  parameters: {
    query: z.string().describe("Search keywords or title/author phrase."),
    limit: z.number().optional().describe("Max number of results to return (default: 5).")
  },
  implementation: async ({ query, limit = 5 }) => {
    let dbObj;
    try {
      dbObj = await openZoteroDB(ZOTERO_DB_PATH);
      const fids = resolveFieldIds(dbObj.db);
      const authorTypeId = resolveAuthorTypeId(dbObj.db);
      const sql = buildMainQuery(fids, authorTypeId);
      const rows = dbObj.db.prepare(sql).all() as any[];

      const terms = query.toLowerCase().split(/\s+/).filter(t => t.length > 2);
      const scored = rows.map(r => {
        let score = 0;
        const title    = (r.title    ?? "").toLowerCase();
        const authors  = (r.authors  ?? "").toLowerCase();
        const abstract = (r.abstract ?? "").toLowerCase();
        const tags     = (r.tags     ?? "").toLowerCase();
        const doi      = (r.doi      ?? "").toLowerCase();

        for (const term of terms) {
          if (title.includes(term))    score += 3;
          if (authors.includes(term))  score += 2;
          if (abstract.includes(term)) score += 1;
          if (tags.includes(term))     score += 1;
          if (doi.includes(term))      score += 3;
        }
        return { ...r, score };
      });

      const results = scored
        .filter(r => r.score > 0)
        .sort((a, b) => b.score - a.score)
        .slice(0, Math.min(limit, 10))
        .map(r => ({
          key: r.key,
          title: r.title,
          authors: r.authors ? r.authors.substring(0, 80) : "N/A",
          year: r.year ? String(r.year).substring(0, 4) : "N/A",
          has_pdf: Boolean(r.pdf_path && r.storage_key),
          doi: r.doi || null,
          abstract_preview: r.abstract ? truncate(r.abstract, 250) : null
        }));

      return { success: true, count: results.length, papers: results };
    } catch (err: any) {
      return { success: false, error: err.message };
    } finally {
      if (dbObj) await closeDB(dbObj.db, dbObj.tmpDir);
    }
  }
});

// ── 2. Tool: zotero_read_paper ──
const zotero_read_paper = tool({
  name: "zotero_read_paper",
  description: "Read full metadata and PDF text/abstract of a specific paper from Zotero using its item key or title.",
  parameters: {
    paper_key_or_title: z.string().describe("The Zotero item key (e.g. '8XYZW123') or exact/partial title.")
  },
  implementation: async ({ paper_key_or_title }) => {
    let dbObj;
    try {
      dbObj = await openZoteroDB(ZOTERO_DB_PATH);
      const fids = resolveFieldIds(dbObj.db);
      const authorTypeId = resolveAuthorTypeId(dbObj.db);
      const sql = buildMainQuery(fids, authorTypeId);
      const rows = dbObj.db.prepare(sql).all() as any[];

      const q = paper_key_or_title.toLowerCase().trim();
      const item = rows.find(r =>
        r.key?.toLowerCase() === q ||
        r.doi?.toLowerCase() === q ||
        r.title?.toLowerCase().includes(q)
      );

      if (!item) {
        return { success: false, error: `Paper not found for query: ${paper_key_or_title}` };
      }

      let pdfContent = "";
      if (item.pdf_path && item.storage_key) {
        const fileName = item.pdf_path.replace("storage:", "");
        const fullPdfPath = join(ZOTERO_STORAGE_PATH, item.storage_key, fileName);
        if (existsSync(fullPdfPath)) {
          try {
            const { stdout } = await execAsync(`pdftotext -l 10 "${fullPdfPath}" - 2>/dev/null || echo ""`, {
              timeout: 15000,
              maxBuffer: 2 * 1024 * 1024
            });
            if (stdout && stdout.trim()) {
              pdfContent = truncate(stdout.trim(), 3500);
            }
          } catch {}
        }
      }

      return {
        success: true,
        item: {
          key: item.key,
          title: item.title,
          authors: item.authors,
          year: item.year,
          doi: item.doi,
          tags: item.tags,
          collections: item.collections,
          abstract: item.abstract,
          pdf_text: pdfContent || "(PDF text not extracted or no PDF attached)"
        }
      };
    } catch (err: any) {
      return { success: false, error: err.message };
    } finally {
      if (dbObj) await closeDB(dbObj.db, dbObj.tmpDir);
    }
  }
});

// ── 3. Tool: zotero_stats ──
const zotero_stats = tool({
  name: "zotero_stats",
  description: "Get general summary and stats from your local Zotero library (collections, tags, paper count).",
  parameters: {},
  implementation: async () => {
    let dbObj;
    try {
      dbObj = await openZoteroDB(ZOTERO_DB_PATH);
      const total = (dbObj.db.prepare("SELECT COUNT(*) AS n FROM items WHERE itemTypeID NOT IN (14,26)").get() as any).n;
      const withPdf = (dbObj.db.prepare("SELECT COUNT(DISTINCT parentItemID) AS n FROM itemAttachments WHERE contentType='application/pdf' AND path LIKE 'storage:%'").get() as any).n;
      const cols = dbObj.db.prepare("SELECT collectionName FROM collections ORDER BY collectionName LIMIT 20").all() as Array<{ collectionName: string }>;
      const tags = dbObj.db.prepare("SELECT t.name, COUNT(*) AS n FROM tags t JOIN itemTags it ON it.tagID=t.tagID GROUP BY t.tagID ORDER BY n DESC LIMIT 15").all() as Array<{ name: string; n: number }>;

      return {
        success: true,
        total_items: total,
        items_with_pdf: withPdf,
        collections: cols.map(c => c.collectionName),
        top_tags: tags.map(t => `${t.name} (${t.n})`)
      };
    } catch (err: any) {
      return { success: false, error: err.message };
    } finally {
      if (dbObj) await closeDB(dbObj.db, dbObj.tmpDir);
    }
  }
});

const toolsProvider: ToolsProvider = async () => {
  return [zotero_search, zotero_read_paper, zotero_stats];
};

export async function main(context: PluginContext) {
  context.withToolsProvider(toolsProvider);
}

import { type PluginContext, tool, type ToolsProvider } from "@lmstudio/sdk";
import { exec } from "child_process";
import { promisify } from "util";
import { copyFile, mkdtemp, rm } from "fs/promises";
import { existsSync } from "fs";
import { join } from "path";
import * as os from "os";
import { z } from "zod";
import Database from "better-sqlite3";

const execAsync = promisify(exec);
const ZOTERO_DB_PATH = "/home/arkantu/Zotero/zotero.sqlite";
const ZOTERO_STORAGE_PATH = "/home/arkantu/Zotero/storage";
const MAX_OUTPUT_CHARS = 3500;

function truncate(text: string, max: number = MAX_OUTPUT_CHARS): string {
  if (!text) return "";
  if (text.length <= max) return text;
  return text.substring(0, max) + "\n...[truncated output to protect context window]";
}

async function openZoteroDB(dbPath: string): Promise<any> {
  try {
    const db = new Database(dbPath, { readonly: true, fileMustExist: true });
    db.pragma("query_only = ON");
    return { db, tmpDir: null };
  } catch {
    const tmpDir = await mkdtemp(join(os.tmpdir(), "zotero-"));
    const tmpPath = join(tmpDir, "zotero.sqlite");
    await copyFile(dbPath, tmpPath);
    const db = new Database(tmpPath, { readonly: true, fileMustExist: true });
    db.pragma("query_only = ON");
    return { db, tmpDir };
  }
}

async function closeDB(db: any, tmpDir: string | null) {
  try { db.close(); } catch {}
  if (tmpDir) {
    await rm(tmpDir, { recursive: true, force: true }).catch(() => {});
  }
}

function resolveFieldIds(db: any): Record<string, number> {
  const rows = db.prepare(
    "SELECT fieldName, fieldID FROM fields WHERE fieldName IN ('title','date','DOI','abstractNote','url','ISBN','publisher')"
  ).all() as Array<{ fieldName: string; fieldID: number }>;
  const fids: Record<string, number> = {};
  for (const r of rows) fids[r.fieldName] = r.fieldID;
  return fids;
}

function resolveAuthorTypeId(db: any): number {
  const row = db.prepare(
    "SELECT creatorTypeID FROM creatorTypes WHERE creatorType = 'author'"
  ).get() as { creatorTypeID: number } | undefined;
  return row?.creatorTypeID ?? 1;
}

function buildMainQuery(fids: Record<string, number>, authorTypeId: number) {
  const titleId    = fids["title"]        ?? -1;
  const dateId     = fids["date"]         ?? -1;
  const doiId      = fids["DOI"]          ?? -1;
  const abstractId = fids["abstractNote"] ?? -1;

  return `
    SELECT
        i.itemID    AS id,
        i.key       AS key,
        tv.value    AS title,
        GROUP_CONCAT(c.lastName || ', ' || COALESCE(c.firstName,''), '; ') AS authors,
        dv.value    AS year,
        doiv.value  AS doi,
        av.value    AS abstract,
        att.path    AS pdf_path,
        atti.key    AS storage_key,
        COALESCE(GROUP_CONCAT(DISTINCT t.name), '') AS tags,
        COALESCE(GROUP_CONCAT(DISTINCT col.collectionName), '') AS collections
    FROM items i
    LEFT JOIN itemData    td   ON td.itemID   = i.itemID AND td.fieldID   = ${titleId}
    LEFT JOIN itemDataValues tv ON tv.valueID = td.valueID
    LEFT JOIN itemData    dd   ON dd.itemID   = i.itemID AND dd.fieldID   = ${dateId}
    LEFT JOIN itemDataValues dv ON dv.valueID = dd.valueID
    LEFT JOIN itemData    doid ON doid.itemID = i.itemID AND doid.fieldID = ${doiId}
    LEFT JOIN itemDataValues doiv ON doiv.valueID = doid.valueID
    LEFT JOIN itemData    ad   ON ad.itemID   = i.itemID AND ad.fieldID   = ${abstractId}
    LEFT JOIN itemDataValues av ON av.valueID = ad.valueID
    LEFT JOIN itemCreators ic  ON ic.itemID   = i.itemID AND ic.creatorTypeID = ${authorTypeId}
    LEFT JOIN creators c       ON c.creatorID = ic.creatorID
    LEFT JOIN (
        SELECT parentItemID, MIN(itemID) AS itemID, path
        FROM itemAttachments
        WHERE contentType = 'application/pdf' AND path LIKE 'storage:%'
        GROUP BY parentItemID
    ) att ON att.parentItemID = i.itemID
    LEFT JOIN items atti ON atti.itemID = att.itemID
    LEFT JOIN itemTags it2 ON it2.itemID = i.itemID
    LEFT JOIN tags t ON t.tagID = it2.tagID
    LEFT JOIN collectionItems ci ON ci.itemID = i.itemID
    LEFT JOIN collections col ON col.collectionID = ci.collectionID
    WHERE i.itemTypeID NOT IN (14, 26)
      AND tv.value IS NOT NULL
    GROUP BY i.itemID
  `;
}

// ── 1. Tool: zotero_search ──
const zotero_search = tool({
  name: "zotero_search",
  description: "Search papers in local Zotero library by keywords, title, author, abstract, tags or DOI.",
  parameters: {
    query: z.string().describe("Search keywords or title/author phrase."),
    limit: z.number().optional().describe("Max number of results to return (default: 5).")
  },
  implementation: async ({ query, limit = 5 }) => {
    let dbObj;
    try {
      dbObj = await openZoteroDB(ZOTERO_DB_PATH);
      const fids = resolveFieldIds(dbObj.db);
      const authorTypeId = resolveAuthorTypeId(dbObj.db);
      const sql = buildMainQuery(fids, authorTypeId);
      const rows = dbObj.db.prepare(sql).all() as any[];

      const terms = query.toLowerCase().split(/\s+/).filter(t => t.length > 2);
      const scored = rows.map(r => {
        let score = 0;
        const title    = (r.title    ?? "").toLowerCase();
        const authors  = (r.authors  ?? "").toLowerCase();
        const abstract = (r.abstract ?? "").toLowerCase();
        const tags     = (r.tags     ?? "").toLowerCase();
        const doi      = (r.doi      ?? "").toLowerCase();

        for (const term of terms) {
          if (title.includes(term))    score += 3;
          if (authors.includes(term))  score += 2;
          if (abstract.includes(term)) score += 1;
          if (tags.includes(term))     score += 1;
          if (doi.includes(term))      score += 3;
        }
        return { ...r, score };
      });

      const results = scored
        .filter(r => r.score > 0)
        .sort((a, b) => b.score - a.score)
        .slice(0, Math.min(limit, 10))
        .map(r => ({
          key: r.key,
          title: r.title,
          authors: r.authors ? r.authors.substring(0, 80) : "N/A",
          year: r.year ? String(r.year).substring(0, 4) : "N/A",
          has_pdf: Boolean(r.pdf_path && r.storage_key),
          doi: r.doi || null,
          abstract_preview: r.abstract ? truncate(r.abstract, 250) : null
        }));

      return { success: true, count: results.length, papers: results };
    } catch (err: any) {
      return { success: false, error: err.message };
    } finally {
      if (dbObj) await closeDB(dbObj.db, dbObj.tmpDir);
    }
  }
});

// ── 2. Tool: zotero_read_paper ──
const zotero_read_paper = tool({
  name: "zotero_read_paper",
  description: "Read full metadata and PDF text/abstract of a specific paper from Zotero using its item key or title.",
  parameters: {
    paper_key_or_title: z.string().describe("The Zotero item key (e.g. '8XYZW123') or exact/partial title.")
  },
  implementation: async ({ paper_key_or_title }) => {
    let dbObj;
    try {
      dbObj = await openZoteroDB(ZOTERO_DB_PATH);
      const fids = resolveFieldIds(dbObj.db);
      const authorTypeId = resolveAuthorTypeId(dbObj.db);
      const sql = buildMainQuery(fids, authorTypeId);
      const rows = dbObj.db.prepare(sql).all() as any[];

      const q = paper_key_or_title.toLowerCase().trim();
      const item = rows.find(r =>
        r.key?.toLowerCase() === q ||
        r.doi?.toLowerCase() === q ||
        r.title?.toLowerCase().includes(q)
      );

      if (!item) {
        return { success: false, error: `Paper not found for query: ${paper_key_or_title}` };
      }

      let pdfContent = "";
      if (item.pdf_path && item.storage_key) {
        const fileName = item.pdf_path.replace("storage:", "");
        const fullPdfPath = join(ZOTERO_STORAGE_PATH, item.storage_key, fileName);
        if (existsSync(fullPdfPath)) {
          try {
            const { stdout } = await execAsync(`pdftotext -l 10 "${fullPdfPath}" - 2>/dev/null || echo ""`, {
              timeout: 15000,
              maxBuffer: 2 * 1024 * 1024
            });
            if (stdout && stdout.trim()) {
              pdfContent = truncate(stdout.trim(), 3500);
            }
          } catch {}
        }
      }

      return {
        success: true,
        item: {
          key: item.key,
          title: item.title,
          authors: item.authors,
          year: item.year,
          doi: item.doi,
          tags: item.tags,
          collections: item.collections,
          abstract: item.abstract,
          pdf_text: pdfContent || "(PDF text not extracted or no PDF attached)"
        }
      };
    } catch (err: any) {
      return { success: false, error: err.message };
    } finally {
      if (dbObj) await closeDB(dbObj.db, dbObj.tmpDir);
    }
  }
});

// ── 3. Tool: zotero_stats ──
const zotero_stats = tool({
  name: "zotero_stats",
  description: "Get general summary and stats from your local Zotero library (collections, tags, paper count).",
  parameters: {},
  implementation: async () => {
    let dbObj;
    try {
      dbObj = await openZoteroDB(ZOTERO_DB_PATH);
      const total = (dbObj.db.prepare("SELECT COUNT(*) AS n FROM items WHERE itemTypeID NOT IN (14,26)").get() as any).n;
      const withPdf = (dbObj.db.prepare("SELECT COUNT(DISTINCT parentItemID) AS n FROM itemAttachments WHERE contentType='application/pdf' AND path LIKE 'storage:%'").get() as any).n;
      const cols = dbObj.db.prepare("SELECT collectionName FROM collections ORDER BY collectionName LIMIT 20").all() as Array<{ collectionName: string }>;
      const tags = dbObj.db.prepare("SELECT t.name, COUNT(*) AS n FROM tags t JOIN itemTags it ON it.tagID=t.tagID GROUP BY t.tagID ORDER BY n DESC LIMIT 15").all() as Array<{ name: string; n: number }>;

      return {
        success: true,
        total_items: total,
        items_with_pdf: withPdf,
        collections: cols.map(c => c.collectionName),
        top_tags: tags.map(t => `${t.name} (${t.n})`)
      };
    } catch (err: any) {
      return { success: false, error: err.message };
    } finally {
      if (dbObj) await closeDB(dbObj.db, dbObj.tmpDir);
    }
  }
});

const toolsProvider: ToolsProvider = async () => {
  return [zotero_search, zotero_read_paper, zotero_stats];
};

export async function main(context: PluginContext) {
  context.withToolsProvider(toolsProvider);
}