// src/toolsProvider.ts
import { text, tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./config";

const RATE_LIMIT = {
  perSecond: 1,
  perMonth: 15_000,
} as const;

const requestCount = {
  second: 0,
  month: 0,
  lastReset: Date.now(),
};

function checkRateLimit(): void {
  const now = Date.now();
  if (now - requestCount.lastReset > 1_000) {
    requestCount.second = 0;
    requestCount.lastReset = now;
  }
  if (
    requestCount.second >= RATE_LIMIT.perSecond ||
    requestCount.month >= RATE_LIMIT.perMonth
  ) {
    throw new Error("Rate limit exceeded");
  }
  requestCount.second += 1;
  requestCount.month += 1;
}

// ─── Shared fetch helper ────────────────────────────────────────────────────

async function braveFetch<T>(url: URL, apiKey: string): Promise<T> {
  checkRateLimit();

  const response = await fetch(url, {
    headers: {
      Accept: "application/json",
      "Accept-Encoding": "gzip",
      "X-Subscription-Token": apiKey,
    },
  });

  if (!response.ok) {
    const body = await response.text().catch(() => "");
    throw new Error(
      `Brave API error: ${response.status} ${response.statusText}${body ? `\n${body}` : ""}`,
    );
  }

  return response.json() as Promise<T>;
}

// ─── Types ──────────────────────────────────────────────────────────────────

interface WebResult {
  title: string;
  description: string;
  url: string;
}

interface BraveWebSearchResponse {
  web?: { results?: Array<{ title?: string; description?: string; url?: string }> };
}

interface BraveLocationsResponse {
  locations?: { results?: Array<{ id?: string }> };
}

interface Poi {
  id: string;
  name: string;
  address?: {
    streetAddress?: string;
    addressLocality?: string;
    addressRegion?: string;
    postalCode?: string;
  };
  phone?: string;
  rating?: { ratingValue?: number; ratingCount?: number };
  priceRange?: string;
  openingHours?: string[];
}

interface PoisResponse {
  results?: Poi[];
}

interface DescriptionsResponse {
  descriptions: Record<string, string>;
}

// ─── Core search functions ──────────────────────────────────────────────────

async function performWebSearch(
  apiKey: string,
  query: string,
  count = 10,
  offset = 0,
): Promise<string> {
  const url = new URL("https://api.search.brave.com/res/v1/web/search");
  url.searchParams.set("q", query);
  url.searchParams.set("count", String(Math.min(count, 20)));
  url.searchParams.set("offset", String(offset));

  const data = await braveFetch<BraveWebSearchResponse>(url, apiKey);

  const results: WebResult[] = (data.web?.results ?? []).map((r) => ({
    title: r.title ?? "",
    description: r.description ?? "",
    url: r.url ?? "",
  }));

  if (results.length === 0) return "No results found.";

  return results
    .map((r) => `Title: ${r.title}\nDescription: ${r.description}\nURL: ${r.url}`)
    .join("\n\n");
}

async function getPoisData(apiKey: string, ids: string[]): Promise<PoisResponse> {
  const url = new URL("https://api.search.brave.com/res/v1/local/pois");
  for (const id of ids) {
    if (id) url.searchParams.append("ids", id);
  }
  return braveFetch<PoisResponse>(url, apiKey);
}

async function getDescriptionsData(
  apiKey: string,
  ids: string[],
): Promise<DescriptionsResponse> {
  const url = new URL("https://api.search.brave.com/res/v1/local/descriptions");
  for (const id of ids) {
    if (id) url.searchParams.append("ids", id);
  }
  return braveFetch<DescriptionsResponse>(url, apiKey);
}

function formatLocalResults(
  poisData: PoisResponse,
  descData: DescriptionsResponse,
): string {
  const pois = poisData.results ?? [];
  if (pois.length === 0) return "No local results found.";

  return pois
    .map((poi) => {
      const address =
        [
          poi.address?.streetAddress,
          poi.address?.addressLocality,
          poi.address?.addressRegion,
          poi.address?.postalCode,
        ]
          .filter(Boolean)
          .join(", ") || "N/A";

      return [
        `Name: ${poi.name}`,
        `Address: ${address}`,
        `Phone: ${poi.phone ?? "N/A"}`,
        `Rating: ${poi.rating?.ratingValue ?? "N/A"} (${poi.rating?.ratingCount ?? 0} reviews)`,
        `Price Range: ${poi.priceRange ?? "N/A"}`,
        `Hours: ${(poi.openingHours ?? []).join(", ") || "N/A"}`,
        `Description: ${descData.descriptions[poi.id] ?? "No description available"}`,
      ].join("\n");
    })
    .join("\n---\n");
}

async function performLocalSearch(
  apiKey: string,
  query: string,
  count = 5,
): Promise<string> {
  const webUrl = new URL("https://api.search.brave.com/res/v1/web/search");
  webUrl.searchParams.set("q", query);
  webUrl.searchParams.set("search_lang", "en");
  webUrl.searchParams.set("result_filter", "locations");
  webUrl.searchParams.set("count", String(Math.min(count, 20)));

  const webData = await braveFetch<BraveLocationsResponse>(webUrl, apiKey);

  const locationIds =
    webData.locations?.results
      ?.map((r) => r.id)
      .filter((id): id is string => Boolean(id)) ?? [];

  if (locationIds.length === 0) {
    return performWebSearch(apiKey, query, count);
  }

  const [poisData, descriptionsData] = await Promise.all([
    getPoisData(apiKey, locationIds),
    getDescriptionsData(apiKey, locationIds),
  ]);

  return formatLocalResults(poisData, descriptionsData);
}

// ─── Tools Provider ─────────────────────────────────────────────────────────

export async function toolsProvider(
  ctl: ToolsProviderController,
): Promise<Tool[]> {
  const config = ctl.getPluginConfig(configSchematics);
  const apiKey = config.get("braveApiKey")?.trim() ?? "";

  if (!apiKey) {
    // Return no tools so the model doesn't try to call them with a missing key
    return [];
  }

  const webSearchTool = tool({
    name: "brave_web_search",
    description: text`
      Performs a web search using the Brave Search API.
      Ideal for general queries, news, articles, and online content.
      Use for broad information gathering, recent events, or diverse web sources.
      Supports pagination (offset) and result count (max 20).
    `,
    parameters: {
      query: z
        .string()
        .min(1)
        .max(400)
        .describe("Search query (max 400 chars)"),
      count: z
        .number()
        .int()
        .min(1)
        .max(20)
        .optional()
        .describe("Number of results (1-20, default 10)"),
      offset: z
        .number()
        .int()
        .min(0)
        .max(9)
        .optional()
        .describe("Pagination offset (0-9, default 0)"),
    },
    implementation: async ({ query, count = 10, offset = 0 }) => {
      try {
        return await performWebSearch(apiKey, query, count, offset);
      } catch (err) {
        return `Error: ${err instanceof Error ? err.message : String(err)}`;
      }
    },
  });

  const localSearchTool = tool({
    name: "brave_local_search",
    description: text`
      Searches for local businesses and places using Brave Local Search API.
      Best for queries about physical locations, businesses, restaurants, services, etc.
      Returns: name, address, phone, rating, review count, price range, hours, and description.
      Use when the query implies "near me" or mentions a specific location.
      Automatically falls back to web search if no local results are found.
    `,
    parameters: {
      query: z
        .string()
        .min(1)
        .describe("Local search query (e.g. 'pizza near Central Park')"),
      count: z
        .number()
        .int()
        .min(1)
        .max(20)
        .optional()
        .describe("Number of results (1-20, default 5)"),
    },
    implementation: async ({ query, count = 5 }) => {
      try {
        return await performLocalSearch(apiKey, query, count);
      } catch (err) {
        return `Error: ${err instanceof Error ? err.message : String(err)}`;
      }
    },
  });

  return [webSearchTool, localSearchTool];
}

// src/toolsProvider.ts
import { text, tool, type Tool, type ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import { configSchematics } from "./config";

const RATE_LIMIT = {
  perSecond: 1,
  perMonth: 15_000,
} as const;

const requestCount = {
  second: 0,
  month: 0,
  lastReset: Date.now(),
};

function checkRateLimit(): void {
  const now = Date.now();
  if (now - requestCount.lastReset > 1_000) {
    requestCount.second = 0;
    requestCount.lastReset = now;
  }
  if (
    requestCount.second >= RATE_LIMIT.perSecond ||
    requestCount.month >= RATE_LIMIT.perMonth
  ) {
    throw new Error("Rate limit exceeded");
  }
  requestCount.second += 1;
  requestCount.month += 1;
}

// ─── Shared fetch helper ────────────────────────────────────────────────────

async function braveFetch<T>(url: URL, apiKey: string): Promise<T> {
  checkRateLimit();

  const response = await fetch(url, {
    headers: {
      Accept: "application/json",
      "Accept-Encoding": "gzip",
      "X-Subscription-Token": apiKey,
    },
  });

  if (!response.ok) {
    const body = await response.text().catch(() => "");
    throw new Error(
      `Brave API error: ${response.status} ${response.statusText}${body ? `\n${body}` : ""}`,
    );
  }

  return response.json() as Promise<T>;
}

// ─── Types ──────────────────────────────────────────────────────────────────

interface WebResult {
  title: string;
  description: string;
  url: string;
}

interface BraveWebSearchResponse {
  web?: { results?: Array<{ title?: string; description?: string; url?: string }> };
}

interface BraveLocationsResponse {
  locations?: { results?: Array<{ id?: string }> };
}

interface Poi {
  id: string;
  name: string;
  address?: {
    streetAddress?: string;
    addressLocality?: string;
    addressRegion?: string;
    postalCode?: string;
  };
  phone?: string;
  rating?: { ratingValue?: number; ratingCount?: number };
  priceRange?: string;
  openingHours?: string[];
}

interface PoisResponse {
  results?: Poi[];
}

interface DescriptionsResponse {
  descriptions: Record<string, string>;
}

// ─── Core search functions ──────────────────────────────────────────────────

async function performWebSearch(
  apiKey: string,
  query: string,
  count = 10,
  offset = 0,
): Promise<string> {
  const url = new URL("https://api.search.brave.com/res/v1/web/search");
  url.searchParams.set("q", query);
  url.searchParams.set("count", String(Math.min(count, 20)));
  url.searchParams.set("offset", String(offset));

  const data = await braveFetch<BraveWebSearchResponse>(url, apiKey);

  const results: WebResult[] = (data.web?.results ?? []).map((r) => ({
    title: r.title ?? "",
    description: r.description ?? "",
    url: r.url ?? "",
  }));

  if (results.length === 0) return "No results found.";

  return results
    .map((r) => `Title: ${r.title}\nDescription: ${r.description}\nURL: ${r.url}`)
    .join("\n\n");
}

async function getPoisData(apiKey: string, ids: string[]): Promise<PoisResponse> {
  const url = new URL("https://api.search.brave.com/res/v1/local/pois");
  for (const id of ids) {
    if (id) url.searchParams.append("ids", id);
  }
  return braveFetch<PoisResponse>(url, apiKey);
}

async function getDescriptionsData(
  apiKey: string,
  ids: string[],
): Promise<DescriptionsResponse> {
  const url = new URL("https://api.search.brave.com/res/v1/local/descriptions");
  for (const id of ids) {
    if (id) url.searchParams.append("ids", id);
  }
  return braveFetch<DescriptionsResponse>(url, apiKey);
}

function formatLocalResults(
  poisData: PoisResponse,
  descData: DescriptionsResponse,
): string {
  const pois = poisData.results ?? [];
  if (pois.length === 0) return "No local results found.";

  return pois
    .map((poi) => {
      const address =
        [
          poi.address?.streetAddress,
          poi.address?.addressLocality,
          poi.address?.addressRegion,
          poi.address?.postalCode,
        ]
          .filter(Boolean)
          .join(", ") || "N/A";

      return [
        `Name: ${poi.name}`,
        `Address: ${address}`,
        `Phone: ${poi.phone ?? "N/A"}`,
        `Rating: ${poi.rating?.ratingValue ?? "N/A"} (${poi.rating?.ratingCount ?? 0} reviews)`,
        `Price Range: ${poi.priceRange ?? "N/A"}`,
        `Hours: ${(poi.openingHours ?? []).join(", ") || "N/A"}`,
        `Description: ${descData.descriptions[poi.id] ?? "No description available"}`,
      ].join("\n");
    })
    .join("\n---\n");
}

async function performLocalSearch(
  apiKey: string,
  query: string,
  count = 5,
): Promise<string> {
  const webUrl = new URL("https://api.search.brave.com/res/v1/web/search");
  webUrl.searchParams.set("q", query);
  webUrl.searchParams.set("search_lang", "en");
  webUrl.searchParams.set("result_filter", "locations");
  webUrl.searchParams.set("count", String(Math.min(count, 20)));

  const webData = await braveFetch<BraveLocationsResponse>(webUrl, apiKey);

  const locationIds =
    webData.locations?.results
      ?.map((r) => r.id)
      .filter((id): id is string => Boolean(id)) ?? [];

  if (locationIds.length === 0) {
    return performWebSearch(apiKey, query, count);
  }

  const [poisData, descriptionsData] = await Promise.all([
    getPoisData(apiKey, locationIds),
    getDescriptionsData(apiKey, locationIds),
  ]);

  return formatLocalResults(poisData, descriptionsData);
}

// ─── Tools Provider ─────────────────────────────────────────────────────────

export async function toolsProvider(
  ctl: ToolsProviderController,
): Promise<Tool[]> {
  const config = ctl.getPluginConfig(configSchematics);
  const apiKey = config.get("braveApiKey")?.trim() ?? "";

  if (!apiKey) {
    // Return no tools so the model doesn't try to call them with a missing key
    return [];
  }

  const webSearchTool = tool({
    name: "brave_web_search",
    description: text`
      Performs a web search using the Brave Search API.
      Ideal for general queries, news, articles, and online content.
      Use for broad information gathering, recent events, or diverse web sources.
      Supports pagination (offset) and result count (max 20).
    `,
    parameters: {
      query: z
        .string()
        .min(1)
        .max(400)
        .describe("Search query (max 400 chars)"),
      count: z
        .number()
        .int()
        .min(1)
        .max(20)
        .optional()
        .describe("Number of results (1-20, default 10)"),
      offset: z
        .number()
        .int()
        .min(0)
        .max(9)
        .optional()
        .describe("Pagination offset (0-9, default 0)"),
    },
    implementation: async ({ query, count = 10, offset = 0 }) => {
      try {
        return await performWebSearch(apiKey, query, count, offset);
      } catch (err) {
        return `Error: ${err instanceof Error ? err.message : String(err)}`;
      }
    },
  });

  const localSearchTool = tool({
    name: "brave_local_search",
    description: text`
      Searches for local businesses and places using Brave Local Search API.
      Best for queries about physical locations, businesses, restaurants, services, etc.
      Returns: name, address, phone, rating, review count, price range, hours, and description.
      Use when the query implies "near me" or mentions a specific location.
      Automatically falls back to web search if no local results are found.
    `,
    parameters: {
      query: z
        .string()
        .min(1)
        .describe("Local search query (e.g. 'pizza near Central Park')"),
      count: z
        .number()
        .int()
        .min(1)
        .max(20)
        .optional()
        .describe("Number of results (1-20, default 5)"),
    },
    implementation: async ({ query, count = 5 }) => {
      try {
        return await performLocalSearch(apiKey, query, count);
      } catch (err) {
        return `Error: ${err instanceof Error ? err.message : String(err)}`;
      }
    },
  });

  return [webSearchTool, localSearchTool];
}