export const LAB1_CONFIDENCE_MODEL_ALIAS = "batiai/gemma4-e2b:q4";
export const LAB1_DEFAULT_MAX_TOKENS = 64;
export const LAB1_DEFAULT_TEMPERATURE = 0.7;
export const LAB1_MAX_CONTEXT_MESSAGES = 10;
export const LAB1_MAX_MESSAGE_LENGTH = 4000;

export type Lab1ConfidenceRole = "assistant" | "user";

export type Lab1ConfidenceMessage = {
  content: string;
  role: Lab1ConfidenceRole;
};

export type Lab1TopAlternative = {
  probability: number;
  token: string;
};

export type Lab1ResponseToken = {
  logprob: number;
  probability: number;
  token: string;
  topAlternatives: Lab1TopAlternative[];
};

export type Lab1ConfidenceResponse = {
  content: string;
  model: string;
  role: "assistant";
  tokens: Lab1ResponseToken[];
};

type OpenAiLogprobAlternative = {
  logprob?: number;
  token?: string;
};

type OpenAiLogprobToken = {
  logprob?: number;
  token?: string;
  top_logprobs?: OpenAiLogprobAlternative[];
};

type OpenAiCompatibilityPayload = {
  choices?: Array<{
    logprobs?: {
      content?: OpenAiLogprobToken[];
    };
    message?: {
      content?: string;
    };
  }>;
  model?: string;
};

export function getLab1SystemPrompt() {
  return [
    "You are helping students inspect token-level confidence in a local language model.",
    "Reply clearly and concisely.",
    "Prefer one compact paragraph unless the user explicitly asks for a list.",
  ].join(" ");
}

export function clampLab1Messages(messages: Lab1ConfidenceMessage[]) {
  return messages
    .filter((message) => {
      return (
        (message.role === "assistant" || message.role === "user") &&
        typeof message.content === "string"
      );
    })
    .map((message) => {
      return {
        content: message.content.slice(0, LAB1_MAX_MESSAGE_LENGTH),
        role: message.role,
      } satisfies Lab1ConfidenceMessage;
    })
    .slice(-LAB1_MAX_CONTEXT_MESSAGES);
}

export function logprobToProbabilityPercent(logprob: number) {
  if (!Number.isFinite(logprob)) {
    return 0;
  }

  return roundProbability(Math.exp(logprob) * 100);
}

export function formatProbabilityPercent(probability: number) {
  if (probability >= 10) {
    return `${probability.toFixed(1)}%`;
  }

  if (probability >= 1) {
    return `${probability.toFixed(2)}%`;
  }

  if (probability > 0) {
    return `${probability.toFixed(3)}%`;
  }

  return "0%";
}

export function getConfidenceBand(probability: number) {
  if (probability >= 70) return "very-high";
  if (probability >= 40) return "high";
  if (probability >= 15) return "medium";
  if (probability >= 5) return "low";
  return "very-low";
}

export function extractLab1AssistantContent(
  payload: OpenAiCompatibilityPayload,
) {
  const content = payload.choices?.[0]?.message?.content?.trim();
  return content || null;
}

export function extractLab1ResponseTokens(
  payload: OpenAiCompatibilityPayload,
): Lab1ResponseToken[] {
  const rawTokens = payload.choices?.[0]?.logprobs?.content;
  if (!Array.isArray(rawTokens)) {
    return [];
  }

  return rawTokens.flatMap((rawToken) => {
    const token = rawToken.token ?? "";
    const logprob = rawToken.logprob;

    if (!token || typeof logprob !== "number" || !Number.isFinite(logprob)) {
      return [];
    }

    const topAlternatives = Array.isArray(rawToken.top_logprobs)
      ? rawToken.top_logprobs
          .flatMap((candidate) => {
            const candidateToken = candidate.token ?? "";
            const candidateLogprob = candidate.logprob;

            if (
              !candidateToken ||
              candidateToken === token ||
              typeof candidateLogprob !== "number" ||
              !Number.isFinite(candidateLogprob)
            ) {
              return [];
            }

            return [
              {
                probability: logprobToProbabilityPercent(candidateLogprob),
                token: candidateToken,
              } satisfies Lab1TopAlternative,
            ];
          })
          .slice(0, 5)
      : [];

    return [
      {
        logprob,
        probability: logprobToProbabilityPercent(logprob),
        token,
        topAlternatives,
      } satisfies Lab1ResponseToken,
    ];
  });
}

function roundProbability(value: number) {
  return Math.max(0, Math.min(100, Math.round(value * 100) / 100));
}
