import { NextResponse } from "next/server";

import { normalizeUpstreamChatEndpoint } from "~/lib/lab2-chat";
import {
  clampLab1Messages,
  extractLab1AssistantContent,
  extractLab1ResponseTokens,
  getLab1SystemPrompt,
  LAB1_CONFIDENCE_MODEL_ALIAS,
  LAB1_DEFAULT_MAX_TOKENS,
  LAB1_DEFAULT_TEMPERATURE,
  type Lab1ConfidenceMessage,
} from "~/lib/lab1-confidence";

type ChatRouteRequestBody = {
  messages?: Lab1ConfidenceMessage[];
};

const LOCAL_OLLAMA_TIMEOUT_MS = 90000;

function getLocalOllamaEndpoint() {
  const configuredBaseUrl =
    process.env.COURSEWARE_OLLAMA_BASE_URL?.trim() || "http://127.0.0.1:11434";

  return normalizeUpstreamChatEndpoint(configuredBaseUrl);
}

function getLab1ModelAlias() {
  return (
    process.env.COURSEWARE_LAB1_OLLAMA_MODEL_ALIAS?.trim() ||
    LAB1_CONFIDENCE_MODEL_ALIAS
  );
}

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

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

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

  const controller = new AbortController();
  const timeoutId = setTimeout(
    () => controller.abort(),
    LOCAL_OLLAMA_TIMEOUT_MS,
  );

  try {
    const upstreamResponse = await fetch(getLocalOllamaEndpoint(), {
      body: JSON.stringify({
        logprobs: true,
        max_tokens: LAB1_DEFAULT_MAX_TOKENS,
        messages: [
          {
            content: getLab1SystemPrompt(),
            role: "system",
          },
          ...clampLab1Messages(body.messages),
        ],
        model: getLab1ModelAlias(),
        stream: false,
        temperature: LAB1_DEFAULT_TEMPERATURE,
        top_logprobs: 5,
      }),
      headers: {
        "Content-Type": "application/json",
      },
      method: "POST",
      signal: controller.signal,
    });

    const responseText = await upstreamResponse.text();
    const parsedBody = JSON.parse(responseText) as unknown;

    if (!upstreamResponse.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 local Ollama endpoint returned ${upstreamResponse.status}.`;

      return NextResponse.json(
        {
          error: message,
        },
        { status: upstreamResponse.status },
      );
    }

    if (!parsedBody || typeof parsedBody !== "object") {
      return NextResponse.json(
        {
          error: "The local Ollama endpoint returned an unreadable response.",
        },
        { status: 502 },
      );
    }

    const tokens = extractLab1ResponseTokens(parsedBody);
    if (tokens.length === 0) {
      return NextResponse.json(
        {
          error:
            "The local Ollama response did not include token logprobs. Confirm the installed Ollama version supports logprobs.",
        },
        { status: 502 },
      );
    }

    const content =
      extractLab1AssistantContent(parsedBody) ||
      tokens.map((token) => token.token).join("");

    return NextResponse.json({
      content,
      model:
        ("model" in parsedBody && typeof parsedBody.model === "string"
          ? parsedBody.model
          : getLab1ModelAlias()),
      role: "assistant",
      tokens,
    });
  } catch (caughtError) {
    if (caughtError instanceof Error && caughtError.name === "AbortError") {
      return NextResponse.json(
        {
          error: `The local Ollama endpoint timed out after ${Math.floor(LOCAL_OLLAMA_TIMEOUT_MS / 1000)} seconds.`,
        },
        { status: 504 },
      );
    }

    return NextResponse.json(
      {
        error: "The Lab 1 confidence route could not reach the local Ollama endpoint.",
      },
      { status: 502 },
    );
  } finally {
    clearTimeout(timeoutId);
  }
}
