/**
 * @file index.ts
 * @description Data-sampler LM Studio plugin entry point.
 *
 * Registers two tools via toolsProvider:
 *
 * Tool: sample
 *   Samples input text into optimised chunks using priority, relevance,
 *   or length-based strategies.
 *   See ./sampler.ts and ./strategies/ for implementation details.
 *
 * Tool: load_dataset
 *   Loads sample rows from HuggingFace datasets via the HF Datasets Server
 *   REST API. Supports config/split discovery, BM25 search, and multiple
 *   output formats (text, json, prompt).
 *   See ./dataset.ts for implementation details.
 *
 * @package data-sampler
 */

import {
  PluginContext,
  tool,
  Tool,
  ToolsProviderController,
} from "@lmstudio/sdk";
import { z } from "zod";

import { loadDataset } from "./dataset";
import { sample } from "./sampler";

// ---------------------------------------------------------------------------
// Plugin entry point
// ---------------------------------------------------------------------------

export async function main(context: PluginContext): Promise<void> {
  context.withToolsProvider(toolsProvider);
}

// ---------------------------------------------------------------------------
// Tools provider
// ---------------------------------------------------------------------------

const SAMPLE_TOOL_DESCRIPTION =
  "Sample input text into optimised chunks using one of three strategies.\n\n" +
  "Strategies:\n" +
  '  - "priority": Scores chunks by auto-extracted keyword frequency and ' +
  "position. Best when you want the most information-dense segments.\n" +
  '  - "relevance": BM25-style scoring against a provided query. ' +
  "Best when you know what you are looking for.\n" +
  '  - "length": Simple fixed-size chunking with optional overlap. ' +
  "Best when you need uniform segments.\n\n" +
  "Use this tool to prepare large texts for downstream processing where " +
  "context window limits matter.";

const LOAD_DATASET_TOOL_DESCRIPTION =
  "Load sample rows from a HuggingFace dataset via the " +
  "HF Datasets Server API.\n\n" +
  "Parameters:\n" +
  "  - dataset (required): HF dataset name, e.g. squad, imdb, " +
  "ibm/duorc.\n" +
  '  - split (optional, default "train"): Which split to use.\n' +
  "  - config (optional): Dataset config/subset. Required when a " +
  "dataset has multiple configs.\n" +
  "  - max_samples (optional, default 5, max 100): Rows to return.\n" +
  "  - query (optional): BM25 full-text search query.\n" +
  '  - format (optional, default "text"): Output format text, json, ' +
  "or prompt.\n" +
  "  - template (optional): Custom template with {text} placeholder.\n\n" +
  "Use this tool when you need sample data from HF datasets for " +
  "few-shot examples, evaluation, or analysis without downloading " +
  "the full dataset.";

async function toolsProvider(ctl: ToolsProviderController): Promise<Tool[]> {
  // -----------------------------------------------------------------------
  // sample tool
  // -----------------------------------------------------------------------

  const sampleTool = tool({
    name: "sample",
    description: SAMPLE_TOOL_DESCRIPTION,
    parameters: {
      text: z
        .string()
        .min(1, "Text must not be empty.")
        .describe("The input text to sample into chunks."),
      strategy: z
        .enum(["priority", "relevance", "length"])
        .optional()
        .default("priority")
        .describe(
          "Sampling strategy. 'priority' (default) scores by keyword " +
            "frequency and position. 'relevance' requires a query. " +
            "'length' produces fixed-size chunks.",
        ),
      chunkSize: z
        .number()
        .int()
        .positive()
        .optional()
        .describe(
          "Target character size per chunk. Default: 2000. Minimum: 1.",
        ),
      overlap: z
        .number()
        .int()
        .min(0)
        .optional()
        .describe(
          "Number of overlapping characters between consecutive chunks. " +
            "Default: 0. Only applies when strategy is 'length'.",
        ),
      maxChunks: z
        .number()
        .int()
        .positive()
        .optional()
        .describe(
          "Maximum number of chunks to return. " +
            "Default: unlimited (all chunks returned).",
        ),
      query: z
        .string()
        .optional()
        .describe(
          "Query string for the 'relevance' strategy (BM25-style scoring). " +
            "Ignored for other strategies.",
        ),
    },

    implementation: async (
      { text, strategy, chunkSize, overlap, maxChunks, query },
      { status, warn, signal },
    ) => {
      if (signal.aborted) {
        return "Sampling cancelled by user.";
      }

      status(
        `Sampling ${text.length.toLocaleString()} chars using "${strategy}" strategy...`,
      );

      try {
        const result = sample(text, {
          strategy,
          chunkSize,
          overlap,
          maxChunks,
          query,
        });

        if (result.chunks.length === 0) {
          return {
            chunks: [],
            totalChunks: 0,
            originalLength: result.originalLength,
            strategy: result.strategy,
            message: "No chunks were produced. The input text may be empty.",
          };
        }

        status(
          `Produced ${result.totalChunks} chunk(s) ` +
            `from ${result.originalLength.toLocaleString()} chars.`,
        );

        return result;
      } catch (err: unknown) {
        const message =
          err instanceof Error ? err.message : String(err ?? "unknown error");
        warn(`Sampling error: ${message}`);
        return {
          chunks: [],
          totalChunks: 0,
          originalLength: text.length ?? 0,
          strategy: strategy ?? "priority",
          error: message,
        };
      }
    },
  });

  // -----------------------------------------------------------------------
  // load_dataset tool
  // -----------------------------------------------------------------------

  const loadDatasetTool = tool({
    name: "load_dataset",
    description: LOAD_DATASET_TOOL_DESCRIPTION,
    parameters: {
      dataset: z
        .string()
        .min(1, "Dataset name must not be empty.")
        .describe(
          'HuggingFace dataset name, e.g. "squad", "imdb", ' +
            '"ibm/duorc", or "bigcode/the-stack-dedup".',
        ),
      split: z
        .enum(["train", "test", "validation"])
        .optional()
        .default("train")
        .describe('Dataset split to load from. Default: "train".'),
      config: z
        .string()
        .optional()
        .describe(
          "Dataset config/subset name. Required when a dataset has " +
            "multiple configs (e.g. 'wikitext-103-raw-v1' for wiki text).",
        ),
      max_samples: z
        .number()
        .int()
        .min(1)
        .max(100)
        .optional()
        .describe("Number of rows to return. Default: 5. Max: 100."),
      query: z
        .string()
        .optional()
        .describe(
          "Optional BM25 full-text search query to filter rows. " +
            "When provided, the server returns the most relevant rows.",
        ),
      format: z
        .enum(["text", "json", "prompt", "fewshot"])
        .optional()
        .default("text")
        .describe(
          'Output format. "text" (default): key:value lines. ' +
            '"json": raw JSON objects. "prompt": === Sample N === format. ' +
            '"fewshot": Input:/Output: pairs ready for prompt injection ' +
            "(use input_columns/output_column to control mapping).",
        ),
      template: z
        .string()
        .optional()
        .describe(
          "Custom template with {text} placeholder for each sample. " +
            'Only used when format is "text".',
        ),
      indices: z
        .string()
        .optional()
        .describe(
          'Comma/space-separated row indices to fetch, e.g. "0,5,10". ' +
            "Takes precedence over max_samples. Only those rows are returned.",
        ),
      input_columns: z
        .array(z.string())
        .optional()
        .describe(
          'Columns to treat as input ("fewshot" format only). ' +
            "Defaults to all columns except output_column.",
        ),
      output_column: z
        .string()
        .optional()
        .describe(
          'Column to treat as output/answer ("fewshot" format only). ' +
            "Defaults to the last column.",
        ),
    },

    implementation: async (params, { status, warn, signal }) => {
      if (signal.aborted) {
        return "Dataset loading cancelled by user.";
      }

      status(
        `Loading up to ${params.max_samples ?? 5} sample(s) ` +
          `from dataset "${params.dataset}"...`,
      );

      try {
        const result = await loadDataset(params, signal);

        status(
          `Loaded ${result.count} sample(s) from ` +
            `${result.dataset}/${result.config} ` +
            `(split: ${result.split}, total rows: ${result.totalRows}).`,
        );

        return result;
      } catch (err: unknown) {
        const message =
          err instanceof Error ? err.message : String(err ?? "unknown error");
        warn(`Dataset loading error: ${message}`);
        return {
          dataset: params.dataset,
          split: params.split ?? "train",
          config: params.config ?? null,
          samples: [],
          totalRows: 0,
          count: 0,
          format: params.format ?? "text",
          formatted: "",
          error: message,
        };
      }
    },
  });

  return [sampleTool, loadDatasetTool];
}