import { type NextRequest, NextResponse } from "next/server";

import { RunpodError } from "@/lib/server/runpod/client";
import {
  rpImg2Img,
  rpInpaint,
  rpRemoveBackground,
  rpUpscale,
} from "@/lib/server/runpod/endpoints";

import { guard } from "../_guard";

export const runtime = "nodejs";
export const maxDuration = 60; // SD / cold-start models can take a while
export const dynamic = "force-dynamic";

/**
 * Generative image-edit proxy. The BACKEND is pluggable — pick one with env
 * `AI_EDIT_PROVIDER` (default `auto`):
 *
 *   - `runpod`      → RunPod serverless GPU endpoints (one per model). Maps each
 *                     op → endpoint: restore/upscale → Real-ESRGAN (#7), colorize
 *                     → img2img (#10), replace-sky / magic-eraser / generative-fill
 *                     → SD 3.5 masked inpaint (#9), prompt → SD 3.5 img2img (#10).
 *                     Env: RUNPOD_API_KEY + per-model RUNPOD_*_URL. Key stays
 *                     server-side.
 *   - `local`       → your own Stable Diffusion server (Automatic1111 / Forge /
 *                     SD.Next img2img API). Env: LOCAL_SD_URL.
 *   - `huggingface` → Hugging Face Inference API. Env: HF_API_TOKEN, HF_IMAGE_MODEL.
 *   - `gemini`      → Google Gemini image model (needs a billed key for image output).
 *                     Env: GEMINI_API_KEY, GEMINI_IMAGE_MODEL.
 *   - `auto`        → first configured of: runpod → local → huggingface → gemini.
 *
 * NOTE: `remove-background` runs in-browser by default (@imgly, no key), so it
 * usually never reaches here. Object detection uses its own route (./classify).
 *
 * POST { imageBase64, mimeType?, op, maskBase64?, params? } -> { imageBase64, mimeType }
 * Auth: session-gated. Rate limit: 12/min (the most expensive route).
 */

const OP_PROMPTS: Record<string, string> = {
  restore:
    "Restore and enhance this photograph: improve sharpness and clarity, correct exposure and white balance, reduce noise and compression artifacts, recover detail. Keep it natural and photorealistic.",
  colorize: "Colorize this image with natural, realistic, well-balanced colors.",
  "replace-sky":
    "Replace the sky with a dramatic, beautiful golden-hour sky with soft clouds. Keep the foreground subject unchanged and the result photorealistic.",
};

const MAX_BASE64 = 4_000_000; // ~3 MB decoded — stays under serverless body limits

type Provider = "runpod" | "local" | "huggingface" | "gemini" | "none";

function resolveProvider(): Provider {
  const explicit = (process.env.AI_EDIT_PROVIDER || "auto").toLowerCase();
  if (
    explicit === "runpod" ||
    explicit === "local" ||
    explicit === "huggingface" ||
    explicit === "gemini"
  )
    return explicit;
  if (explicit === "none") return "none";
  // auto: prefer RunPod GPU endpoints, then a private local server, then HF, then Gemini.
  // Detect RunPod when the key + ANY image endpoint URL is set (an upscale/colorize-only
  // deployment is valid — not just the SD ones).
  if (
    process.env.RUNPOD_API_KEY &&
    (process.env.RUNPOD_SD_IMG2IMG_URL ||
      process.env.RUNPOD_SD_INPAINT_URL ||
      process.env.RUNPOD_UPSCALE_URL ||
      process.env.RUNPOD_COLORIZE_URL ||
      process.env.RUNPOD_BG_REMOVE_URL)
  )
    return "runpod";
  if (process.env.LOCAL_SD_URL) return "local";
  if (process.env.HF_API_TOKEN) return "huggingface";
  if (process.env.GEMINI_API_KEY) return "gemini";
  return "none";
}

/** Ops that only the RunPod (mask/fixed-function) backend can serve. */
const RUNPOD_ONLY_OPS = new Set(["upscale", "magic-eraser", "generative-fill"]);

interface EditResult {
  imageBase64: string;
  mimeType: string;
}

export async function POST(req: NextRequest) {
  const gate = await guard(req, "edit");
  if (!gate.ok) return gate.response;

  const provider = resolveProvider();
  if (provider === "none") {
    return NextResponse.json(
      {
        error:
          "AI image editing is not configured. Set AI_EDIT_PROVIDER=runpod + RUNPOD_API_KEY + the per-model RUNPOD_*_URL vars (RunPod GPU), or LOCAL_SD_URL (own Stable Diffusion), HF_API_TOKEN (Hugging Face), or GEMINI_API_KEY. Background removal and all analysis still work with no key.",
      },
      { status: 503 },
    );
  }

  let body: unknown;
  try {
    body = await req.json();
  } catch {
    return NextResponse.json({ error: "Invalid request body." }, { status: 400 });
  }

  const { imageBase64, mimeType, op, maskBase64, params } = (body ?? {}) as {
    imageBase64?: unknown;
    mimeType?: unknown;
    op?: { type?: string; prompt?: string; factor?: number };
    maskBase64?: unknown;
    params?: unknown;
  };

  if (typeof imageBase64 !== "string" || imageBase64.length === 0) {
    return NextResponse.json({ error: "Invalid image." }, { status: 400 });
  }
  const hasMask = typeof maskBase64 === "string" && maskBase64.length > 0;
  // Image + mask share one request body — budget them together against the cap.
  if (imageBase64.length + (hasMask ? (maskBase64 as string).length : 0) > MAX_BASE64) {
    return NextResponse.json(
      { error: "Image (plus mask) is too large — try a smaller image." },
      { status: 400 },
    );
  }
  const safeMime =
    typeof mimeType === "string" && /^image\/(jpeg|png|webp)$/.test(mimeType)
      ? mimeType
      : "image/jpeg";

  const opType = op?.type ?? "";
  if (provider !== "runpod" && RUNPOD_ONLY_OPS.has(opType)) {
    return NextResponse.json(
      { error: "This edit needs the RunPod backend (set AI_EDIT_PROVIDER=runpod)." },
      { status: 400 },
    );
  }

  // Build the instruction from an allow-listed op (never trust arbitrary server prompts).
  let instruction = "";
  if (opType === "prompt" || opType === "generative-fill") {
    const p = typeof op?.prompt === "string" ? op.prompt.trim() : "";
    if (!p) return NextResponse.json({ error: "Empty prompt." }, { status: 400 });
    instruction = p.slice(0, 500);
  } else if (opType === "replace-sky") {
    instruction =
      typeof op?.prompt === "string" && op.prompt.trim()
        ? `Replace the sky with: ${op.prompt.trim().slice(0, 300)}. Keep the foreground unchanged and photorealistic.`
        : OP_PROMPTS["replace-sky"]!;
  } else if (opType === "magic-eraser") {
    instruction =
      "Fill the selected region with a clean, seamless, plausible background. Photorealistic.";
  } else if (OP_PROMPTS[opType]) {
    instruction = OP_PROMPTS[opType]!;
  } else if (opType !== "upscale" && opType !== "remove-background") {
    return NextResponse.json({ error: "Unsupported operation." }, { status: 400 });
  }

  try {
    let result: EditResult;
    if (provider === "runpod")
      result = await editRunPod(
        op ?? {},
        imageBase64,
        instruction,
        hasMask ? (maskBase64 as string) : undefined,
        params,
      );
    else if (provider === "local") result = await editLocal(instruction, imageBase64);
    else if (provider === "huggingface") result = await editHuggingFace(instruction, imageBase64);
    else result = await editGemini(instruction, imageBase64, safeMime);
    return NextResponse.json(result);
  } catch (err) {
    const message = err instanceof Error ? err.message : "AI request failed.";
    const status = err instanceof AiError || err instanceof RunpodError ? err.status : 502;
    return NextResponse.json({ error: message }, { status });
  }
}

// ---------------------------------------------------------------------------
// Backend: RunPod serverless GPU endpoints (one model per endpoint).
// Each op maps to its endpoint; the API key + URLs stay server-side.
// ---------------------------------------------------------------------------
interface SdParams {
  negativePrompt?: string;
  strength?: number;
  steps?: number;
  seed?: number;
  guidanceScale?: number;
}

function sanitizeParams(raw: unknown): SdParams {
  const p = (raw ?? {}) as Record<string, unknown>;
  const out: SdParams = {};
  if (typeof p.negativePrompt === "string" && p.negativePrompt.trim())
    out.negativePrompt = p.negativePrompt.trim().slice(0, 300);
  const strength = Number(p.strength);
  if (Number.isFinite(strength)) out.strength = Math.max(0, Math.min(1, strength));
  const steps = Number(p.steps);
  if (Number.isFinite(steps)) out.steps = Math.max(1, Math.min(60, Math.round(steps)));
  const guidance = Number(p.guidanceScale);
  if (Number.isFinite(guidance)) out.guidanceScale = Math.max(1, Math.min(20, guidance));
  const seed = Number(p.seed);
  if (Number.isFinite(seed)) out.seed = Math.max(0, Math.min(2_147_483_647, Math.round(seed)));
  return out;
}

async function editRunPod(
  op: { type?: string; prompt?: string; factor?: number },
  imageBase64: string,
  instruction: string,
  maskBase64: string | undefined,
  rawParams: unknown,
): Promise<EditResult> {
  const params = sanitizeParams(rawParams);
  switch (op.type) {
    case "remove-background":
      // U²-Net via rembg (#6) — a real endpoint replacing the flaky in-browser remover.
      return rpRemoveBackground(imageBase64);
    case "restore":
      // Real-ESRGAN (#7) with the GFPGAN face pass = "Restore & Enhance".
      return rpUpscale(imageBase64, 4, true);
    case "upscale":
      return rpUpscale(imageBase64, op.factor === 4 ? 4 : 2, false);
    case "colorize":
      // The dedicated DDColor endpoint kept hard-crashing (modelscope). Route
      // colorize through the img2img model as an instruction instead.
      return rpImg2Img({ imageB64: imageBase64, prompt: instruction, ...params });
    case "prompt":
      return rpImg2Img({ imageB64: imageBase64, prompt: instruction, ...params }); // SD 3.5 img2img (#10)
    case "replace-sky":
      // True sky replacement is masked inpaint (#9). Without a mask (no in-app sky
      // segmentation yet) degrade to a low-strength img2img (#10) so the foreground
      // is mostly preserved.
      if (maskBase64)
        return rpInpaint({
          imageB64: imageBase64,
          maskB64: maskBase64,
          prompt: instruction,
          ...params,
        });
      return rpImg2Img({
        imageB64: imageBase64,
        prompt: instruction,
        ...params,
        strength: params.strength ?? 0.4,
      });
    case "magic-eraser":
    case "generative-fill":
      // SD 3.5 masked inpaint (#9) — white in the mask = the region to regenerate.
      if (!maskBase64) throw new AiError("This edit needs a mask/selection.", 400);
      return rpInpaint({
        imageB64: imageBase64,
        maskB64: maskBase64,
        prompt: instruction,
        ...params,
      });
    default:
      throw new AiError("Unsupported operation.", 400);
  }
}

class AiError extends Error {
  status: number;
  constructor(message: string, status = 502) {
    super(message);
    this.status = status;
  }
}

// ---------------------------------------------------------------------------
// Backend: local Stable Diffusion (Automatic1111 / Forge / SD.Next img2img API)
// ---------------------------------------------------------------------------
async function editLocal(instruction: string, imageBase64: string): Promise<EditResult> {
  const base = process.env.LOCAL_SD_URL;
  if (!base || !/^https?:\/\//i.test(base)) {
    throw new AiError("LOCAL_SD_URL is not a valid http(s) URL.", 500);
  }
  const url = `${base.replace(/\/$/, "")}/sdapi/v1/img2img`;
  let res: Response;
  try {
    res = await fetch(url, {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({
        init_images: [imageBase64],
        prompt: instruction,
        denoising_strength: Number(process.env.LOCAL_SD_DENOISE ?? 0.55),
        steps: Number(process.env.LOCAL_SD_STEPS ?? 25),
        cfg_scale: 7,
        sampler_name: process.env.LOCAL_SD_SAMPLER || "Euler a",
      }),
      cache: "no-store",
    });
  } catch {
    throw new AiError("Could not reach your local Stable Diffusion server (LOCAL_SD_URL).", 502);
  }
  if (!res.ok) {
    throw new AiError(`Local SD server error (${res.status}).`, 502);
  }
  const data = (await res.json().catch(() => null)) as { images?: string[] } | null;
  const out = data?.images?.[0];
  if (!out) throw new AiError("Local SD server did not return an image.", 502);
  // A1111 returns raw base64 PNG (no data: prefix).
  return { imageBase64: out.includes(",") ? out.split(",")[1]! : out, mimeType: "image/png" };
}

// ---------------------------------------------------------------------------
// Backend: Hugging Face Inference API — instruction image editing.
// ---------------------------------------------------------------------------
async function editHuggingFace(instruction: string, imageBase64: string): Promise<EditResult> {
  const token = process.env.HF_API_TOKEN;
  if (!token) throw new AiError("HF_API_TOKEN is not set.", 500);
  const model = process.env.HF_IMAGE_MODEL || "timbrooks/instruct-pix2pix";
  let res: Response;
  try {
    res = await fetch(`https://api-inference.huggingface.co/models/${model}`, {
      method: "POST",
      headers: {
        authorization: `Bearer ${token}`,
        "content-type": "application/json",
        // Wait for the model to warm up instead of a fast 503.
        "x-wait-for-model": "true",
      },
      body: JSON.stringify({
        inputs: imageBase64,
        parameters: { prompt: instruction, guidance_scale: 7, image_guidance_scale: 1.5 },
      }),
      cache: "no-store",
    });
  } catch {
    throw new AiError("Could not reach the Hugging Face Inference API.", 502);
  }
  if (!res.ok) {
    // Truncated on purpose — never surface a full upstream body.
    const detail = (await res.text().catch(() => "")).slice(0, 160);
    if (res.status === 503) throw new AiError("The model is loading — try again in ~20s.", 503);
    throw new AiError(`Hugging Face error (${res.status}). ${detail}`, 502);
  }
  // Success returns raw image bytes.
  const outMime = res.headers.get("content-type") || "image/png";
  if (outMime.startsWith("application/json")) {
    const j = (await res.json().catch(() => null)) as { error?: string } | null;
    throw new AiError(
      j?.error ? `Hugging Face: ${j.error}` : "Hugging Face returned no image.",
      502,
    );
  }
  const buf = await res.arrayBuffer();
  return { imageBase64: Buffer.from(buf).toString("base64"), mimeType: outMime };
}

// ---------------------------------------------------------------------------
// Backend: Google Gemini image model (needs a billed key for image output).
// ---------------------------------------------------------------------------
async function editGemini(
  instruction: string,
  imageBase64: string,
  safeMime: string,
): Promise<EditResult> {
  const apiKey = process.env.GEMINI_API_KEY;
  if (!apiKey) throw new AiError("GEMINI_API_KEY is not set.", 500);
  const model = process.env.GEMINI_IMAGE_MODEL || "gemini-2.5-flash-image";
  const prompt = `Edit this image as follows: ${instruction}. Preserve realism unless explicitly asked otherwise.`;
  let res: Response;
  try {
    res = await fetch(
      `https://generativelanguage.googleapis.com/v1beta/models/${model}:generateContent`,
      {
        method: "POST",
        headers: { "content-type": "application/json", "x-goog-api-key": apiKey },
        body: JSON.stringify({
          contents: [
            {
              role: "user",
              parts: [
                { inlineData: { mimeType: safeMime, data: imageBase64 } },
                { text: prompt },
              ],
            },
          ],
          generationConfig: { responseModalities: ["IMAGE"] },
        }),
        cache: "no-store",
      },
    );
  } catch {
    throw new AiError("Could not reach the AI service.", 502);
  }
  if (!res.ok) {
    // Truncated on purpose — never surface a full upstream body.
    const detail = (await res.text().catch(() => "")).slice(0, 160);
    throw new AiError(`AI service error (${res.status}). ${detail}`, 502);
  }
  const data = (await res.json().catch(() => null)) as GeminiResponse | null;
  const parts = data?.candidates?.[0]?.content?.parts ?? [];
  const imgPart = parts.find((p) => p.inlineData?.data || p.inline_data?.data);
  const out = imgPart?.inlineData?.data ?? imgPart?.inline_data?.data;
  if (!out)
    throw new AiError(
      "The model did not return an image (the free Gemini tier has no image output — use LOCAL_SD_URL or HF_API_TOKEN instead).",
      502,
    );
  const outMime = imgPart?.inlineData?.mimeType ?? imgPart?.inline_data?.mime_type ?? "image/png";
  return { imageBase64: out, mimeType: outMime };
}

interface GeminiPart {
  text?: string;
  inlineData?: { mimeType?: string; data?: string };
  inline_data?: { mime_type?: string; data?: string };
}
interface GeminiResponse {
  candidates?: Array<{ content?: { parts?: GeminiPart[] } }>;
}