/**
 * XPI — exploit_search tool for Pi Agent.
 *
 * Searches the preview.is security corpus (api.preview.is) for exploit
 * techniques, attack primitives, and bypasses. Results are grounded in
 * real write-ups with source URLs — the agent should cite them and never
 * invent exploit details from memory.
 *
 * Authentication: reads the API key from ~/.pi/preview-key.json
 * or the `PREVIEW_IS_API_KEY` environment variable (header: `X-API-Key`).
 */

import { createHash } from "node:crypto";
import { existsSync, readFileSync, statSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { abortableSleep, TtlLruCache } from "@xaccefy/pi-shared";
import { Type } from "typebox";

const PREVIEW_IS_API = "https://api.preview.is/search";

const ExploitSearchSchema = Type.Object(
  {
    query: Type.String({
      description: "Search query for exploit techniques, CVEs, or attack primitives",
    }),
    limit: Type.Optional(
      Type.Number({ description: "Max results to return (default 5, maps to API `k`)" }),
    ),
    minScore: Type.Optional(
      Type.Number({
        description: "Minimum relevance score 0–1 (default 0.1). Higher = stricter matching.",
      }),
    ),
  },
  { additionalProperties: false },
);

// ── API response types ────────────────────────────────────────────────

type MatchedSection = {
  heading?: string;
  score?: number;
  text?: string;
};

type PreviewIsResult = {
  rank?: number;
  score?: number;
  title?: string;
  url?: string;
  file?: string;
  matched_sections?: MatchedSection[];
  content?: string | null;
};

type PreviewIsResponse = {
  query?: string;
  count?: number;
  results?: PreviewIsResult[];
};

// ── Normalised result surfaced to the agent ───────────────────────────

type ExploitResult = {
  title: string;
  url: string;
  score: number;
  technique: string;
  context: string;
};

/** Read the API key from ~/.pi/preview-key.json or the environment, returning a clear error if missing. */
function getApiKey(): string {
  const envKey = process.env.PREVIEW_IS_API_KEY?.trim();
  if (envKey) return envKey;

  const jsonPath = join(homedir(), ".pi", "preview-key.json");
  if (existsSync(jsonPath)) {
    // Permission gate BEFORE the try: the mode error must reach the caller,
    // not fall through to the generic "no API key" message.
    if (process.platform !== "win32") {
      const mode = statSync(jsonPath).mode & 0o077;
      if (mode !== 0) {
        throw new Error(
          `~/.pi/preview-key.json is readable by other users (mode ${(mode & 0o777).toString(8)}) — ` +
            "an API key must not be group/world-readable. Fix with: chmod 600 ~/.pi/preview-key.json",
        );
      }
    }
    try {
      const parsed = JSON.parse(readFileSync(jsonPath, "utf-8"));
      if (parsed && typeof parsed === "object") {
        const fileKey = (parsed.apiKey ?? "").trim();
        if (fileKey) return fileKey;
      }
    } catch {
      // fall through to the error below
    }
  }

  throw new Error(
    'No preview.is API key found. Set it in ~/.pi/preview-key.json ({ "apiKey": "rk_..." }) or export:\n' +
      '  export PREVIEW_IS_API_KEY="rk_..."',
  );
}

/** Build a compact technique string from the best matched section heading. */
function extractMatch(sections: MatchedSection[] | undefined): {
  technique: string;
  context: string;
} {
  const best = sections?.length
    ? [...sections].sort((a, b) => (b.score ?? 0) - (a.score ?? 0))[0]
    : undefined;
  const text = best?.text?.trim() || "";
  return {
    technique: best?.heading?.trim() || "",
    context: text.length > 400 ? `${text.slice(0, 400)}…` : text,
  };
}

/**
 * Query preview.is with a timeout, one retry on transient failures (429/5xx), and a
 * bounded singleflight TTL cache so repeated/parallel identical queries don't burn quota.
 */
const REQUEST_TIMEOUT_MS = 30_000;
const CACHE_TTL_MS = 10 * 60 * 1000;
const CACHE_MAX = 64;
const cache = new TtlLruCache<PreviewIsResponse>(CACHE_TTL_MS, CACHE_MAX);

async function queryPreviewIs(
  apiKey: string,
  query: string,
  limit: number,
  minScore: number,
  parentSignal?: AbortSignal,
): Promise<PreviewIsResponse> {
  // Include the API key (hashed) so a key rotation mid-session cannot serve
  // results fetched under the old key's quota/permissions.
  const keyHash = createHash("sha256").update(apiKey).digest("hex").slice(0, 12);
  const cacheKey = `${keyHash}|${query}|${limit}|${minScore}`;

  return cache.getOrLoad(cacheKey, async () => {
    const doFetch = (): Promise<Response> =>
      fetch(PREVIEW_IS_API, {
        method: "POST",
        headers: {
          "Content-Type": "application/json",
          Accept: "application/json",
          "X-API-Key": apiKey,
        },
        body: JSON.stringify({ query, k: limit, min_score: minScore }),
        signal: AbortSignal.any([
          ...(parentSignal ? [parentSignal] : []),
          AbortSignal.timeout(REQUEST_TIMEOUT_MS),
        ]),
      });

    let lastErr: Error | undefined;
    for (let attempt = 0; attempt < 2; attempt++) {
      const response = await doFetch();
      if (response.status === 401 || response.status === 403) {
        throw new Error(
          `preview.is rejected the API key (HTTP ${response.status}). Check PREVIEW_IS_API_KEY is valid.`,
        );
      }
      if (response.status === 429) {
        lastErr = new Error("preview.is rate limit reached. Wait for the weekly/monthly reset.");
        try {
          response.body?.cancel();
        } catch {
          // ignore
        }
        await abortableSleep(500 * (attempt + 1), parentSignal);
        continue;
      }
      if (!response.ok) {
        if (response.status >= 500 && attempt === 0) {
          lastErr = new Error(`Exploit search failed: HTTP ${response.status}`);
          try {
            response.body?.cancel();
          } catch {
            // ignore
          }
          await abortableSleep(500, parentSignal);
          continue;
        }
        throw new Error(`Exploit search failed: HTTP ${response.status}`);
      }
      return (await response.json()) as PreviewIsResponse;
    }
    throw lastErr ?? new Error("Exploit search failed after retries");
  });
}

export default function exploitSearchExtension(pi: ExtensionAPI) {
  pi.registerTool({
    name: "exploit_search",
    label: "Search Exploit Techniques",
    description:
      "Search the preview.is security corpus for exploit techniques, attack primitives, and bypasses. " +
      "Results are grounded in real write-ups with source URLs. " +
      'Focus on the technique demonstrated, not just the POC. Requires PREVIEW_IS_API_KEY env var or ~/.pi/preview-key.json ("apiKey" field).',
    promptSnippet: "Search for exploit techniques",
    promptGuidelines: [
      "Cite the source URL for every technique you surface from exploit_search.",
      "Do not invent exploit details from memory — ground answers in the returned write-ups.",
    ],
    parameters: ExploitSearchSchema,

    async execute(_id, params, signal, _onUpdate, _ctx) {
      const query = (params.query as string).trim();
      const limit = Math.max(1, Math.min((params.limit as number | undefined) ?? 5, 50));
      const minScore = Math.max(0, Math.min((params.minScore as number | undefined) ?? 0.1, 1));

      try {
        if (!query) {
          throw new Error("query must be a non-empty string");
        }
        const apiKey = getApiKey();
        const data = await queryPreviewIs(apiKey, query, limit, minScore, signal);

        const seen = new Set<string>();
        const results: ExploitResult[] = (data.results || [])
          .filter((r) => {
            const url = r.url || "";
            if (!url) return true;
            if (seen.has(url)) return false;
            seen.add(url);
            return true;
          })
          .map((r) => ({
            title: r.title || "Untitled",
            url: r.url || "",
            score: r.score ?? 0,
            ...extractMatch(r.matched_sections),
          }));

        if (results.length === 0) {
          return {
            content: [{ type: "text" as const, text: `No results for "${query}".` }],
            details: { results: [], query, count: 0 },
          };
        }

        const lines = [`Results for "${query}" (${results.length} hits, min_score ${minScore}):`];
        for (const r of results) {
          lines.push("");
          lines.push(`→ ${r.title}  [score: ${r.score.toFixed(3)}]`);
          if (r.technique) lines.push(`  Technique: ${r.technique}`);
          if (r.context) lines.push(`  Context: ${r.context}`);
          if (r.url) lines.push(`  URL: ${r.url}`);
        }

        return {
          content: [{ type: "text" as const, text: lines.join("\n") }],
          details: { results, query, count: results.length },
        };
      } catch (error) {
        throw new Error(`Exploit search error: ${(error as Error).message}`, { cause: error });
      }
    },

    renderCall(args, theme) {
      return new Text(
        theme.fg("toolTitle", theme.bold("exploit_search ")) +
          theme.fg("dim", (args.query as string) ?? ""),
        0,
        0,
      );
    },

    renderResult(result, _options, theme, context) {
      if (context.isError) {
        return new Text(theme.fg("error", "✗ Exploit search failed"), 0, 0);
      }
      const details = result.details as { results?: ExploitResult[] } | undefined;
      const count = details?.results?.length ?? 0;
      return new Text(theme.fg("success", `✓ ${count} results`), 0, 0);
    },
  });
}