import { NextResponse } from "next/server";

import { getLlmTracePayloadStoreFromRuntime } from "./trace-payload-store";
import { getLlmTraceRepositoryFromRuntime } from "./trace-repository";
import type {
  LlmTraceCall,
  LlmTraceCallPublic,
  LlmTraceManifest,
} from "./trace-types";

export function noStoreJson(body: unknown, status = 200) {
  return NextResponse.json(body, {
    status,
    headers: { "cache-control": "no-store" },
  });
}

export function noStoreResponse<T extends Response>(response: T): T {
  response.headers.set("cache-control", "no-store");
  return response;
}

export function toPublicTraceCall(call: LlmTraceCall): LlmTraceCallPublic {
  const {
    request_object_key: requestObjectKey,
    response_object_key: responseObjectKey,
    ...publicFields
  } = call;
  return {
    ...publicFields,
    request_available: Boolean(requestObjectKey),
    response_available: Boolean(responseObjectKey),
  };
}

export async function getTraceManifest(
  jobId: string,
): Promise<LlmTraceManifest | null> {
  const repository = getLlmTraceRepositoryFromRuntime();
  const run = await repository.getRun(jobId);
  if (!run) return null;
  const calls = await repository.listCalls(jobId);
  return { run, calls: calls.map(toPublicTraceCall) };
}

export async function readTracePayload({
  jobId,
  callId,
  kind,
}: {
  jobId: string;
  callId: string;
  kind: "request" | "response";
}) {
  const repository = getLlmTraceRepositoryFromRuntime();
  const call = (await repository.listCalls(jobId)).find(
    (candidate) => candidate.call_id === callId,
  );
  if (!call) return noStoreJson({ error: "追踪调用不存在" }, 404);

  const key = kind === "request"
    ? call.request_object_key
    : call.response_object_key;
  if (!key) {
    if (kind === "response" && call.status === "started") {
      return noStoreJson({ state: "waiting" }, 202);
    }
    return noStoreJson({
      error: kind === "response" ? "该调用未产生响应" : "追踪请求正文不存在",
    }, 404);
  }

  const payload = await getLlmTracePayloadStoreFromRuntime().getJson(key);
  return payload == null
    ? noStoreJson({ error: "追踪正文不存在" }, 404)
    : noStoreJson(payload);
}