// Agnes AI provider (OpenAI-compatible) — https://agnes-ai.com
// Docs: https://agnes-ai.com/zh-Hans/docs/overview
// Auth: AGNES_API_KEY env var (user-level Windows env var)
//
// Model discovery: registers a fast seed list on startup, then refreshes
// from /v1/models in the background.  Discovered models are persisted to disk
// and re-used on subsequent starts when the API is unreachable.

import { mkdir, writeFile } from "node:fs/promises";
import { readFileSync } from "node:fs";
import { join } from "node:path";
import {
  createAssistantMessageEventStream,
  openAICompletionsApi,
} from "@earendil-works/pi-ai";
import { Image, Markdown } from "@earendil-works/pi-tui";

let appendAgnesImage = null;

// ---------------------------------------------------------------------------
// Model helpers
// ---------------------------------------------------------------------------

function detectLimits(id) {
  if (id.startsWith("agnes-2.5")) return { contextWindow: 1048576, maxTokens: 65536 };
  if (id.startsWith("agnes-2.0")) return { contextWindow: 1048576, maxTokens: 32768 };
  // Sensible defaults for unknown models
  return { contextWindow: 131072, maxTokens: 32768 };
}

const IMAGE_MODELS = new Set(["agnes-image-2.0-flash", "agnes-image-2.1-flash"]);
const VIDEO_MODELS = new Set(["agnes-video-v2.0", "agnes-video-2.5", "agnes-video-2.5-flash"]);

// Prefix checks let newly released Agnes image/video versions route correctly
// as soon as they appear in /v1/models, without requiring a plugin update.
function isImageModel(id) {
  return IMAGE_MODELS.has(id) || id.startsWith("agnes-image-");
}

function isVideoModel(id) {
  return VIDEO_MODELS.has(id) || id.startsWith("agnes-video-");
}

function convertModel(model) {
  const id = model.id;
  const imageModel = isImageModel(id);
  return {
    id,
    name: model.id || id,
    reasoning: false,
    input: imageModel ? ["text", "image"] : ["text", "image"],
    cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
    ...detectLimits(id),
    agnesImageModel: imageModel,
    agnesVideoModel: VIDEO_MODELS.has(id),
  };
}

// ---------------------------------------------------------------------------
// Seed models (available immediately on startup)
// ---------------------------------------------------------------------------

const AGNES_SEED = [
  "agnes-2.5-flash",
  "agnes-2.5-pro",
  "agnes-2.5-pro-alpha",
  "agnes-2.0-flash",
  "agnes-video-v2.0",
  "agnes-video-2.5",
  "agnes-video-2.5-flash",
];

// ---------------------------------------------------------------------------
// Dynamic model fetch (shared by startup & refreshModels)
// ---------------------------------------------------------------------------

async function fetchModels(baseUrl, apiKey, signal) {
  const headers = {};
  if (apiKey) headers["Authorization"] = `Bearer ${apiKey}`;

  const res = await fetch(`${baseUrl}/models`, { headers, redirect: "follow", signal });
  if (!res.ok) throw new Error(`HTTP ${res.status} ${res.statusText}`);

  const payload = await res.json();
  // OpenAI /v1/models returns { data: [{ id, ... }] }
  const data = Array.isArray(payload?.data) ? payload.data : Array.isArray(payload) ? payload : [];
  return data
    .filter((m) => m && m.id)
    .map(convertModel);
}

// ---------------------------------------------------------------------------
// Agnes image generation API
// ---------------------------------------------------------------------------

function latestUser(context) {
  const messages = Array.isArray(context?.messages) ? context.messages : [];
  return [...messages].reverse().find((message) => message?.role === "user");
}

function imageRequestContent(context) {
  const user = latestUser(context);
  if (!user) return { prompt: "", images: [] };
  if (typeof user.content === "string") return { prompt: user.content, images: [] };
  return {
    prompt: user.content.filter((part) => part?.type === "text").map((part) => part.text ?? "").join("\n"),
    images: user.content.filter((part) => part?.type === "image"),
  };
}

async function saveImage(image, modelId) {
  const directory = join(process.cwd(), ".pi", "generated-images");
  await mkdir(directory, { recursive: true });
  const mime = image?.mime_type ?? "image/png";
  const extension = mime.includes("jpeg") ? "jpg" : mime.includes("webp") ? "webp" : "png";
  const filePath = join(directory, `${modelId}-${Date.now()}.${extension}`);
  if (image?.b64_json) await writeFile(filePath, Buffer.from(image.b64_json, "base64"));
  else if (image?.url) {
    const response = await fetch(image.url);
    if (!response.ok) throw new Error(`Unable to download image: HTTP ${response.status}`);
    await writeFile(filePath, Buffer.from(await response.arrayBuffer()));
  } else throw new Error("Agnes image API returned no url or b64_json");
  return { filePath, mimeType: mime };
}

function providerKey(model, options) {
  return options?.apiKey ?? process.env[model.provider === "agnes-cn" ? "AGNES_CN_API_KEY" : "AGNES_API_KEY"] ?? "";
}

function streamAgnesImage(model, context, options) {
  const stream = createAssistantMessageEventStream();
  const output = {
    role: "assistant", content: [], api: model.api, provider: model.provider, model: model.id,
    usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0,
      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
    stopReason: "pending", timestamp: Date.now(),
  };
  (async () => {
    try {
      stream.push({ type: "start", partial: output });
      const { prompt, images } = imageRequestContent(context);
      if (!prompt) throw new Error("Image generation requires a text prompt");
      const body = {
        model: model.id,
        prompt,
        ...(images.length ? { image: images.map((image) => `data:${image.mimeType};base64,${image.data}`) } : {}),
        extra_body: { response_format: "url" },
      };
      const baseUrl = model.baseUrl ?? (model.provider === "agnes-cn" ? "https://api.agnes-ai.cn/v1" : "https://apihub.agnes-ai.com/v1");
      const response = await fetch(`${baseUrl}/images/generations`, {
        method: "POST",
        headers: { Authorization: `Bearer ${providerKey(model, options)}`, "Content-Type": "application/json" },
        body: JSON.stringify(body), signal: options?.signal,
      });
      const payload = await response.json();
      if (!response.ok) throw new Error(payload?.error?.message ?? `Agnes image API HTTP ${response.status}`);
      const image = payload?.data?.[0];
      if (!image) throw new Error("Agnes image API returned no image data");
      const saved = await saveImage(image, model.id);
      const filePath = saved.filePath;
      appendAgnesImage?.({ path: filePath, mimeType: saved.mimeType });
      const text = image.url
        ? `![Generated image](${image.url})\n\nSaved local copy: ${filePath}\n\nImage URL may expire according to Agnes retention policy.`
        : `Generated image saved to: ${filePath}`;
      output.content.push({ type: "text", text });
      stream.push({ type: "text_start", contentIndex: 0, partial: output });
      stream.push({ type: "text_delta", contentIndex: 0, delta: text, partial: output });
      stream.push({ type: "text_end", contentIndex: 0, content: text, partial: output });
      output.stopReason = "stop";
      stream.push({ type: "done", reason: "stop", message: output });
      stream.end();
    } catch (error) {
      output.stopReason = options?.signal?.aborted ? "aborted" : "error";
      output.errorMessage = error instanceof Error ? error.message : String(error);
      stream.push({ type: "error", reason: output.stopReason, error: output });
      stream.end();
    }
  })();
  return stream;
}

function streamAgnes(model, context, options) {
  if (model.agnesVideoModel || isVideoModel(model.id)) return streamAgnesVideo(model, context, options);
  if (model.agnesImageModel || isImageModel(model.id)) return streamAgnesImage(model, context, options);
  return openAICompletionsApi().streamSimple(model, context, options);
}

async function waitForVideo(baseUrl, videoId, apiKey, signal) {
  const deadline = Date.now() + 30 * 60 * 1000;
  while (Date.now() < deadline) {
    if (signal?.aborted) throw new Error("Video generation aborted");
    await new Promise((resolve, reject) => {
      const timer = setTimeout(resolve, 5000);
      signal?.addEventListener("abort", () => { clearTimeout(timer); reject(new Error("Video generation aborted")); }, { once: true });
    });
    const apiRoot = baseUrl.replace(/\/v1\/?$/, "");
    const response = await fetch(`${apiRoot}/agnesapi?video_id=${encodeURIComponent(videoId)}`, {
      headers: { Authorization: `Bearer ${apiKey}` }, signal,
    });
    const payload = await response.json();
    if (!response.ok) throw new Error(payload?.error?.message ?? `Agnes video status HTTP ${response.status}`);
    if (payload.status === "completed") return payload;
    if (payload.status === "failed") throw new Error(payload?.error?.message ?? "Agnes video generation failed");
  }
  throw new Error("Agnes video generation timed out after 30 minutes");
}

async function saveVideo(url, modelId) {
  const directory = join(process.cwd(), ".pi", "generated-videos");
  await mkdir(directory, { recursive: true });
  const filePath = join(directory, `${modelId}-${Date.now()}.mp4`);
  const response = await fetch(url);
  if (!response.ok) throw new Error(`Unable to download generated video: HTTP ${response.status}`);
  await writeFile(filePath, Buffer.from(await response.arrayBuffer()));
  return filePath;
}

function streamAgnesVideo(model, context, options) {
  const stream = createAssistantMessageEventStream();
  const output = {
    role: "assistant", content: [], api: model.api, provider: model.provider, model: model.id,
    usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0,
      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
    stopReason: "pending", timestamp: Date.now(),
  };
  (async () => {
    try {
      stream.push({ type: "start", partial: output });
      const { prompt, images } = imageRequestContent(context);
      if (!prompt) throw new Error("Video generation requires a text prompt");
      const baseUrl = model.baseUrl ?? (model.provider === "agnes-cn" ? "https://api.agnes-ai.cn/v1" : "https://apihub.agnes-ai.com/v1");
      const apiKey = providerKey(model, options);
      const body = {
        model: model.id,
        prompt,
        ...(images.length === 1 ? { image: `data:${images[0].mimeType};base64,${images[0].data}` } : {}),
        ...(images.length > 1 ? { extra_body: { image: images.map((image) => `data:${image.mimeType};base64,${image.data}`), mode: "keyframes" } } : {}),
        num_frames: 121,
        frame_rate: 24,
      };
      const response = await fetch(`${baseUrl}/videos`, {
        method: "POST", headers: { Authorization: `Bearer ${apiKey}`, "Content-Type": "application/json" },
        body: JSON.stringify(body), signal: options?.signal,
      });
      const task = await response.json();
      if (!response.ok) throw new Error(task?.error?.message ?? `Agnes video API HTTP ${response.status}`);
      const videoId = task.video_id ?? task.id ?? task.task_id;
      if (!videoId) throw new Error("Agnes video API returned no video_id");
      const result = task.status === "completed" ? task : await waitForVideo(baseUrl, videoId, apiKey, options?.signal);
      const url = result?.metadata?.url;
      if (!url) throw new Error("Agnes video API returned no metadata.url");
      const filePath = await saveVideo(url, model.id);
      const text = `Generated video saved to: ${filePath}\n\nVideo URL: ${url}`;
      output.content.push({ type: "text", text });
      stream.push({ type: "text_start", contentIndex: 0, partial: output });
      stream.push({ type: "text_delta", contentIndex: 0, delta: text, partial: output });
      stream.push({ type: "text_end", contentIndex: 0, content: text, partial: output });
      output.stopReason = "stop";
      stream.push({ type: "done", reason: "stop", message: output });
      stream.end();
    } catch (error) {
      output.stopReason = options?.signal?.aborted ? "aborted" : "error";
      output.errorMessage = error instanceof Error ? error.message : String(error);
      stream.push({ type: "error", reason: output.stopReason, error: output });
      stream.end();
    }
  })();
  return stream;
}

// ---------------------------------------------------------------------------
// Extension entry point (synchronous — no network on startup)
// ---------------------------------------------------------------------------

export default function (pi) {
  appendAgnesImage = (image) => pi.appendEntry("agnes-generated-image", image);
  pi.registerEntryRenderer("agnes-generated-image", (entry, _options, theme) => {
    const image = entry.data ?? {};
    try {
      const data = readFileSync(image.path).toString("base64");
      return new Image(data, image.mimeType || "image/png", theme, { maxWidthCells: 80, maxHeightCells: 30 });
    } catch {
      return new Markdown(`Generated image unavailable: ${image.path ?? "unknown path"}`, 1, 0, theme);
    }
  });

  const providers = [
    { id: "agnes", name: "Agnes AI", baseUrl: "https://apihub.agnes-ai.com/v1", apiKeyEnv: "AGNES_API_KEY" },
    { id: "agnes-cn", name: "Agnes AI (CN)", baseUrl: "https://api.agnes-ai.cn/v1", apiKeyEnv: "AGNES_CN_API_KEY" },
  ];

  for (const p of providers) {
    const baseUrl = p.baseUrl;
    const apiKeyEnv = p.apiKeyEnv;
    // Let /login provide the key when the environment variable is absent.
    // A literal placeholder would make pi consider the provider configured,
    // while still sending an invalid key during model refresh.
    let currentModels = AGNES_SEED.map((id) => convertModel({ id }));
    const provider = {
      id: p.id,
      name: p.name,
      baseUrl,
      auth: {
        apiKey: {
          name: `${p.name} API Key`,
          async login(interaction) {
            const key = await interaction.prompt({ type: "secret", message: `${p.name} API Key` });
            if (!key.trim()) throw new Error(`${p.name} API Key cannot be empty`);
            return { type: "api_key", key: key.trim() };
          },
          async resolve({ credential, ctx }) {
            const key = credential?.key ?? await ctx.env(apiKeyEnv);
            return key ? { auth: { apiKey: key }, source: credential?.key ? "stored API key" : apiKeyEnv } : undefined;
          },
        },
      },
      getModels: () => currentModels,
      get models() { return currentModels; },
      streamSimple: streamAgnes,
      async refreshModels({ signal, stored, publish, allowNetwork, credential }) {
        const cachedModels = Array.isArray(stored?.models) ? stored.models : undefined;
        if (cachedModels?.length) currentModels = cachedModels;
        if (!allowNetwork || signal.aborted) return currentModels;
        const apiKey = credential?.key ?? process.env[apiKeyEnv];
        try {
          const models = await fetchModels(baseUrl, apiKey, signal);
          if (models.length > 0) {
            currentModels = models;
            await publish({ persist: { provider: p.id, models } });
          }
        } catch (error) {
          if (!currentModels.length) throw error;
        }
        return currentModels;
      },
    };
    pi.registerProvider(provider.id, {
      name: provider.name,
      baseUrl: provider.baseUrl,
      api: "openai-completions",
      apiKey: `$${p.apiKeyEnv}`,
      streamSimple: provider.streamSimple,
      models: provider.getModels(),
      refreshModels: provider.refreshModels,
    });
  }
}