import { readFile } from "node:fs/promises"
import type { ExtensionAPI } from "@oh-my-pi/pi-coding-agent"

// viking_context — authenticated, local-audience access to the federated
// OpenViking service on M1. No direct OpenViking CLI access remains here.

const VIKING = "http://100.117.228.68:8093"
const TOKEN_FILE_ENV = "AHEL_VIKING_TOOL_TOKEN_FILE"
const FIND_CAP = 2500
const READ_CAP = 3500
const RESPONSE_CAP = 64 * 1024
const TIMEOUT_MS = 60000

async function loadToken(): Promise<string> {
  const tokenFile = process.env[TOKEN_FILE_ENV]
  if (!tokenFile) throw new Error("knowledge authentication unavailable")
  const token = (await readFile(tokenFile, "utf8")).trim()
  if (!token) throw new Error("knowledge authentication unavailable")
  return token
}

async function readResponseBody(response: Response): Promise<string> {
  if (!response.body) return ""
  const reader = response.body.getReader()
  const decoder = new TextDecoder()
  let text = ""
  let bytes = 0
  for (;;) {
    const { done, value } = await reader.read()
    if (done) break
    bytes += value.byteLength
    if (bytes > RESPONSE_CAP) {
      try { await reader.cancel() } catch {}
      throw new Error("knowledge response too large")
    }
    text += decoder.decode(value, { stream: true })
  }
  return text + decoder.decode()
}

async function postViking(endpoint: string, body: object, signal?: AbortSignal): Promise<any> {
  const token = await loadToken()
  const controller = new AbortController()
  const timer = setTimeout(() => controller.abort(), TIMEOUT_MS)
  const abort = () => controller.abort()
  signal?.addEventListener("abort", abort, { once: true })
  try {
    const response = await fetch(`${VIKING}${endpoint}`, {
      method: "POST",
      headers: {
        "Authorization": `Bearer ${token}`,
        "Content-Type": "application/json",
      },
      signal: controller.signal,
      body: JSON.stringify(body),
    })
    if (!response.ok) throw new Error(`knowledge service http ${response.status}`)
    const text = await readResponseBody(response)
    try { return JSON.parse(text) } catch { throw new Error("invalid knowledge response") }
  } finally {
    clearTimeout(timer)
    signal?.removeEventListener("abort", abort)
  }
}

function normalizeResult(result: any): object | null {
  if (!result || (result.realm !== "cloud" && result.realm !== "private")) return null
  if (typeof result.uri !== "string" || !result.uri.startsWith("viking://")) return null
  return {
    realm: result.realm,
    uri: result.uri,
    score: Number.isFinite(Number(result.score)) ? Number(result.score) : 0,
    abstract: String(result.abstract || "").slice(0, FIND_CAP),
  }
}

export default function vikingExtension(pi: ExtensionAPI) {
  const { z } = pi.zod

  pi.registerTool({
    name: "viking_context",
    label: "Viking Context",
    description:
      "Ahel knowledge base. Give query to search (returns ranked viking:// URIs with abstracts); give uri to read that node's content. Drill down: search first, then read the best URI.",
    parameters: z.object({
      query: z.string().optional().describe("Natural-language search"),
      uri: z.string().optional().describe("viking:// URI to read"),
      realm: z.enum(["cloud", "private"]).optional().describe("Realm returned by search; required with uri"),
    }),
    async execute(_toolCallId, params, signal) {
      try {
        if (params.uri) {
          if (!params.realm) {
            return { content: [{ type: "text", text: "error: provide realm from search result" }], isError: true }
          }
          const payload = await postViking("/v2/read", {
            audience: "local", realm: params.realm, uri: params.uri,
          }, signal)
          if (!payload || payload.realm !== params.realm || payload.uri !== params.uri) {
            throw new Error("knowledge response rejected")
          }
          return { content: [{ type: "text", text: String(payload.content || "").slice(0, READ_CAP) }] }
        }
        if (params.query) {
          const payload = await postViking("/v2/find", {
            query: params.query, audience: "local", limit: 5,
          }, signal)
          const results = Array.isArray(payload?.results)
            ? payload.results.map(normalizeResult).filter(Boolean)
            : []
          const text = JSON.stringify({ results, partial: payload?.partial === true }).slice(0, FIND_CAP)
          return { content: [{ type: "text", text }] }
        }
      } catch {
        // Never surface credential values, paths, response bodies, or transport
        // diagnostics through the model-visible tool result.
        return { content: [{ type: "text", text: "error: knowledge request failed" }], isError: true }
      }
      return {
        content: [{ type: "text", text: "error: provide query or uri" }],
        isError: true,
      }
    },
  })
}
