/**
 * Fast Jev Compaction Extension for Pi Coding Agent
 *
 * Uses Jev System One judgments to prune stale tool calls and tool results
 * during session compaction without rewriting user or assistant messages.
 */

import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { readFileSync } from "node:fs";
import { homedir } from "node:os";
import { compactMessages } from "./src/messages.js";
import { reductionRatio } from "./src/compact.js";
import type { Message, ToolResult, ToolUse } from "./src/types.js";

/** Extract text content from a Pi content string or ContentBlock array */
function extractText(content: unknown): string {
  if (typeof content === "string") return content;
  if (Array.isArray(content)) {
    return content
      .map((block) => {
        if (typeof block === "string") return block;
        if (block && typeof block === "object") {
          if ("text" in block && typeof block.text === "string") return block.text;
          if ("thinking" in block && typeof block.thinking === "string") return block.thinking;
        }
        return "";
      })
      .filter(Boolean)
      .join("\n");
  }
  return "";
}

/** Convert Pi AgentMessages to fast-jev-compaction Message array */
export function convertPiToJevMessages(messages: AgentMessage[]): Message[] {
  const result: Message[] = [];

  for (const msg of messages) {
    if (msg.role === "user") {
      result.push({
        role: "user",
        text: extractText((msg as any).content),
        toolUses: [],
      });
    } else if (msg.role === "assistant") {
      const assistantMsg = msg as any;
      const content = Array.isArray(assistantMsg.content) ? assistantMsg.content : [];

      const textParts: string[] = [];
      const toolUses: ToolUse[] = [];

      for (const block of content) {
        if (!block || typeof block !== "object") continue;
        if (block.type === "text" && typeof block.text === "string") {
          textParts.push(block.text);
        } else if (block.type === "thinking" && typeof block.thinking === "string") {
          textParts.push(`[thinking: ${block.thinking.slice(0, 200)}...]`);
        } else if (block.type === "toolCall") {
          toolUses.push({
            tool_use_id: block.id || `tool_${Math.random().toString(36).slice(2, 9)}`,
            tool: block.name || "unknown",
            input: block.arguments || {},
          });
        }
      }

      result.push({
        role: "assistant",
        text: textParts.join("\n"),
        toolUses,
      });
    } else if (msg.role === "toolResult") {
      const toolMsg = msg as any;
      const tool_use_id = toolMsg.toolCallId || "unknown";
      const text = extractText(toolMsg.content);
      const isError = Boolean(toolMsg.isError);

      result.push({
        role: "user",
        text: "",
        toolUses: [],
        toolResults: [
          {
            tool_use_id,
            text,
            isError,
          },
        ],
      });
    } else if (msg.role === "bashExecution") {
      const bashMsg = msg as any;
      const tool_use_id = `bash_${Math.random().toString(36).slice(2, 9)}`;
      const command = bashMsg.command || "";
      const output = bashMsg.output || "";
      const isError = (bashMsg.exitCode ?? 0) !== 0;

      // Associate as assistant call + user result
      result.push({
        role: "assistant",
        text: "",
        toolUses: [
          {
            tool_use_id,
            tool: "bash",
            input: { command },
          },
        ],
      });
      result.push({
        role: "user",
        text: "",
        toolUses: [],
        toolResults: [
          {
            tool_use_id,
            text: output,
            isError,
          },
        ],
      });
    } else if (msg.role === "branchSummary" || msg.role === "compactionSummary") {
      const sumMsg = msg as any;
      result.push({
        role: "user",
        text: `[Previous Summary]: ${sumMsg.summary || ""}`,
        toolUses: [],
      });
    }
  }

  return result;
}

/** Serialize the Jev-pruned message list into an exact, verbatim transcript.
 *  Dropped calls/results vanish; kept content stays exact (no LLM summarization). */
export function serializePrunedTranscript(messages: readonly Message[]): string {
  const lines: string[] = [];
  for (const m of messages) {
    if (m.role === "user") {
      const text = m.text.trim();
      const results = (m.toolResults ?? []).map((r) =>
        `  [tool_result${r.isError ? " ERROR" : ""}] ${r.text}`
      );
      if (text) lines.push(`USER: ${text}`);
      if (results.length) lines.push(...results);
    } else if (m.role === "assistant") {
      const text = m.text.trim();
      if (text) lines.push(`ASSISTANT: ${text}`);
      for (const t of m.toolUses) {
        const input = JSON.stringify(t.input ?? {});
        const clipped = input.length > 300 ? `${input.slice(0, 300)}…` : input;
        lines.push(`  [tool_call ${t.tool}] ${clipped}`);
      }
    }
  }
  return lines.join("\n");
}

/** Format decisions + stats as a compact header (audit trail only, content is the transcript). */
function formatJevHeader(result: any): string {
  const { decisions, stats } = result;
  const kept = decisions.filter((d: any) => d.action === "keep").length;
  const truncated = decisions.filter((d: any) => d.action === "drop_result").length;
  const dropped = decisions.filter((d: any) => d.action === "drop_call").length;
  return (
    `## Context pruned by fast-jev-compaction (Jev System One)\n\n` +
    `Dropped ${dropped} stale tool call(s), truncated ${truncated} result(s), kept ${kept} intact. ` +
    `Messages ${stats.messagesBefore}\u2192${stats.messagesAfter}, chars ${stats.charsBefore}\u2192${stats.charsAfter} ` +
    `(${(reductionRatio(result) * 100).toFixed(1)}% reduction). User/assistant text is verbatim; nothing was summarized by an LLM.`
  );
}

export default function fastJevCompactionExtension(pi: ExtensionAPI) {
  /** Resolve the TypeSafe/Jev API key from env, falling back to the secrets file. */
  function resolveJevKey(): string | undefined {
    const fromEnv = process.env.JEV_API_KEY || process.env.TYPESAFE_API_KEY;
    if (fromEnv) return fromEnv;
    // Fallback: read the NixOS secrets file directly (works without a re-login).
    try {
      const path = `${homedir()}/.config/environment.d/10-secrets.conf`;
      const content = readFileSync(path, "utf8");
      const m = content.match(/^\s*(?:export\s+)?(?:JEV_API_KEY|TYPESAFE_API_KEY)\s*=\s*"?([^"\n]+)"?/m);
      if (m?.[1]) return m[1].trim();
    } catch {
      /* ignore */
    }
    return undefined;
  }

  pi.on("session_before_compact", async (event, ctx) => {
    const apiKey = resolveJevKey();

    if (!apiKey) {
      ctx.ui.notify("Fast-Jev: JEV_API_KEY not set, using default compaction", "warning");
      return;
    }

    const { preparation, signal } = event;
    const { messagesToSummarize, turnPrefixMessages, tokensBefore, firstKeptEntryId } = preparation;

    const allPiMessages = [...messagesToSummarize, ...(turnPrefixMessages || [])];

    if (allPiMessages.length === 0) {
      return;
    }

    ctx.ui.notify(
      `Fast-Jev Compaction: evaluating ${allPiMessages.length} messages with Jev System One...`,
      "info"
    );

    try {
      const jevMessages = convertPiToJevMessages(allPiMessages);

      const result = await compactMessages(jevMessages, {
        apiKey,
        preserveRecentMessages: 4,
        keepThreshold: 0.5,
      });

      const savedRatio = reductionRatio(result);

      if (savedRatio < 0.15) {
        ctx.ui.notify(
          `Fast-Jev: low reduction (${(savedRatio * 100).toFixed(1)}%), using default compaction`,
          "info"
        );
        return;
      }

      // The summary holds the EXACT pruned transcript (kept content verbatim,
      // dropped tool calls/results removed) — not an LLM summary.
      const transcript = serializePrunedTranscript(result.messages);
      const summary = `${formatJevHeader(result)}\n\n---\n\n${transcript}`;

      ctx.ui.notify(
        `Fast-Jev Compaction complete: reduced ${(savedRatio * 100).toFixed(1)}% of tool content`,
        "info"
      );

      return {
        compaction: {
          summary,
          firstKeptEntryId,
          tokensBefore,
        },
      };
    } catch (error: any) {
      const message = error instanceof Error ? error.message : String(error);
      ctx.ui.notify(`Fast-Jev Compaction failed: ${message}, falling back to default`, "warning");
      return;
    }
  });
}