/**
 * Traced LLM Wrapper
 *
 * Wraps AI SDK's generateText with automatic event logging.
 * When called within a run context, emits llm_request_* events.
 */

import { generateText } from "ai";
import { getRunContext } from "./run-context";
import { appendRunEvent, getOrCreatePromptVersion } from "./queries";

// ============================================================================
// Types
// ============================================================================

/**
 * Options for traced LLM calls
 */
export interface TracedLLMOptions {
  /** Operation name for logging (e.g., "planner", "synthesizer", "judge") */
  llmOperation?: string;
  /** Prompt template key for versioning (e.g., "pipeline:enhance-prompt:system") */
  promptTemplateKey?: string;
  /** Step name if within a pipeline step */
  stepName?: string;
}

/**
 * Extended generateText options with tracing
 */
export type GenerateTextTracedOptions = Parameters<typeof generateText>[0] & TracedLLMOptions;

/**
 * Result type from generateTextTraced - mirrors AI SDK's GenerateTextResult
 */
export type GenerateTextTracedResult = Awaited<ReturnType<typeof generateText>>;

// ============================================================================
// Traced generateText
// ============================================================================

/**
 * Traced version of generateText that emits events to agent_run_events
 *
 * Usage:
 * ```ts
 * const result = await generateTextTraced({
 *   model: getUtilityModel(),
 *   system: SYSTEM_PROMPT,
 *   prompt: userPrompt,
 *   llmOperation: "synthesizer",
 *   promptTemplateKey: "pipeline:web-browse:synthesizer",
 * });
 * ```
 */
export async function generateTextTraced(
  options: GenerateTextTracedOptions
): Promise<GenerateTextTracedResult> {
  const ctx = getRunContext();
  const startTime = Date.now();

  // Extract tracing options
  const { llmOperation, promptTemplateKey, stepName, ...generateOptions } = options;

  // Get model name for logging - try multiple ways to access it
  let modelName = "unknown";
  if (options.model) {
    const model = options.model as Record<string, unknown>;
    if (typeof model.modelId === "string") {
      modelName = model.modelId;
    } else if (typeof model.id === "string") {
      modelName = model.id;
    } else if (typeof model === "string") {
      modelName = model;
    }
  }

  // Track prompt version if template key provided
  let promptVersionId: string | undefined;
  if (promptTemplateKey && options.system) {
    try {
      const { promptVersion } = await getOrCreatePromptVersion({
        templateKey: promptTemplateKey,
        content: typeof options.system === "string" ? options.system : JSON.stringify(options.system),
      });
      promptVersionId = promptVersion.id;
    } catch (error) {
      console.error("[generateTextTraced] Failed to track prompt version:", error);
    }
  }

  // Emit start event if in run context
  if (ctx) {
    try {
      // Access optional properties via unknown to avoid TS index signature errors
      const opts = options as unknown as Record<string, unknown>;
      await appendRunEvent({
        runId: ctx.runId,
        eventType: "llm_request_started",
        level: "info",
        pipelineName: ctx.pipelineName,
        stepName,
        llmOperation,
        promptVersionId,
        data: {
          model: modelName,
          maxOutputTokens: opts.maxOutputTokens,
          temperature: opts.temperature,
          hasTools: !!options.tools,
          toolCount: options.tools ? Object.keys(options.tools).length : 0,
        },
      });
    } catch (error) {
      console.error("[generateTextTraced] Failed to emit start event:", error);
    }
  }

  try {
    // Call the actual generateText
    const result = await generateText(generateOptions as Parameters<typeof generateText>[0]);

    const durationMs = Date.now() - startTime;

    // Emit completion event if in run context
    if (ctx) {
      try {
        await appendRunEvent({
          runId: ctx.runId,
          eventType: "llm_request_completed",
          level: "info",
          durationMs,
          pipelineName: ctx.pipelineName,
          stepName,
          llmOperation,
          promptVersionId,
          data: {
            model: modelName,
            finishReason: result.finishReason,
            usage: result.usage ? {
              inputTokens: result.usage.inputTokens,
              outputTokens: result.usage.outputTokens,
              totalTokens: result.usage.totalTokens,
            } : undefined,
            responseLength: result.text?.length,
            toolCallCount: result.steps?.reduce((acc, s) => acc + (s.toolCalls?.length || 0), 0) || 0,
          },
        });
      } catch (error) {
        console.error("[generateTextTraced] Failed to emit completion event:", error);
      }
    }

    return result;

  } catch (error) {
    const durationMs = Date.now() - startTime;

    // Emit failure event if in run context
    if (ctx) {
      try {
        await appendRunEvent({
          runId: ctx.runId,
          eventType: "llm_request_failed",
          level: "error",
          durationMs,
          pipelineName: ctx.pipelineName,
          stepName,
          llmOperation,
          promptVersionId,
          data: {
            model: modelName,
            error: error instanceof Error ? error.message : String(error),
            errorType: error instanceof Error ? error.name : "unknown",
          },
        });
      } catch (logError) {
        console.error("[generateTextTraced] Failed to emit error event:", logError);
      }
    }

    throw error;
  }
}

