import { NodeSDK } from '@opentelemetry/sdk-node';
import { LangfuseSpanProcessor } from '@langfuse/otel';
import { propagateAttributes, startObservation } from '@langfuse/tracing';
import type { LangfuseGeneration, LangfuseSpan } from '@langfuse/tracing';
import { Tracer } from '../llm/Tracer';
import type {
  GenerationHandle,
  GenerationInput,
  GenerationOutput,
  SessionInfo,
  TraceHandle,
} from '../llm/Tracer';

interface OpenAiUsageDetails {
  prompt_tokens?: number;
  completion_tokens?: number;
  total_tokens?: number;
  prompt_tokens_details?: { cached_tokens: number };
}

/**
 * Langfuse recognizes usage in either a flat generic schema (input/output/
 * cache_read_input_tokens/...) or OpenAI's native nested schema
 * (prompt_tokens/completion_tokens/prompt_tokens_details.cached_tokens/...)
 * - but NOT a mix of the two ("do not mix formats" per Langfuse's docs).
 * Mixing them previously caused the whole generation.update() payload
 * (including output) to land wrong, so this stays entirely within the
 * OpenAI-native schema.
 *
 * The installed @langfuse/tracing SDK's usageDetails type is
 * `Record<string, number> | OpenAiUsage` (a 3-field camelCase shape with
 * no _details nesting) even though the wire schema (@langfuse/core's
 * OpenAiCompletionUsageSchema) supports this nested form - a gap in the
 * SDK's public types, not a real runtime restriction, hence the cast.
 */
function toUsageDetails(usage: GenerationOutput['usage']): OpenAiUsageDetails | undefined {
  if (!usage) return undefined;
  const details: OpenAiUsageDetails = {};
  if (usage.promptTokens !== undefined) details.prompt_tokens = usage.promptTokens;
  if (usage.completionTokens !== undefined) details.completion_tokens = usage.completionTokens;
  if (usage.totalTokens !== undefined) details.total_tokens = usage.totalTokens;
  if (usage.cachedTokens !== undefined) details.prompt_tokens_details = { cached_tokens: usage.cachedTokens };
  return Object.keys(details).length > 0 ? details : undefined;
}

export class LangfuseTracer extends Tracer {
  private readonly processor: LangfuseSpanProcessor;
  private readonly sdk: NodeSDK;

  constructor(opts: { publicKey: string; secretKey: string; baseUrl: string }) {
    super();
    this.processor = new LangfuseSpanProcessor({
      publicKey: opts.publicKey,
      secretKey: opts.secretKey,
      baseUrl: opts.baseUrl,
    });
    this.sdk = new NodeSDK({ spanProcessors: [this.processor] });
    this.sdk.start();
  }

  startTrace(name: string, metadata?: Record<string, unknown>): TraceHandle {
    const span: LangfuseSpan = startObservation(name, { metadata });
    return {
      id: span.traceId,
      startGeneration(input: GenerationInput): GenerationHandle {
        const generation: LangfuseGeneration = span.startObservation(
          input.name,
          { model: input.model, input: input.input, metadata: input.metadata },
          { asType: 'generation' }
        );
        return {
          end(output: GenerationOutput) {
            generation.update({
              output: output.output,
              usageDetails: toUsageDetails(output.usage) as Record<string, number> | undefined,
              costDetails: output.usage?.cost !== undefined ? { total: output.usage.cost } : undefined,
              metadata: output.metadata,
            });
            generation.end();
          },
        };
      },
    };
  }

  async flush(): Promise<void> {
    await this.processor.forceFlush();
  }

  runWithSession<T>(session: SessionInfo, fn: () => Promise<T>): Promise<T> {
    return propagateAttributes({ sessionId: session.sessionId, tags: session.tags }, fn);
  }

  async shutdown(): Promise<void> {
    await this.sdk.shutdown();
  }
}