import type {
  AgentMessage,
  ExtensionAPI,
  ExtensionContext,
  InputEvent,
  MessageEndEvent,
} from "@earendil-works/pi-coding-agent";
import type { AssistantMessage, TextContent, UserMessage } from "@earendil-works/pi-ai";
import { writeFile } from "node:fs/promises";
import { createTranslator, safeTranslate } from "./backends.ts";
import { loadConfig, saveConfig, type BackendName, type Config, type OutputMode } from "./config.ts";
import { openOriginalPanel } from "./panel.ts";
import type { PiTranslateMeta } from "./types.ts";

const STATUS_PREFIX = "⇄";

function updateStatus(cfg: Config, ctx: ExtensionContext): void {
  if (!cfg.showFooterStatus || !ctx.hasUI) {
    ctx.ui.setStatus("translate", undefined);
    return;
  }
  const label = cfg.enabled
    ? `${STATUS_PREFIX} ${cfg.sourceLang}`
    : `${STATUS_PREFIX} off`;
  ctx.ui.setStatus("translate", ctx.ui.theme.fg("accent", label));
}

function isUserMessage(msg: AgentMessage): msg is UserMessage {
  return msg.role === "user";
}

function isAssistantMessage(msg: AgentMessage): msg is AssistantMessage {
  return msg.role === "assistant";
}

function extractText(content: string | (TextContent | { type: "image" })[]): string {
  if (typeof content === "string") return content;
  return content
    .filter((c): c is TextContent => c.type === "text")
    .map((c) => c.text)
    .join("\n");
}

async function translateAssistantContent(
  content: AssistantMessage["content"],
  translator: ReturnType<typeof createTranslator>,
  modelLang: string,
  targetLang: string,
  protectCode: boolean,
  signal: AbortSignal | undefined,
): Promise<AssistantMessage["content"]> {
  const results: AssistantMessage["content"] = [];
  for (const c of content) {
    if (c.type === "text") {
      const result = await safeTranslate(
        translator,
        c.text,
        modelLang,
        targetLang,
        { protectCode, signal },
      );
      results.push({ type: "text", text: result.text });
    } else {
      results.push(c);
    }
  }
  return results;
}

function originalToText(original: AssistantMessage["content"]): string {
  return original
    .map((c) => (c.type === "text" ? c.text : ""))
    .filter(Boolean)
    .join("\n\n");
}

function setTranslatedContent(
  msg: UserMessage,
  modelText: string,
): void {
  if (typeof msg.content === "string") {
    (msg as UserMessage & { piTranslate: PiTranslateMeta }).piTranslate.model = modelText;
    return;
  }

  // Array content: preserve non-text blocks (e.g. images), replace the first
  // text block with the full translation and drop subsequent text blocks.
  const translatedContent: (TextContent | { type: "image" })[] = [];
  let textReplaced = false;
  for (const c of msg.content) {
    if (c.type === "text") {
      if (!textReplaced) {
        translatedContent.push({ type: "text", text: modelText });
        textReplaced = true;
      }
    } else {
      translatedContent.push(c);
    }
  }
  if (!textReplaced) {
    translatedContent.push({ type: "text", text: modelText });
  }
  (msg as UserMessage & { piTranslate: PiTranslateMeta }).piTranslate.model = translatedContent;
}

export default function (pi: ExtensionAPI): void {
  let cfg: Config = loadConfig();
  const pending = new Map<string, string>();
  const contextCache = new Map<string, string>();
  let lastErrorNotified = false;
  let lastAssistantOriginal: (AssistantMessage & { piTranslate: PiTranslateMeta }) | undefined;

  pi.on("session_start", async (_event, ctx) => {
    cfg = loadConfig();
    pending.clear();
    contextCache.clear();
    lastErrorNotified = false;
    lastAssistantOriginal = undefined;
    updateStatus(cfg, ctx);
  });

  pi.registerCommand("translate", {
    description: "Toggle translation on/off",
    handler: async (_args, ctx) => {
      cfg.enabled = !cfg.enabled;
      saveConfig(cfg);
      updateStatus(cfg, ctx);
      ctx.ui.notify(
        `Translation ${cfg.enabled ? "enabled" : "disabled"}`,
        "info",
      );
    },
  });

  pi.registerCommand("translate-backend", {
    description: "Set translation backend",
    handler: async (args, ctx) => {
      const name = args.trim() as BackendName;
      const valid: BackendName[] = ["google", "mymemory", "libretranslate", "llm"];
      if (!valid.includes(name)) {
        ctx.ui.notify(`Usage: /translate-backend <${valid.join("|")}>`, "warning");
        return;
      }
      cfg.backend = name;
      contextCache.clear();
      saveConfig(cfg);
      ctx.ui.notify(`Translation backend set to ${name}`, "info");
    },
  });

  pi.registerCommand("translate-lang", {
    description: "Set source language code",
    handler: async (args, ctx) => {
      const code = args.trim().toLowerCase();
      if (!code || code.length > 5) {
        ctx.ui.notify("Usage: /translate-lang <code> (e.g. it, es, fr)", "warning");
        return;
      }
      cfg.sourceLang = code;
      contextCache.clear();
      saveConfig(cfg);
      updateStatus(cfg, ctx);
      ctx.ui.notify(`Source language set to ${code}`, "info");
    },
  });

  pi.registerCommand("translate-model-lang", {
    description: "Set the language the model reasons in",
    handler: async (args, ctx) => {
      const code = args.trim().toLowerCase();
      if (!code || code.length > 5) {
        ctx.ui.notify("Usage: /translate-model-lang <code> (e.g. en, es, fr)", "warning");
        return;
      }
      cfg.modelLang = code;
      contextCache.clear();
      saveConfig(cfg);
      ctx.ui.notify(`Model language set to ${code}`, "info");
    },
  });

  pi.registerCommand("translate-protect", {
    description: "Toggle code/path protection during translation",
    handler: async (_args, ctx) => {
      cfg.protectCode = !cfg.protectCode;
      saveConfig(cfg);
      ctx.ui.notify(`Code protection ${cfg.protectCode ? "enabled" : "disabled"}`, "info");
    },
  });

  pi.registerCommand("translate-mode", {
    description: "Toggle output mode (translate or native)",
    handler: async (_args, ctx) => {
      cfg.outputMode = cfg.outputMode === "translate" ? "native" : "translate";
      saveConfig(cfg);
      ctx.ui.notify(`Output mode: ${cfg.outputMode}`, "info");
    },
  });

  pi.registerCommand("translate-status", {
    description: "Show current translation settings",
    handler: async (_args, ctx) => {
      const lines = [
        `enabled: ${cfg.enabled}`,
        `backend: ${cfg.backend}`,
        `sourceLang: ${cfg.sourceLang}`,
        `modelLang: ${cfg.modelLang}`,
        `outputMode: ${cfg.outputMode}`,
        `protectCode: ${cfg.protectCode}`,
        `shortcut: ${cfg.originalShortcut}`,
      ];
      ctx.ui.notify(lines.join(" | "), "info");
    },
  });

  pi.registerCommand("translate-original", {
    description: "Show the original model response in a panel",
    handler: async (_args, ctx) => {
      if (!lastAssistantOriginal) {
        ctx.ui.notify("No original model response available yet", "warning");
        return;
      }
      openOriginalPanel(ctx, lastAssistantOriginal);
    },
  });

  pi.registerCommand("translate-mirror", {
    description: "Write the original model response to a file",
    handler: async (args, ctx) => {
      const path = args.trim();
      if (!path) {
        ctx.ui.notify("Usage: /translate-mirror <file-path>", "warning");
        return;
      }
      if (!lastAssistantOriginal) {
        ctx.ui.notify("No original model response available yet", "warning");
        return;
      }
      const text = originalToText(lastAssistantOriginal.piTranslate.original ?? []);
      try {
        await writeFile(path, text, "utf8");
        ctx.ui.notify(`Original model response written to ${path}`, "info");
      } catch (err) {
        const message = err instanceof Error ? err.message : String(err);
        ctx.ui.notify(`Failed to write mirror: ${message}`, "error");
      }
    },
  });

  pi.registerShortcut("ctrl+shift+e", async (ctx) => {
    if (!lastAssistantOriginal) {
      ctx.ui.notify("No original model response available yet", "warning");
      return;
    }
    openOriginalPanel(ctx, lastAssistantOriginal);
  });

  pi.on("input", async (event: InputEvent, ctx) => {
    if (!cfg.enabled) return { action: "continue" };
    if (event.source === "extension") return { action: "continue" };

    const text = event.text.trim();
    if (!text) return { action: "continue" };

    // Skip slash commands/templates and bash prefixes
    if (text.startsWith("/") || text.startsWith("!")) {
      return { action: "continue" };
    }

    const translator = createTranslator(cfg, ctx.modelRegistry);
    const result = await safeTranslate(
      translator,
      text,
      cfg.sourceLang,
      cfg.modelLang,
      { protectCode: cfg.protectCode, signal: ctx.signal },
    );

    if (result.error) {
      if (!lastErrorNotified) {
        lastErrorNotified = true;
        ctx.ui.notify(`Translation failed: ${result.error}`, "warning");
      }
    } else {
      lastErrorNotified = false;
    }

    pending.set(text, result.text);
    return { action: "continue" };
  });

  pi.on("message_end", async (event: MessageEndEvent, ctx) => {
    const msg = event.message;
    if (!cfg.enabled) return;

    if (isUserMessage(msg)) {
      const originalText = extractText(msg.content);
      const modelText = pending.get(originalText) ?? contextCache.get(originalText);
      if (!modelText) return;

      const meta: PiTranslateMeta = {
        sourceLang: cfg.sourceLang,
        targetLang: cfg.modelLang,
        backend: cfg.backend,
      };
      (msg as UserMessage & { piTranslate: PiTranslateMeta }).piTranslate = meta;
      setTranslatedContent(msg, modelText);
      pending.delete(originalText);
      return { message: msg };
    }

    if (isAssistantMessage(msg)) {
      if (msg.stopReason === "error" || cfg.outputMode === "native") {
        return;
      }

      const hasText = msg.content.some((c) => c.type === "text" && c.text.trim());
      if (!hasText) {
        return;
      }

      const translator = createTranslator(cfg, ctx.modelRegistry);
      const originalContent = msg.content;

      ctx.ui.setWorkingMessage("Translating response…");
      try {
        const translatedContent = await translateAssistantContent(
          originalContent,
          translator,
          cfg.modelLang,
          cfg.sourceLang,
          cfg.protectCode,
          ctx.signal,
        );

        const meta: PiTranslateMeta = {
          sourceLang: cfg.modelLang,
          targetLang: cfg.sourceLang,
          backend: cfg.backend,
          original: originalContent,
        };
        (msg as AssistantMessage & { piTranslate: PiTranslateMeta }).piTranslate = meta;
        msg.content = translatedContent;
        lastAssistantOriginal = msg as AssistantMessage & { piTranslate: PiTranslateMeta };
        return { message: msg };
      } finally {
        ctx.ui.setWorkingMessage();
      }
    }

    return;
  });

  pi.on("context", async (event, ctx) => {
    if (!cfg.enabled) return;

    const messages = event.messages;
    const translator = createTranslator(cfg, ctx.modelRegistry);

    for (const msg of messages) {
      if (isAssistantMessage(msg)) {
        const original = (msg as AssistantMessage & { piTranslate?: PiTranslateMeta }).piTranslate?.original;
        if (original !== undefined) {
          msg.content = original;
        }
        continue;
      }

      if (!isUserMessage(msg)) continue;

      const translated = (msg as UserMessage & { piTranslate?: PiTranslateMeta }).piTranslate?.model;
      if (translated !== undefined) {
        msg.content = translated;
        continue;
      }

      const originalText = extractText(msg.content);
      const cached = contextCache.get(originalText);
      if (cached) {
        msg.content = cached;
        continue;
      }

      const result = await safeTranslate(
        translator,
        originalText,
        cfg.sourceLang,
        cfg.modelLang,
        { protectCode: cfg.protectCode },
      );
      contextCache.set(originalText, result.text);
      if (typeof msg.content === "string") {
        msg.content = result.text;
      } else {
        const fallbackContent: (TextContent | { type: "image" })[] = [
          { type: "text", text: result.text },
        ];
        for (const c of msg.content) {
          if (c.type !== "text") fallbackContent.push(c);
        }
        msg.content = fallbackContent;
      }
    }

    return { messages };
  });
}