import type { ExpectationResult } from "../../expect/types";
import type { LlmCheckObservation, ResolvedLlmExpectConfig, ResolvedLlmUsageExpect } from "./types";

import { checkContentExpectations } from "../../expect/content";
import { mismatchFailure } from "../../expect/failure";
import { checkHeaderExpectations } from "../../expect/headers";
import { checkStatusCode } from "../../expect/status";
import { checkValueExpectation } from "../../expect/value";

export function checkStreamExpect(
  observation: LlmCheckObservation,
  expect: ResolvedLlmExpectConfig,
): ExpectationResult {
  if (!observation.stream || !expect.stream) return { failure: null, matched: true };

  const expectedCompleted = expect.stream.completed;
  if (observation.stream.completed !== expectedCompleted) {
    return {
      failure: mismatchFailure(
        "stream",
        "stream.completed",
        expectedCompleted,
        observation.stream.completed,
        "stream.completed mismatch",
      ),
      matched: false,
    };
  }

  if (expect.stream.firstTokenMs && observation.stream.firstTokenMs !== null) {
    return checkValueExpectation(observation.stream.firstTokenMs, expect.stream.firstTokenMs, {
      message: "stream.firstTokenMs mismatch",
      path: "stream.firstTokenMs",
      phase: "stream",
    });
  } else if (expect.stream.firstTokenMs && observation.stream.firstTokenMs === null) {
    return {
      failure: mismatchFailure(
        "stream",
        "stream.firstTokenMs",
        expect.stream.firstTokenMs,
        null,
        "stream.firstTokenMs missing",
      ),
      matched: false,
    };
  }

  return { failure: null, matched: true };
}

export function runExpects(
  observation: LlmCheckObservation,
  expect: ResolvedLlmExpectConfig | undefined,
): ExpectationResult {
  if (!expect) {
    const defaultStatus = checkStatusCode(observation.http?.status ?? 0, [200]);
    if (!defaultStatus.matched) return defaultStatus;
    return { failure: null, matched: true };
  }

  const http = observation.http;

  const statusResult = checkStatusCode(http?.status ?? 0, expect.status);
  if (!statusResult.matched) return statusResult;

  if (http && expect.headers) {
    const headersResult = checkHeaderExpectations(http.headers, expect.headers);
    if (!headersResult.matched) return headersResult;
  }

  if (observation.stream !== null) {
    const streamResult = checkStreamExpect(observation, expect);
    if (!streamResult.matched) return streamResult;
  }

  const outputResult = checkContentExpectations(observation.outputText, expect.output, {
    path: "output",
    phase: "output",
  });
  if (!outputResult.matched) return outputResult;

  if (expect.finishReason !== undefined) {
    const result = checkValueExpectation(observation.finishReason, expect.finishReason, {
      message: "finishReason mismatch",
      path: "finishReason",
      phase: "finishReason",
    });
    if (!result.matched) return result;
  }

  if (expect.rawFinishReason !== undefined) {
    const result = checkValueExpectation(observation.rawFinishReason, expect.rawFinishReason, {
      message: "rawFinishReason mismatch",
      path: "rawFinishReason",
      phase: "rawFinishReason",
    });
    if (!result.matched) return result;
  }

  if (expect.usage && observation.usage) {
    const usageResult = checkUsageExpect(observation.usage, expect.usage);
    if (!usageResult.matched) return usageResult;
  }

  return { failure: null, matched: true };
}

function checkUsageExpect(
  usage: { inputTokens: number; outputTokens: number; totalTokens: number },
  expectUsage: ResolvedLlmUsageExpect,
): ExpectationResult {
  if (expectUsage.inputTokens !== undefined) {
    const result = checkValueExpectation(usage.inputTokens, expectUsage.inputTokens, {
      message: "usage.inputTokens mismatch",
      path: "usage.inputTokens",
      phase: "usage",
    });
    if (!result.matched) return result;
  }
  if (expectUsage.outputTokens !== undefined) {
    const result = checkValueExpectation(usage.outputTokens, expectUsage.outputTokens, {
      message: "usage.outputTokens mismatch",
      path: "usage.outputTokens",
      phase: "usage",
    });
    if (!result.matched) return result;
  }
  if (expectUsage.totalTokens !== undefined) {
    const result = checkValueExpectation(usage.totalTokens, expectUsage.totalTokens, {
      message: "usage.totalTokens mismatch",
      path: "usage.totalTokens",
      phase: "usage",
    });
    if (!result.matched) return result;
  }
  return { failure: null, matched: true };
}
