import { mkdir, writeFile } from "node:fs/promises";
import { dirname, extname, resolve } from "node:path";
import {
  type ExtensionAPI,
  withFileMutationQueue,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

type GatewayImageGenerationConfig = {
  endpoint: string;
};

type GatewayConfig = {
  baseUrl: string;
  models: Array<{ id: string; name: string }>;
  tools?: {
    image_generation?: GatewayImageGenerationConfig;
  };
};

type GatewayProvider = {
  id: string;
  name: string;
  gateway: string;
  apiKey?: string;
  config: GatewayConfig;
  authenticatedConfig?: GatewayConfig;
};

type ImageGenerationToolInput = {
  prompt: string;
  output_path?: string;
  input_images?: Array<{
    mediaType: string;
    data: string;
  }>;
};

type SavedImage = {
  path: string;
  mediaType: string;
  bytes: number;
};

type GatewayImageGenerationResponse = {
  kind: "radius#image_generation";
  model: string;
  prompt: string;
  images: Array<{
    dataUrl: string;
    mediaType: string;
  }>;
  text?: string;
};

const DEFAULT_GATEWAY = "https://radius.pi.dev";
const DEFAULT_DEV_GATEWAY = "http://localhost:8788";
const GATEWAY_PROVIDER_ID = "radius";
const DEV_GATEWAY_PROVIDER_ID = "radius-dev";
const GATEWAY_API_KEY_ENV = "PI_GATEWAY_API_KEY";
const DEV_GATEWAY_API_KEY_ENV = "PI_DEV_GATEWAY_API_KEY";
const IMAGE_GENERATION_TOOL_NAME = "image_generation";

export default async function (pi: ExtensionAPI) {
  const providers: GatewayProvider[] = [
    {
      id: GATEWAY_PROVIDER_ID,
      name: "Radius",
      gateway: process.env.PI_GATEWAY || DEFAULT_GATEWAY,
      apiKey: process.env[GATEWAY_API_KEY_ENV],
      config: createInitialGatewayConfig(process.env.PI_GATEWAY || DEFAULT_GATEWAY),
    },
  ];

  const devGateway = process.env.PI_DEV_GATEWAY || DEFAULT_DEV_GATEWAY;
  if (devGateway) {
    providers.push({
      id: DEV_GATEWAY_PROVIDER_ID,
      name: "Radius (dev)",
      gateway: devGateway,
      apiKey: process.env[DEV_GATEWAY_API_KEY_ENV],
      config: createInitialGatewayConfig(devGateway),
    });
  }

  pi.registerTool({
    name: IMAGE_GENERATION_TOOL_NAME,
    label: "Image Generation",
    description:
      "Generate or edit an image through Radius using Gemini 2.5 Flash Image via OpenRouter. Generated images are written to disk and the tool returns their paths.",
    promptSnippet: "Generate or edit an image through Radius and save it to disk",
    parameters: Type.Object(
      {
        prompt: Type.String({ description: "Image generation prompt" }),
        output_path: Type.Optional(
          Type.String({
            minLength: 1,
            description:
              "File path for the generated image, relative to the working directory unless absolute. For multiple images, a numeric suffix is added. Defaults to generated-image-<tool-call-id>.<extension>.",
          }),
        ),
        input_images: Type.Optional(
          Type.Array(
            Type.Object(
              {
                mediaType: Type.String(),
                data: Type.String(),
              },
              { additionalProperties: false },
            ),
          ),
        ),
      },
      { additionalProperties: false },
    ),
    async execute(
      toolCallId: string,
      params: ImageGenerationToolInput,
      signal: AbortSignal | undefined,
      _onUpdate: unknown,
      ctx: any,
    ) {
      const provider = selectProvider(providers, ctx);
      const config = getEffectiveConfig(provider).tools?.image_generation;
      if (!config) {
        throw new Error("Radius image generation is not configured");
      }
      const apiKey = await getApiKey(provider, ctx);
      const response = await fetch(config.endpoint, {
        method: "POST",
        headers: {
          authorization: `Bearer ${apiKey}`,
          accept: "application/json",
          "content-type": "application/json",
        },
        body: JSON.stringify({
          prompt: params.prompt,
          input_images: params.input_images,
        }),
        signal,
      });
      if (!response.ok) {
        throw new Error(`${response.status} ${response.statusText}: ${await response.text()}`);
      }
      const payload = (await response.json()) as GatewayImageGenerationResponse;
      const images = await saveGeneratedImages(payload, params.output_path, toolCallId, ctx.cwd);
      const summary = [
        `Generated ${images.length} ${images.length === 1 ? "image" : "images"} with ${payload.model}.`,
        ...images.map((image) => `Saved ${image.mediaType} (${image.bytes} bytes) to ${image.path}`),
      ].join("\n");
      return {
        content: [{ type: "text", text: summary }],
        details: {
          provider: provider.id,
          model: payload.model,
          images,
        },
      };
    },
  });
}

async function saveGeneratedImages(
  payload: GatewayImageGenerationResponse,
  requestedPath: string | undefined,
  toolCallId: string,
  cwd: string,
): Promise<SavedImage[]> {
  if (payload.images.length === 0) {
    throw new Error("Radius returned no generated images");
  }

  const basePath = requestedPath
    ? resolve(cwd, requestedPath.replace(/^@/u, ""))
    : undefined;
  const safeToolCallId = toolCallId.replace(/[^a-zA-Z0-9_-]/gu, "-").slice(0, 80) || Date.now();
  const savedImages: SavedImage[] = [];

  for (const [index, image] of payload.images.entries()) {
    const data = decodeImageDataUrl(image.dataUrl);
    const path = basePath
      ? addImageIndex(basePath, index, payload.images.length)
      : resolve(
          cwd,
          `generated-image-${safeToolCallId}${payload.images.length > 1 ? `-${index + 1}` : ""}${extensionForMediaType(image.mediaType)}`,
        );

    await withFileMutationQueue(path, async () => {
      await mkdir(dirname(path), { recursive: true });
      await writeFile(path, data);
    });
    savedImages.push({ path, mediaType: image.mediaType, bytes: data.length });
  }

  return savedImages;
}

function decodeImageDataUrl(dataUrl: string): Buffer {
  const match = /^data:[^;,]+;base64,(.*)$/su.exec(dataUrl);
  if (!match?.[1]) {
    throw new Error("Radius returned an invalid image data URL");
  }
  return Buffer.from(match[1], "base64");
}

function addImageIndex(path: string, index: number, total: number): string {
  if (total === 1) {
    return path;
  }
  const extension = extname(path);
  const stem = extension ? path.slice(0, -extension.length) : path;
  return `${stem}-${index + 1}${extension}`;
}

function extensionForMediaType(mediaType: string): string {
  switch (mediaType.toLowerCase()) {
    case "image/avif":
      return ".avif";
    case "image/gif":
      return ".gif";
    case "image/jpeg":
    case "image/jpg":
      return ".jpg";
    case "image/png":
      return ".png";
    case "image/svg+xml":
      return ".svg";
    case "image/webp":
      return ".webp";
    default:
      return ".img";
  }
}

function selectProvider(providers: GatewayProvider[], ctx: any): GatewayProvider {
  const provider = providers.find((candidate) => candidate.id === ctx.model?.provider);
  if (!provider) {
    throw new Error("image_generation is only available when the current model is a Radius model");
  }
  return provider;
}

async function getApiKey(provider: GatewayProvider, ctx: any): Promise<string> {
  if (provider.apiKey) {
    return provider.apiKey;
  }
  const result = await ctx.modelRegistry.getApiKeyAndHeaders(ctx.model);
  if (!result.ok || !result.apiKey) {
    throw new Error(result.ok ? "Missing gateway API key" : result.error);
  }
  return result.apiKey;
}

function getEffectiveConfig(provider?: GatewayProvider): GatewayConfig {
  return (
    provider?.authenticatedConfig ?? provider?.config ?? createInitialGatewayConfig(DEFAULT_GATEWAY)
  );
}

function createInitialGatewayConfig(gateway: string): GatewayConfig {
  return {
    baseUrl: `${gateway.replace(/\/+$/u, "")}/v1`,
    models: [],
    tools: {
      image_generation: {
        endpoint: `${gateway.replace(/\/+$/u, "")}/v1/images/generate`,
      },
    },
  };
}