/**
 * Seltz.ai Search Extension for Pi Coding Agent
 *
 * Provides two tools for web-grounded search and answers using the Seltz.ai API:
 * - seltz_search: Search the web for relevant context and documents
 * - seltz_answer: Get direct, natural-language answers with citations (RAG over search results)
 *
 * Priority for API key resolution:
 * 1. `api_key` tool parameter (explicit per-call override)
 * 2. `SELTZ_API_KEY` environment variable
 * 3. `~/.seltz/config.json` → `{ "api_key": "..." }`
 * 4. `<cwd>/.seltz/config.json` → `{ "api_key": "..." }`
 */

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";
import { StringEnum } from "@earendil-works/pi-ai";
import {
  VERSION,
  truncateHead,
  DEFAULT_MAX_BYTES,
  DEFAULT_MAX_LINES,
} from "@earendil-works/pi-coding-agent";
import { existsSync, readFileSync } from "node:fs";
import { join } from "node:path";
import { homedir } from "node:os";

const SELTZ_API_BASE = "https://api.seltz.ai";
const USER_AGENT = `pi-seltz-search/1.0.1 (pi/${VERSION})`;

/**
 * Resolve the Seltz API key with fallback priority:
 * 1. `api_key` tool parameter (explicit per-call override)
 * 2. `SELTZ_API_KEY` environment variable
 * 3. `~/.seltz/config.json` → `{ "api_key": "..." }`
 * 4. `<cwd>/.seltz/config.json` → `{ "api_key": "..." }`
 */
function getApiKey(paramsApiKey: string | undefined, cwd: string): string | undefined {
  if (paramsApiKey) return paramsApiKey;
  if (process.env.SELTZ_API_KEY) return process.env.SELTZ_API_KEY;

  const configPaths = [
    join(homedir(), ".seltz", "config.json"),
    join(cwd, ".seltz", "config.json"),
  ];

  for (const configPath of configPaths) {
    if (existsSync(configPath)) {
      try {
        const raw = readFileSync(configPath, "utf8");
        const config = JSON.parse(raw) as Record<string, unknown>;
        if (typeof config.api_key === "string" && config.api_key.length > 0) {
          return config.api_key;
        }
      } catch {
        // Invalid JSON — skip silently
      }
    }
  }

  return undefined;
}

async function seltzFetch<T>(path: string, apiKey: string, body: Record<string, unknown>): Promise<T> {
  const response = await fetch(`${SELTZ_API_BASE}${path}`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      "x-api-key": apiKey,
      "User-Agent": USER_AGENT,
    },
    body: JSON.stringify(body),
  });

  if (!response.ok) {
    const errorText = await response.text().catch(() => "");
    throw new Error(`Seltz API error ${response.status}: ${errorText || response.statusText}`);
  }

  return (await response.json()) as T;
}

// ─── Search Tool ───────────────────────────────────────────────────────────────

interface SearchDocument {
  url: string | null;
  content: string | null;
  published_date: string | null;
}

interface SearchResponse {
  documents: SearchDocument[];
}

const searchSchema = Type.Object({
  query: Type.String({
    description:
      'The search query. Keep it short for best results (e.g., "TypeScript generic types").',
  }),
  max_results: Type.Optional(
    Type.Integer({
      description: "Maximum number of results to return. Default: 10.",
      minimum: 1,
      maximum: 50,
    }),
  ),
  include_domains: Type.Optional(
    Type.Array(Type.String(), {
      description:
        'Only include results from these domains (e.g., ["typescriptlang.org", "github.com"]).',
    }),
  ),
  exclude_domains: Type.Optional(
    Type.Array(Type.String(), {
      description: "Exclude results from these domains.",
    }),
  ),
  scope: Type.Optional(
    StringEnum(["news"] as const, {
      description:
        'Limit search to a specific vertical. Use "news" for news-only results.',
    }),
  ),
  from_date: Type.Optional(
    Type.String({
      description: 'Only include results on or after this date (ISO 8601, e.g., "2025-01-01").',
    }),
  ),
  to_date: Type.Optional(
    Type.String({
      description: 'Only include results on or before this date (ISO 8601, e.g., "2025-12-31").',
    }),
  ),
  api_key: Type.Optional(
    Type.String({
      description:
        "Seltz API key. If omitted, the SELTZ_API_KEY environment variable or ~/.seltz/config.json is used.",
    }),
  ),
});

export type SeltzSearchInput = typeof searchSchema.static;

// ─── Answer Tool ───────────────────────────────────────────────────────────────

interface HttpCitation {
  url: string;
  content: string | null;
}

interface AnswerHttpResponse {
  answer: string;
  citations: HttpCitation[];
}

const answerSchema = Type.Object({
  query: Type.String({
    description:
      "The natural-language question to answer (e.g., \"Who is reported to be Apple's next CEO?\").",
  }),
  model: Type.Optional(
    Type.String({
      description:
        "Optional tier selector (e.g., 'seltz-base'). Selects the answer tier for behavior and billing.",
    }),
  ),
  include_content: Type.Optional(
    Type.Boolean({
      description:
        "When true, citations include the source document content text. Default: false.",
    }),
  ),
  scope: Type.Optional(
    StringEnum(["news"] as const, {
      description:
        'Limit the grounding search to a specific vertical. Use "news" for news-only.',
    }),
  ),
  api_key: Type.Optional(
    Type.String({
      description:
        "Seltz API key. If omitted, the SELTZ_API_KEY environment variable or ~/.seltz/config.json is used.",
    }),
  ),
});

export type SeltzAnswerInput = typeof answerSchema.static;

// ─── Extension Entry Point ─────────────────────────────────────────────────────

export default function (pi: ExtensionAPI) {
  // Register the seltz_search tool
  pi.registerTool({
    name: "seltz_search",
    label: "Seltz Search",
    description:
      "Search the live web using Seltz AI. Returns structured documents with URLs and content, optimized for LLM reasoning. Use this for current information, API documentation, specifications, pricing, or any data that may have changed since your training cutoff.",
    promptSnippet:
      "Search the web for up-to-date information using Seltz AI",
    promptGuidelines: [
      "Use seltz_search when the user asks about recent events, current documentation, or live data that may not be in your training set.",
      "Prefer seltz_search over generic web searches when precise, context-engineered results with source URLs are needed.",
    ],
    parameters: searchSchema,
    async execute(_toolCallId, params, signal, onUpdate, ctx) {
      if (signal?.aborted) {
        return { content: [{ type: "text", text: "Search cancelled." }], details: {} };
      }

      const apiKey = getApiKey(params.api_key, ctx.cwd);
      if (!apiKey) {
        throw new Error(
          "No Seltz API key provided. Set the SELTZ_API_KEY environment variable, create ~/.seltz/config.json, or pass api_key in the tool parameters. Get a key at https://console.seltz.ai/api-keys",
        );
      }

      const {
        api_key: _apiKeyParam,
        ...requestBody
      } = params;

      const maxResults = requestBody.max_results ?? 10;
      onUpdate?.({
        content: [{ type: "text", text: `Searching Seltz for "${requestBody.query}" (up to ${maxResults} results)...` }],
      });

      const result = await seltzFetch<SearchResponse>("/v1/search", apiKey, {
        ...requestBody,
        max_results: maxResults,
        api_key: null,
      });

      if (!result.documents || result.documents.length === 0) {
        return {
          content: [{
            type: "text",
            text: `No results found for query: "${requestBody.query}"`,
          }],
          details: { query: requestBody.query, resultCount: 0 },
        };
      }

      const formattedDocs = result.documents
        .map((doc, i) => {
          const parts: string[] = [];
          parts.push(`**[${i + 1}] ${doc.url || "Unknown source"}**`);
          if (doc.published_date) {
            parts.push(`*Published: ${doc.published_date}*`);
          }
          if (doc.content) {
            parts.push(doc.content);
          }
          return parts.join("\n\n");
        })
        .join("\n\n---\n\n");

      let outputText = `Found ${result.documents.length} result(s) for "${requestBody.query}":\n\n${formattedDocs}`;

      const truncation = truncateHead(outputText, {
        maxLines: DEFAULT_MAX_LINES,
        maxBytes: DEFAULT_MAX_BYTES,
      });
      outputText = truncation.content;

      if (truncation.truncated) {
        outputText += `\n\n[Output truncated: ${truncation.outputLines} of ${truncation.totalLines} lines shown]`;
      }

      return {
        content: [
          {
            type: "text",
            text: outputText,
          },
        ],
        details: {
          query: requestBody.query,
          resultCount: result.documents.length,
          documents: result.documents.map((d) => ({
            url: d.url,
            published_date: d.published_date,
          })),
        },
      };
    },
  });

  // Register the seltz_answer tool
  pi.registerTool({
    name: "seltz_answer",
    label: "Seltz Answer",
    description:
      "Get a direct, cited answer from the live web using Seltz AI. Use this when you need a concise answer with source citations rather than full document content. Best for factual questions, 'who/what/when/where' queries, and quick lookups.",
    promptSnippet:
      "Get a direct, cited answer to a question using Seltz AI's RAG engine",
    promptGuidelines: [
      "Use seltz_answer when the user asks a specific factual question that needs a direct answer with citations.",
      "Prefer seltz_answer over seltz_search for simple factual queries; use seltz_search when deeper document context is needed.",
    ],
    parameters: answerSchema,
    async execute(_toolCallId, params, signal, onUpdate, ctx) {
      if (signal?.aborted) {
        return { content: [{ type: "text", text: "Answer cancelled." }], details: {} };
      }

      const apiKey = getApiKey(params.api_key, ctx.cwd);
      if (!apiKey) {
        throw new Error(
          "No Seltz API key provided. Set the SELTZ_API_KEY environment variable, create ~/.seltz/config.json, or pass api_key in the tool parameters. Get a key at https://console.seltz.ai/api-keys",
        );
      }

      const {
        api_key: _apiKeyParam,
        ...requestBody
      } = params;

      onUpdate?.({
        content: [{ type: "text", text: `Getting answer for "${requestBody.query}" from Seltz...` }],
      });

      const result = await seltzFetch<AnswerHttpResponse>("/v1/answer", apiKey, {
        ...requestBody,
        stream: false,
        api_key: null,
      });

      const citationList = result.citations
        .map((c, i) => {
          const line = `[${i + 1}] ${c.url}`;
          return c.content ? `${line}\n    ${c.content.substring(0, 200)}` : line;
        })
        .join("\n\n");

      let outputText = [
        result.answer,
        citationList ? `\n\n### Sources\n\n${citationList}` : "",
      ].join("");

      const truncation = truncateHead(outputText, {
        maxLines: DEFAULT_MAX_LINES,
        maxBytes: DEFAULT_MAX_BYTES,
      });
      outputText = truncation.content;

      if (truncation.truncated) {
        outputText += `\n\n[Output truncated: ${truncation.outputLines} of ${truncation.totalLines} lines shown]`;
      }

      return {
        content: [{ type: "text", text: outputText }],
        details: {
          query: requestBody.query,
          citationCount: result.citations.length,
          citations: result.citations.map((c) => ({ url: c.url })),
        },
      };
    },
  });
}