/**
 * Prompt Translate Extension for Pi Coding Agent
 *
 * When enabled, user text is tagged as `[Translating your Prompt...] :` plus the
 * prompt on the next line; translation runs asynchronously, then we update the same bubble via
 * Pi TUI `Markdown.setText`. The `context` hook still translates for the LLM if this
 * finishes late. Interactive-only: needs a captured root `TUI` from `setWidget`.
 *
 * Features:
 * - Toggle on/off via /translate-toggle
 * - Set target language via /translate-lang <language>
 * - Check status via /translate-status
 * - Lightweight: uses a single-turn LLM call with minimal token usage
 */

import { ModelSelectorComponent, SettingsManager, type ExtensionContext, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { completeSimple, type Model, type SimpleStreamOptions } from "@mariozechner/pi-ai";
import { Container, Markdown, type TUI } from "@mariozechner/pi-tui";
import { existsSync, readFileSync, writeFileSync } from "node:fs";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";

// ------------------------------------------------------------------------------
// State
// ------------------------------------------------------------------------------

interface ExtensionState {
  enabled: boolean;
  targetLang: string;
  translateModel?: string;
}

const DEFAULT_STATE: ExtensionState = {
  enabled: false,
  targetLang: "English",
};

const STATE_PATH = join(dirname(fileURLToPath(import.meta.url)), "state.json");

function loadState(): ExtensionState {
  if (!existsSync(STATE_PATH)) {
    return { ...DEFAULT_STATE };
  }
  try {
    const raw = readFileSync(STATE_PATH, "utf-8");
    const parsed = JSON.parse(raw) as Partial<ExtensionState>;
    return {
      enabled: typeof parsed.enabled === "boolean" ? parsed.enabled : DEFAULT_STATE.enabled,
      targetLang: typeof parsed.targetLang === "string" && parsed.targetLang.trim().length > 0
        ? parsed.targetLang.trim()
        : DEFAULT_STATE.targetLang,
      translateModel: typeof parsed.translateModel === "string" && parsed.translateModel.trim().length > 0
        ? parsed.translateModel.trim()
        : undefined,
    };
  } catch {
    return { ...DEFAULT_STATE };
  }
}

function saveState(state: ExtensionState): void {
  try {
    writeFileSync(STATE_PATH, JSON.stringify(state, null, 2) + "\n", "utf-8");
  } catch {
    // Silently fail if state cannot be persisted
  }
}

// ------------------------------------------------------------------------------
// Translation
// ------------------------------------------------------------------------------

/**
 * First line marker; body follows (optional newlines). Space before `:` so CommonMark
 * does not treat `[…]:` as a link reference definition (that swallows the next lines → empty bubble).
 */
const TRANSLATING_HEADER = "[Translating your Prompt...] :";

function wrapTranslatingPrompt(text: string): string {
  return `${TRANSLATING_HEADER}\n${text}`;
}

function unwrapTranslatingPrompt(text: string): string | undefined {
  if (!text.startsWith(TRANSLATING_HEADER)) {
    return undefined;
  }
  return text.slice(TRANSLATING_HEADER.length).replace(/^\n+/, "");
}

function getUserMessagePlainText(message: {
  role: string;
  content: unknown;
}): string {
  if (message.role !== "user") {
    return "";
  }
  const content = message.content;
  if (typeof content === "string") {
    return content;
  }
  if (!Array.isArray(content)) {
    return "";
  }
  const parts: string[] = [];
  for (const block of content) {
    if (typeof block === "object" && block !== null && "type" in block && block.type === "text" && "text" in block) {
      parts.push(String((block as { text: string }).text));
    }
  }
  return parts.join("\n");
}

function setUserMessagePlainText(
  message: { role: string; content: unknown },
  text: string,
): void {
  if (message.role !== "user") {
    return;
  }
  const content = message.content;
  if (typeof content === "string") {
    (message as { content: string }).content = text;
    return;
  }
  if (!Array.isArray(content)) {
    return;
  }
  const next: typeof content = [];
  let mergedIntoFirstText = false;
  for (const block of content) {
    if (
      typeof block === "object" &&
      block !== null &&
      "type" in block &&
      block.type === "text"
    ) {
      if (!mergedIntoFirstText) {
        next.push({ ...block, text });
        mergedIntoFirstText = true;
      }
      continue;
    }
    next.push(block);
  }
  if (!mergedIntoFirstText) {
    next.push({ type: "text", text });
  }
  (message as { content: typeof next }).content = next;
}

async function translateText(
  text: string,
  targetLang: string,
  model: NonNullable<Parameters<typeof completeSimple>[0]>,
  streamOptions?: SimpleStreamOptions,
): Promise<string> {
  const context = {
    systemPrompt:
      `Translate the following user message to ${targetLang}. ` +
      `Preserve all code blocks, file paths, variable names, technical terms, and concepts exactly. ` +
      `Only translate natural language portions. ` +
      `Do NOT add explanations, thinking, or commentary. ` +
      `Output ONLY the translated text.`,
    messages: [
      {
        role: "user" as const,
        content: text,
        timestamp: Date.now(),
      },
    ],
  };

  const response = await completeSimple(model, context, {
    reasoning: "low",
    ...streamOptions,
  });

  if (response.errorMessage) {
    throw new Error(response.errorMessage);
  }

  let translated = "";
  for (const item of response.content) {
    if (item.type === "text") {
      translated += item.text;
    }
  }

  const trimmed = translated.trim();
  return trimmed.length > 0 ? trimmed : text;
}

/** Dedupe concurrent translates (`message_start` + `context` may overlap). */
function translationInflightKey(targetLang: string, original: string): string {
  return `${targetLang}\u0000${original}`;
}

function translateOnce(
  original: string,
  targetLang: string,
  model: NonNullable<Parameters<typeof completeSimple>[0]>,
  streamOptions: SimpleStreamOptions | undefined,
  inflight: Map<string, Promise<string>>,
): Promise<string> {
  const key = translationInflightKey(targetLang, original);
  let pending = inflight.get(key);
  if (!pending) {
    pending = translateText(original, targetLang, model, streamOptions).finally(() => {
      inflight.delete(key);
    });
    inflight.set(key, pending);
  }
  return pending;
}

function getTranslateModel(
  state: ExtensionState,
  ctx: Pick<ExtensionContext, "modelRegistry" | "model">,
): Model | undefined {
  if (state.translateModel) {
    const [provider, ...idParts] = state.translateModel.split("/");
    const modelId = idParts.join("/");
    if (provider && modelId) {
      const found = ctx.modelRegistry.find(provider, modelId);
      if (found) return found;
    }
  }
  return ctx.model;
}

async function resolveStreamOptions(
  model: Model | undefined,
  ctx: Pick<ExtensionContext, "modelRegistry" | "signal">,
): Promise<
  | { ok: true; model: Model; streamOptions: SimpleStreamOptions }
  | { ok: false; reason: string }
> {
  if (!model) {
    return { ok: false, reason: "no_model" };
  }
  const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
  if (!auth.ok) {
    return { ok: false, reason: auth.error };
  }
  const streamOptions: SimpleStreamOptions = {
    reasoning: "low",
  };
  if (auth.apiKey !== undefined) {
    streamOptions.apiKey = auth.apiKey;
  }
  if (auth.headers) {
    streamOptions.headers = auth.headers;
  }
  const sig = ctx.signal;
  if (sig) {
    streamOptions.signal = sig;
  }
  return { ok: true, model, streamOptions };
}

function findLastUserMessage(
  messages: ReadonlyArray<{ role: string; content: unknown }>,
): { role: string; content: unknown } | undefined {
  for (let i = messages.length - 1; i >= 0; i--) {
    const m = messages[i];
    if (m.role === "user") {
      return m;
    }
  }
  return undefined;
}

/** Translate in-place when content is our `[Translating your Prompt...] :` wrapper (live message object). */
async function translateWrappedUserMessageIfNeeded(
  msg: { role: string; content: unknown },
  ctx: Pick<ExtensionContext, "modelRegistry" | "model" | "signal" | "hasUI" | "ui">,
  state: ExtensionState,
  translateInflight: Map<string, Promise<string>>,
  uiErrors: boolean,
): Promise<void> {
  if (msg.role !== "user") {
    return;
  }
  const plain = getUserMessagePlainText(msg);
  const original = unwrapTranslatingPrompt(plain);
  if (original === undefined) {
    return;
  }

  const model = getTranslateModel(state, ctx);
  const resolved = await resolveStreamOptions(model, ctx);
  if (!resolved.ok) {
    if (uiErrors && ctx.hasUI) {
      if (resolved.reason === "no_model") {
        ctx.ui.notify("Prompt translate: no model selected, using original text", "warning");
      }
      else {
        ctx.ui.notify(
          `Translation skipped: ${resolved.reason}. Using original text.`,
          "warning",
        );
      }
    }
    setUserMessagePlainText(msg, original);
    return;
  }

  try {
    const translated = await translateOnce(
      original,
      state.targetLang,
      resolved.model,
      resolved.streamOptions,
      translateInflight,
    );
    setUserMessagePlainText(msg, translated);
  }
  catch (err) {
    const message = err instanceof Error ? err.message : String(err);
    const aborted =
      err instanceof Error &&
      (err.name === "AbortError" || message.toLowerCase().includes("abort"));
    if (aborted) {
      setUserMessagePlainText(msg, original);
      return;
    }
    const hint =
      message.includes("No API key")
        ? " Add an API key in Pi settings or via an environment variable."
        : "";
    if (uiErrors && ctx.hasUI) {
      ctx.ui.notify(`Translation failed: ${message}${hint}`, "error");
    }
    setUserMessagePlainText(msg, original);
  }
}

// ------------------------------------------------------------------------------
// TUI: same user bubble text after translate (Markdown#setText)
// ------------------------------------------------------------------------------

const TUI_CAPTURE_WIDGET_KEY = "__pi_prompt_translate_tui_capture__";

let piRootTui: TUI | undefined;

function capturePiRootTui(ctx: Pick<ExtensionContext, "hasUI" | "ui">): void {
  if (piRootTui || !ctx.hasUI) {
    return;
  }
  try {
    ctx.ui.setWidget(TUI_CAPTURE_WIDGET_KEY, (tui) => {
      piRootTui = tui;
      return new Container();
    });
    ctx.ui.setWidget(TUI_CAPTURE_WIDGET_KEY, undefined);
  }
  catch {
    piRootTui = undefined;
  }
}

function collectMarkdownInstances(node: unknown, out: Markdown[]): void {
  if (!node || typeof node !== "object") {
    return;
  }
  if (node instanceof Markdown) {
    out.push(node);
    return;
  }
  const children = (node as { children?: unknown[] }).children;
  if (!Array.isArray(children)) {
    return;
  }
  for (const ch of children) {
    collectMarkdownInstances(ch, out);
  }
}

/** Match the bubble that still shows `previousPlain`, then swap to `newPlain`. */
function patchUserBubbleMarkdown(root: unknown, previousPlain: string, newPlain: string): void {
  if (!root || previousPlain === newPlain) {
    return;
  }
  const blocks: Markdown[] = [];
  collectMarkdownInstances(root, blocks);
  const hits = blocks.filter((md) => md.text === previousPlain);
  const target = hits[hits.length - 1];
  target?.setText(newPlain);
}

// ------------------------------------------------------------------------------
// Extension Factory
// ------------------------------------------------------------------------------

export default function (pi: ExtensionAPI) {
  const state = loadState();
  const translateInflight = new Map<string, Promise<string>>();

  pi.on("session_shutdown", () => {
    piRootTui = undefined;
  });

  pi.on("message_start", (event, ctx) => {
    if (!state.enabled || event.message.role !== "user") {
      return;
    }
    const bubblePlainAtStart = getUserMessagePlainText(event.message);
    if (unwrapTranslatingPrompt(bubblePlainAtStart) === undefined) {
      return;
    }
    capturePiRootTui(ctx);
    void (async () => {
      await translateWrappedUserMessageIfNeeded(event.message, ctx, state, translateInflight, true);
      const plainAfter = getUserMessagePlainText(event.message);
      patchUserBubbleMarkdown(piRootTui, bubblePlainAtStart, plainAfter);
      piRootTui?.requestRender();
    })();
  });

  pi.on("context", async (event, ctx) => {
    if (!state.enabled) {
      return;
    }
    const messages = event.messages;
    const msg = findLastUserMessage(messages);
    if (!msg) {
      return;
    }
    const plain = getUserMessagePlainText(msg);
    if (unwrapTranslatingPrompt(plain) === undefined) {
      return;
    }
    await translateWrappedUserMessageIfNeeded(msg, ctx, state, translateInflight, true);
    return { messages };
  });

  // Notify on session start
  pi.on("session_start", async (_event, ctx) => {
    if (state.enabled) {
      capturePiRootTui(ctx);
      ctx.ui.notify(
        `Prompt translate: ON -> ${state.targetLang}`,
        "info",
      );
    }
  });

  // Tag input; bubble shows wrapper first; async translate + Markdown#setText; `context` = LLM fallback.
  pi.on("input", (event, _ctx) => {
    if (event.source === "extension") {
      return { action: "continue" };
    }

    if (!state.enabled) {
      return { action: "continue" };
    }

    return {
      action: "transform",
      text: wrapTranslatingPrompt(event.text),
      images: event.images,
    };
  });

  // Commands
  pi.registerCommand("translate-toggle", {
    description: "Toggle prompt translation on/off",
    handler: async (_args, ctx) => {
      state.enabled = !state.enabled;
      saveState(state);
      ctx.ui.notify(
        `Prompt translate: ${state.enabled ? "ON" : "OFF"}`,
        "info",
      );
    },
  });

  pi.registerCommand("translate-lang", {
    description: "Set target translation language (e.g. English, French)",
    handler: async (args, ctx) => {
      const lang = args.trim();
      if (!lang) {
        ctx.ui.notify(
          `Target: ${state.targetLang}. Usage: /translate-lang <language>`,
          "warning",
        );
        return;
      }
      state.targetLang = lang;
      saveState(state);
      ctx.ui.notify(`Target language: ${lang}`, "info");
    },
  });

  pi.registerCommand("translate-status", {
    description: "Show prompt translation status",
    handler: async (_args, ctx) => {
      const model = getTranslateModel(state, ctx);
      const modelLabel = model
        ? state.translateModel
          ? `${model.name} (${model.provider}/${model.id})`
          : `${model.name} (session model)`
        : "none";
      ctx.ui.notify(
        `Translate: ${state.enabled ? "ON" : "OFF"} | Target: ${state.targetLang} | Model: ${modelLabel}`,
        "info",
      );
    },
  });

  pi.registerCommand("translate-model", {
    description: "Select a model for translation, or type 'reset' to use the session model",
    handler: async (args, ctx) => {
      const arg = args.trim().toLowerCase();
      if (arg === "reset" || arg === "default") {
        state.translateModel = undefined;
        saveState(state);
        ctx.ui.notify("Translate model reset to session model", "info");
        return;
      }

      if (!ctx.hasUI) {
        ctx.ui.notify("Model selector requires interactive mode", "warning");
        return;
      }

      const model = await ctx.ui.custom<Model | undefined>((tui, _theme, _kb, done) => {
        const settingsManager = SettingsManager.inMemory();
        const currentModel = getTranslateModel(state, ctx);
        const selector = new ModelSelectorComponent(
          tui,
          currentModel,
          settingsManager,
          ctx.modelRegistry,
          [],
          (selected) => done(selected),
          () => done(undefined),
        );
        return selector;
      }, { overlay: true });

      if (!model) {
        return;
      }

      state.translateModel = `${model.provider}/${model.id}`;
      saveState(state);
      ctx.ui.notify(`Translate model: ${model.name} (${model.provider}/${model.id})`, "info");
    },
  });
}