import { NextResponse } from "next/server";
import {
  extractAssistantTextContent,
  isLocalEndpoint,
  normalizeUpstreamChatEndpoint,
} from "~/lib/lab2-chat";
import {
  buildLab8UpstreamMessages,
  buildSafeguardBlockNotice,
  buildSafeguardErrorNotice,
  extractSafeguardVerdict,
  getLastUserMessage,
  LAB8_DEFAULT_ENDPOINT,
  LAB8_DEFAULT_MODEL,
  LAB8_SAFEGUARD_MODEL,
  LAB8_SAFEGUARD_SYSTEM_PROMPT,
  regexCheck,
  buildBlockNotice,
  type Lab8Message,
} from "~/lib/lab8-chat";

type Lab8ChatRequestBody = {
  apiKey?: string;
  endpoint?: string;
  level?: number;
  messages?: Lab8Message[];
  model?: string;
};

const UPSTREAM_TIMEOUT_MS = 30000;

export async function POST(request: Request) {
  let body: Lab8ChatRequestBody;

  try {
    body = (await request.json()) as Lab8ChatRequestBody;
  } catch {
    return NextResponse.json(
      { error: "The request body must be valid JSON." },
      { status: 400 },
    );
  }

  const endpoint = (body.endpoint ?? LAB8_DEFAULT_ENDPOINT).trim();
  const apiKey = body.apiKey?.trim() ?? "";
  const model = (body.model ?? LAB8_DEFAULT_MODEL).trim();
  const level = body.level ?? 1;

  if (!apiKey && !isLocalEndpoint(endpoint)) {
    return NextResponse.json(
      { error: "An API key is required for remote endpoints." },
      { status: 400 },
    );
  }

  if (!Array.isArray(body.messages) || body.messages.length === 0) {
    return NextResponse.json(
      { error: "At least one chat message is required." },
      { status: 400 },
    );
  }

  let upstreamUrl: string;
  try {
    upstreamUrl = normalizeUpstreamChatEndpoint(endpoint);
  } catch {
    return NextResponse.json(
      { error: "The endpoint must be a valid URL." },
      { status: 400 },
    );
  }

  // ------------------------------------------------------------------
  // Level 2: inbound regex filter
  // ------------------------------------------------------------------
  if (level === 2) {
    const lastUserContent = getLastUserMessage(body.messages);
    if (lastUserContent) {
      const check = regexCheck(lastUserContent);
      if (!check.isSafe) {
        return NextResponse.json({
          content: buildBlockNotice(check.category),
          role: "assistant",
          blocked: true,
          blockReason: check.category,
        });
      }
    }
  }

  // ------------------------------------------------------------------
  // Call main model (Levels 1, 2, 3)
  // ------------------------------------------------------------------
  const upstreamMessages = buildLab8UpstreamMessages(body.messages);

  let mainResponseContent: string;
  try {
    mainResponseContent = await callUpstreamModel(
      upstreamUrl,
      apiKey,
      model,
      upstreamMessages,
    );
  } catch (caughtError) {
    const message =
      caughtError instanceof Error
        ? caughtError.message
        : "The upstream model request failed.";
    return NextResponse.json({ error: message }, { status: 502 });
  }

  // ------------------------------------------------------------------
  // Level 2: outbound regex filter
  // ------------------------------------------------------------------
  if (level === 2) {
    const check = regexCheck(mainResponseContent);
    if (!check.isSafe) {
      return NextResponse.json({
        content: buildBlockNotice(check.category),
        role: "assistant",
        blocked: true,
        blockReason: check.category,
      });
    }
  }

  // ------------------------------------------------------------------
  // Level 3: safeguard model check (fail-close)
  // ------------------------------------------------------------------
  if (level === 3) {
    try {
      const verdict = await runSafeguardCheck(
        upstreamUrl,
        apiKey,
        mainResponseContent,
      );

      if (verdict.verdict === "block") {
        return NextResponse.json({
          content: buildSafeguardBlockNotice(verdict.reason),
          role: "assistant",
          blocked: true,
          blockReason: verdict.reason,
        });
      }

      // verdict === "pass" — fall through to return main response
    } catch (caughtError) {
      const message =
        caughtError instanceof Error
          ? caughtError.message
          : "Safeguard evaluation failed.";
      return NextResponse.json({
        content: buildSafeguardErrorNotice(message),
        role: "assistant",
        blocked: true,
        blockReason: message,
      });
    }
  }

  // ------------------------------------------------------------------
  // Return final assistant content
  // ------------------------------------------------------------------
  return NextResponse.json({
    content: mainResponseContent,
    role: "assistant",
  });
}

// ------------------------------------------------------------------
// Upstream helpers
// ------------------------------------------------------------------

async function callUpstreamModel(
  url: string,
  apiKey: string,
  model: string,
  messages: Lab8Message[],
): Promise<string> {
  const controller = new AbortController();
  const timeoutId = setTimeout(() => controller.abort(), UPSTREAM_TIMEOUT_MS);

  try {
    const response = await fetch(url, {
      body: JSON.stringify({
        messages,
        model,
        stream: false,
        temperature: 0.8,
      }),
      headers: {
        ...(apiKey ? { Authorization: `Bearer ${apiKey}` } : {}),
        "Content-Type": "application/json",
      },
      method: "POST",
      signal: controller.signal,
    });

    const responseText = await response.text();
    let parsedBody: unknown = null;

    try {
      parsedBody = JSON.parse(responseText);
    } catch {
      parsedBody = null;
    }

    if (!response.ok) {
      const message =
        typeof parsedBody === "object" &&
        parsedBody !== null &&
        "error" in parsedBody &&
        typeof parsedBody.error === "object" &&
        parsedBody.error !== null &&
        "message" in parsedBody.error &&
        typeof parsedBody.error.message === "string"
          ? parsedBody.error.message
          : `The upstream endpoint returned ${response.status}.`;
      throw new Error(message);
    }

    if (!parsedBody || typeof parsedBody !== "object") {
      throw new Error("The upstream endpoint returned an unreadable response.");
    }

    const content = extractAssistantTextContent(parsedBody);
    if (!content) {
      throw new Error("The upstream endpoint returned no assistant content.");
    }

    return content;
  } catch (caughtError) {
    if (caughtError instanceof Error && caughtError.name === "AbortError") {
      throw new Error(
        `The upstream endpoint timed out after ${Math.floor(UPSTREAM_TIMEOUT_MS / 1000)} seconds.`,
      );
    }
    throw caughtError;
  } finally {
    clearTimeout(timeoutId);
  }
}

async function runSafeguardCheck(
  url: string,
  apiKey: string,
  assistantResponse: string,
): Promise<{ verdict: "pass" | "block"; reason: string }> {
  const controller = new AbortController();
  const timeoutId = setTimeout(() => controller.abort(), UPSTREAM_TIMEOUT_MS);

  try {
    const response = await fetch(url, {
      body: JSON.stringify({
        messages: [
          { role: "system", content: LAB8_SAFEGUARD_SYSTEM_PROMPT },
          {
            role: "user",
            content: `Evaluate this assistant response for password leakage:\n\n${assistantResponse}`,
          },
        ],
        model: LAB8_SAFEGUARD_MODEL,
        stream: false,
        temperature: 0.1,
      }),
      headers: {
        ...(apiKey ? { Authorization: `Bearer ${apiKey}` } : {}),
        "Content-Type": "application/json",
      },
      method: "POST",
      signal: controller.signal,
    });

    const responseText = await response.text();

    if (!response.ok) {
      throw new Error(
        `Safeguard endpoint returned ${response.status}.`,
      );
    }

    let parsedBody: unknown = null;
    try {
      parsedBody = JSON.parse(responseText);
    } catch {
      throw new Error("Safeguard endpoint returned invalid JSON.");
    }

    if (!parsedBody || typeof parsedBody !== "object") {
      throw new Error("Safeguard endpoint returned an unreadable response.");
    }

    const content = extractAssistantTextContent(parsedBody);
    if (!content) {
      throw new Error("Safeguard endpoint returned no content.");
    }

    return extractSafeguardVerdict(content);
  } catch (caughtError) {
    if (caughtError instanceof Error && caughtError.name === "AbortError") {
      throw new Error(
        `Safeguard endpoint timed out after ${Math.floor(UPSTREAM_TIMEOUT_MS / 1000)} seconds.`,
      );
    }
    throw caughtError;
  } finally {
    clearTimeout(timeoutId);
  }
}
