import {
    CallToolRequestSchema,
    ListToolsRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";
import { db, EMBEDDING_DIM, knexDb } from "./db.js";
import { llamaCompletion, llamaEmbedding } from "./llama-client.js";
import { getToolDefinitions } from "./mcp-tools.js";

function assertEmbeddingDim(vector) {
  if (!Array.isArray(vector) || vector.length !== EMBEDDING_DIM) {
    throw new Error(`Embedding dimension mismatch. Expected ${EMBEDDING_DIM}`);
  }
}

export function registerMcpHandlers(server) {
  server.setRequestHandler(ListToolsRequestSchema, async () => {
    return {
      tools: getToolDefinitions(),
    };
  });

  server.setRequestHandler(CallToolRequestSchema, async (request) => {
    switch (request.params.name) {
      case "add_item_text": {
        const { content, path } = request.params.arguments;
        const embedding = await llamaEmbedding(content);
        assertEmbeddingDim(embedding);
        const insertIds = await knexDb("items").insert({ content, path });
        const id = Array.isArray(insertIds) ? insertIds[0] : insertIds;
        db.prepare("INSERT INTO vec_items(id, embedding) VALUES (?, ?)")
          .run(id, new Float32Array(embedding));
        return { content: [{ type: "text", text: `Added item with id ${id}` }] };
      }
      case "add_item": {
        const { content, path, vector } = request.params.arguments;
        assertEmbeddingDim(vector);
        const insertIds = await knexDb("items").insert({ content, path });
        const id = Array.isArray(insertIds) ? insertIds[0] : insertIds;
        db.prepare("INSERT INTO vec_items(id, embedding) VALUES (?, ?)")
          .run(id, new Float32Array(vector));
        return { content: [{ type: "text", text: `Added item with id ${id}` }] };
      }
      case "search_text": {
        const { content } = request.params.arguments;
        const embedding = await llamaEmbedding(content);
        assertEmbeddingDim(embedding);
        const results = db.prepare(`
          SELECT
            i.content,
            i.path,
            i.created_at,
            i.updated_at,
            v.distance
          FROM vec_items v
          JOIN items i ON v.id = i.id
          WHERE embedding MATCH ?
          ORDER BY distance
          LIMIT 5
        `).all(new Float32Array(embedding));

        return {
          content: [{ type: "text", text: JSON.stringify(results, null, 2) }],
        };
      }
      case "search_vector": {
        const { vector } = request.params.arguments;
        assertEmbeddingDim(vector);
        const results = db.prepare(`
          SELECT
            i.content,
            i.path,
            i.created_at,
            i.updated_at,
            v.distance
          FROM vec_items v
          JOIN items i ON v.id = i.id
          WHERE embedding MATCH ?
          ORDER BY distance
          LIMIT 5
        `).all(new Float32Array(vector));

        return {
          content: [{ type: "text", text: JSON.stringify(results, null, 2) }],
        };
      }
      case "llm_generate": {
        const { prompt, n_predict, temperature } = request.params.arguments;
        const text = await llamaCompletion(prompt, { n_predict, temperature });
        return { content: [{ type: "text", text }] };
      }
      default:
        throw new Error("Unknown tool");
    }
  });
}