import { BraintrustState, initLogger } from "braintrust";
import type { Env, ExecutionTrace } from "./types";

const APP_URL = "https://www.braintrust.dev";
const API_URL = "https://api-eu.braintrust.dev";
export const BRAINTRUST_PROJECT_ID = "c2631407-ef97-405f-8655-14ff8b2ac4aa";
const ORG_ID = "5ddc77ce-16f9-4715-aa21-0405aae67681";

export interface OperationalRun {
  requestId: string;
  mode: "chat" | "assessment";
  evaluation: boolean;
  startedAt: number;
  endedAt: number;
  outcome: "success" | "provider_error" | "timeout" | "agent_error";
  trace?: ExecutionTrace;
}

// Deliberately separate from model instrumentation: no prompts, completions,
// retrieval queries/passages, tool arguments, headers or visitor identifiers.
export async function exportOperationalRun(env: Env, run: OperationalRun) {
  if (!env.BRAINTRUST_API_KEY) return;
  const deadline = AbortSignal.timeout(10_000);
  let failed = false;
  const boundedFetch: typeof fetch = async (input, options) => {
    const url = new URL(
      typeof input === "string"
        ? input
        : input instanceof URL
        ? input.href
        : input.url
    );
    if (![APP_URL, API_URL].includes(url.origin))
      throw new Error("Unexpected telemetry destination");
    const controller = new AbortController();
    const signals = [
      deadline,
      AbortSignal.timeout(2500),
      ...(options?.signal ? [options.signal] : []),
    ];
    const abort = () => controller.abort();
    for (const signal of signals) {
      if (signal.aborted) controller.abort();
      else signal.addEventListener("abort", abort, { once: true });
    }
    let response: Response;
    try {
      const upstream = await fetch(input, {
        ...options,
        redirect: "manual",
        signal: controller.signal,
      });
      if (upstream.status >= 300 && upstream.status < 400)
        throw new Error("Telemetry redirects are disabled");
      // Keep cancellation active while reading the response body too.
      response = new Response(await upstream.arrayBuffer(), {
        status: upstream.status,
        headers: upstream.headers,
      });
    } finally {
      for (const signal of signals) signal.removeEventListener("abort", abort);
    }
    if (url.pathname === "/api/apikey/login" && response.ok) {
      const info = await response.json<any>();
      const org = info.org_info?.find((item: any) => item.id === ORG_ID);
      if (!org || org.api_url !== API_URL)
        throw new Error("Unexpected telemetry organization or region");
      return Response.json({ ...info, org_info: [org] });
    }
    return response;
  };
  try {
    // Each export owns its state and spans. Concurrent requests never share an
    // ambient logger/context, and a rotated secret takes effect immediately.
    const state = new BraintrustState({
      apiKey: env.BRAINTRUST_API_KEY,
      appUrl: APP_URL,
      fetch: boundedFetch,
      noExitFlush: true,
      debugLogLevel: false,
      onFlushError: () => {
        failed = true;
      },
    });
    const logger = initLogger({
      projectId: BRAINTRUST_PROJECT_ID,
      projectName: "aleph.fi",
      state,
      setCurrent: false,
      asyncFlush: true,
      debugLogLevel: false,
    });
    const trace = run.trace;
    const metrics: Record<string, number> = {
      duration_ms: run.endedAt - run.startedAt,
      errors: run.outcome === "success" ? 0 : 1,
    };
    if (trace) {
      metrics.retrieval_ms = trace.retrievalMs;
      metrics.model_calls = trace.modelCalls;
      metrics.source_count = trace.sources.length;
      metrics.tool_calls = trace.tools.length;
      if (trace.firstTokenMs !== null)
        metrics.time_to_first_token = trace.firstTokenMs / 1000;
      if (trace.usageAvailable) {
        metrics.prompt_tokens = trace.inputTokens;
        metrics.completion_tokens = trace.outputTokens;
        metrics.tokens = trace.inputTokens + trace.outputTokens;
      }
      if (trace.costUsd !== null) metrics.cost = trace.costUsd;
    }
    const span = logger.startSpan({
      name: "portfolio.agent",
      startTime: run.startedAt / 1000,
      spanAttributes: { type: "task" },
      event: {
        metadata: {
          request_id: run.requestId,
          mode: run.mode,
          environment: run.evaluation ? "evaluation" : "production",
          outcome: run.outcome,
          telemetry_policy: "operational-metadata-v1",
          code_version: env.BUILD_COMMIT,
          ...(trace
            ? {
                model: trace.model,
                models: trace.models ?? [trace.model],
                embedding_model: trace.embeddingModel,
                prompt_version: trace.promptVersion,
                corpus_version: trace.corpusVersion,
                evidence_version: trace.evidenceVersion,
                usage_available: trace.usageAvailable,
                optimization: trace.optimization,
                tools: trace.tools.map(tool => ({
                  name: [
                    "get_project_evidence",
                    "recommend_content",
                    "show_contact_card",
                  ].includes(tool.name)
                    ? tool.name
                    : "unknown_tool",
                  status: tool.status === "ok" ? "ok" : "unavailable",
                  duration_ms: tool.durationMs,
                  ...(tool.httpStatus !== undefined
                    ? { http_status: tool.httpStatus }
                    : {}),
                })),
              }
            : {}),
        },
        metrics,
        scores: { completed: run.outcome === "success" ? 1 : 0 },
      },
    });
    for (const stage of trace?.stages ?? []) {
      const allowedNames = [
        "retrieval.context",
        "model.call",
        "model.answer",
        "model.assessment",
        "model.followups",
        "model.judge",
        "validation.assessment",
        "get_project_evidence",
        "recommend_content",
        "show_contact_card",
        "unknown_tool",
      ];
      const child = span.startSpan({
        name: allowedNames.includes(stage.name) ? stage.name : "operation",
        startTime: stage.startedAt / 1000,
        spanAttributes: {
          type:
            stage.kind === "model"
              ? "llm"
              : stage.kind === "tool"
              ? "tool"
              : "task",
        },
        event: {
          metadata: {
            stage: stage.kind,
            status: stage.status,
            ...(stage.cacheHit !== undefined
              ? { cache_hit: stage.cacheHit }
              : {}),
            ...(stage.failure ? { failure: stage.failure } : {}),
            ...(stage.model ? { model: stage.model } : {}),
          },
          metrics: {
            errors: ["failed", "unavailable", "aborted"].includes(stage.status)
              ? 1
              : 0,
            ...(stage.inputTokens !== undefined
              ? {
                  prompt_tokens: stage.inputTokens,
                  completion_tokens: stage.outputTokens ?? 0,
                  tokens: stage.inputTokens + (stage.outputTokens ?? 0),
                }
              : {}),
            ...(stage.costUsd !== undefined ? { cost: stage.costUsd } : {}),
          },
        },
      });
      child.end({ endTime: (stage.endedAt ?? run.endedAt) / 1000 });
    }
    span.end({ endTime: run.endedAt / 1000 });
    await logger.flush();
    if (failed) throw new Error("Telemetry export failed");
  } catch {
    // Fixed text only: SDK/provider errors may contain response bodies.
    console.warn("Braintrust operational telemetry export failed");
  }
}