/**
 * Integration tests for Illuma Observability SDK integration in agents.
 *
 * These tests verify that the ObservabilityCallbackHandler and IllumaSpanProcessor
 * correctly route LangChain/OTel events to trace payloads with proper structure,
 * parent-child relationships, and token usage metrics.
 *
 * Uses a mock client to capture enqueued events without hitting a real server.
 */
/* eslint-disable @typescript-eslint/no-explicit-any, @typescript-eslint/explicit-function-return-type, @typescript-eslint/no-unused-vars, @typescript-eslint/no-require-imports, import/no-unresolved */

import { ObservabilityCallbackHandler } from '@illuma-ai/observability-langchain';
import type { Serialized } from '@langchain/core/load/serializable';
import type { LLMResult } from '@langchain/core/outputs';
import type { Document } from '@langchain/core/documents';

/** Minimal IngestionEvent type for our mock (avoids importing core) */
interface IngestionEvent {
  id: string;
  type: string;
  timestamp: string;
  body: Record<string, unknown>;
}

// ---------------------------------------------------------------------------
// Mock Client - captures enqueued events for assertions
// ---------------------------------------------------------------------------

class MockObservabilityClient {
  events: IngestionEvent[] = [];
  flushed = false;
  shutdownCalled = false;

  enqueue(event: IngestionEvent): void {
    this.events.push(event);
  }

  async flush(): Promise<void> {
    this.flushed = true;
  }

  async shutdown(): Promise<void> {
    this.shutdownCalled = true;
    await this.flush();
  }

  /** Get events filtered by type */
  getByType(type: string): IngestionEvent[] {
    return this.events.filter((e) => e.type === type);
  }

  /** Get all event types in order */
  getEventTypes(): string[] {
    return this.events.map((e) => e.type);
  }

  /** Pretty-print all events for debugging */
  printEvents(): void {
    for (const event of this.events) {
      const body = event.body as Record<string, unknown>;
      const usage = body.usage as Record<string, unknown> | undefined;
      console.log(
        `  ${event.type} | name=${body.name ?? '—'} | id=${body.id ?? '—'} | traceId=${body.traceId ?? '—'} | parentObsId=${body.parentObservationId ?? '—'}` +
          (usage
            ? ` | usage={prompt:${usage.promptTokens ?? '—'},completion:${usage.completionTokens ?? '—'},total:${usage.totalTokens ?? '—'}}`
            : '') +
          (body.model ? ` | model=${body.model}` : '') +
          (body.provider ? ` | provider=${body.provider}` : '') +
          (body.level === 'ERROR' ? ` | ERROR: ${body.statusMessage}` : '')
      );
    }
  }
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

const serialized = (name: string): Serialized => ({
  lc: 1,
  type: 'not_implemented',
  id: ['langchain', name],
});

function makeUUID(): string {
  return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => {
    const r = (Math.random() * 16) | 0;
    const v = c === 'x' ? r : (r & 0x3) | 0x8;
    return v.toString(16);
  });
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

describe('ObservabilityCallbackHandler integration', () => {
  let client: MockObservabilityClient;
  let handler: ObservabilityCallbackHandler;

  beforeEach(() => {
    client = new MockObservabilityClient();
    handler = new ObservabilityCallbackHandler({
      client,
      traceName: 'test-agent-run',
      userId: 'user-42',
      sessionId: 'session-abc',
      tags: ['integration-test'],
      metadata: { messageId: 'msg-001' },
      environment: 'test',
      debug: false,
    });
  });

  // -----------------------------------------------------------------------
  // 1. Root trace creation
  // -----------------------------------------------------------------------

  describe('root trace creation', () => {
    it('should create a trace on the first chain start', async () => {
      const chainRunId = makeUUID();
      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'Hello agent' },
        chainRunId
      );

      const traceEvents = client.getByType('trace-create');
      expect(traceEvents).toHaveLength(1);

      const trace = traceEvents[0].body as Record<string, unknown>;
      expect(trace.name).toBe('test-agent-run');
      expect(trace.userId).toBe('user-42');
      expect(trace.sessionId).toBe('session-abc');
      expect(trace.tags).toEqual(['integration-test']);
      expect(trace.metadata).toEqual({ messageId: 'msg-001' });
      expect(trace.environment).toBe('test');
      expect(trace.input).toEqual({ input: 'Hello agent' });
    });

    it('should not create a second trace for nested chains', async () => {
      const outerRunId = makeUUID();
      const innerRunId = makeUUID();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'outer' },
        outerRunId
      );
      await handler.handleChainStart(
        serialized('ChatPromptTemplate'),
        { input: 'inner' },
        innerRunId,
        outerRunId
      );

      const traceEvents = client.getByType('trace-create');
      expect(traceEvents).toHaveLength(1); // Only one trace, not two
    });

    it('should return the trace ID after creation', async () => {
      expect(handler.getTraceId()).toBeNull();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        makeUUID()
      );

      expect(handler.getTraceId()).toBeTruthy();
      expect(typeof handler.getTraceId()).toBe('string');
    });
  });

  // -----------------------------------------------------------------------
  // 2. LLM generation events
  // -----------------------------------------------------------------------

  describe('LLM generation tracing', () => {
    it('should create generation-create with input prompts', async () => {
      const chainRunId = makeUUID();
      const llmRunId = makeUUID();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        chainRunId
      );

      await handler.handleLLMStart(
        serialized('ChatAnthropic'),
        ['You are a helpful assistant.\n\nHuman: Hello'],
        llmRunId,
        chainRunId,
        undefined,
        undefined,
        { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' }
      );

      const genEvents = client.getByType('generation-create');
      expect(genEvents).toHaveLength(1);

      const gen = genEvents[0].body as Record<string, unknown>;
      expect(gen.name).toBe('ChatAnthropic');
      expect(gen.model).toBe('claude-sonnet-4-20250514');
      expect(gen.provider).toBe('anthropic');
      expect(gen.input).toEqual([
        'You are a helpful assistant.\n\nHuman: Hello',
      ]);
      expect(gen.traceId).toBe(handler.getTraceId());

      // Parent should be the chain
      const chainEvents = client.getByType('chain-create');
      const chainObsId = (chainEvents[0].body as Record<string, unknown>).id;
      expect(gen.parentObservationId).toBe(chainObsId);
    });

    it('should create generation-update with token usage on LLM end', async () => {
      const chainRunId = makeUUID();
      const llmRunId = makeUUID();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleLLMStart(
        serialized('ChatOpenAI'),
        ['Hello'],
        llmRunId,
        chainRunId
      );

      const llmResult: LLMResult = {
        generations: [
          [{ text: 'Hi there! How can I help?', generationInfo: {} }],
        ],
        llmOutput: {
          tokenUsage: {
            promptTokens: 150,
            completionTokens: 42,
            totalTokens: 192,
          },
          modelName: 'gpt-4o',
          provider: 'openai',
        },
      };

      await handler.handleLLMEnd(llmResult, llmRunId);

      const genUpdates = client.getByType('generation-update');
      expect(genUpdates).toHaveLength(1);

      const update = genUpdates[0].body as Record<string, unknown>;
      expect(update.output).toBe('Hi there! How can I help?');
      expect(update.model).toBe('gpt-4o');
      expect(update.provider).toBe('openai');
      expect(update.endTime).toBeDefined();

      const usage = update.usage as Record<string, unknown>;
      expect(usage.promptTokens).toBe(150);
      expect(usage.completionTokens).toBe(42);
      expect(usage.totalTokens).toBe(192);
    });

    it('should handle LLM errors with ERROR level', async () => {
      const chainRunId = makeUUID();
      const llmRunId = makeUUID();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleLLMStart(
        serialized('ChatAnthropic'),
        ['Hello'],
        llmRunId,
        chainRunId
      );

      await handler.handleLLMError(new Error('Rate limit exceeded'), llmRunId);

      const genUpdates = client.getByType('generation-update');
      expect(genUpdates).toHaveLength(1);

      const update = genUpdates[0].body as Record<string, unknown>;
      expect(update.level).toBe('ERROR');
      expect(update.statusMessage).toBe('Rate limit exceeded');
      expect(update.endTime).toBeDefined();
    });
  });

  // -----------------------------------------------------------------------
  // 3. Chain events
  // -----------------------------------------------------------------------

  describe('chain tracing', () => {
    it('should create chain-create events with input', async () => {
      const chainRunId = makeUUID();
      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { query: 'What is AI?' },
        chainRunId
      );

      const chainEvents = client.getByType('chain-create');
      expect(chainEvents).toHaveLength(1);

      const chain = chainEvents[0].body as Record<string, unknown>;
      expect(chain.name).toBe('RunnableSequence');
      expect(chain.input).toEqual({ query: 'What is AI?' });
      expect(chain.traceId).toBe(handler.getTraceId());
    });

    it('should create span-update on chain end with output', async () => {
      const chainRunId = makeUUID();
      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        chainRunId
      );

      await handler.handleChainEnd(
        { output: 'AI is artificial intelligence.' },
        chainRunId
      );

      const spanUpdates = client.getByType('span-update');
      expect(spanUpdates).toHaveLength(1);

      const update = spanUpdates[0].body as Record<string, unknown>;
      expect(update.output).toEqual({
        output: 'AI is artificial intelligence.',
      });
      expect(update.endTime).toBeDefined();
    });

    it('should handle chain error with ERROR level', async () => {
      const chainRunId = makeUUID();
      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        chainRunId
      );

      await handler.handleChainError(
        new Error('Chain failed: missing tool'),
        chainRunId
      );

      const spanUpdates = client.getByType('span-update');
      const errorUpdate = spanUpdates.find(
        (e) => (e.body as Record<string, unknown>).level === 'ERROR'
      );
      expect(errorUpdate).toBeDefined();
      expect((errorUpdate!.body as Record<string, unknown>).statusMessage).toBe(
        'Chain failed: missing tool'
      );
    });
  });

  // -----------------------------------------------------------------------
  // 4. Tool events
  // -----------------------------------------------------------------------

  describe('tool tracing', () => {
    it('should create tool-create events', async () => {
      const chainRunId = makeUUID();
      const toolRunId = makeUUID();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleToolStart(
        serialized('web_search'),
        '{"query": "latest AI news"}',
        toolRunId,
        chainRunId
      );

      const toolEvents = client.getByType('tool-create');
      expect(toolEvents).toHaveLength(1);

      const tool = toolEvents[0].body as Record<string, unknown>;
      expect(tool.name).toBe('Tool: web_search');
      expect(tool.input).toBe('{"query": "latest AI news"}');
      expect(tool.traceId).toBe(handler.getTraceId());
    });

    it('should create span-update on tool end', async () => {
      const chainRunId = makeUUID();
      const toolRunId = makeUUID();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleToolStart(
        serialized('calculator'),
        '{"expression": "2+2"}',
        toolRunId,
        chainRunId
      );

      await handler.handleToolEnd('4', toolRunId);

      const spanUpdates = client.getByType('span-update');
      expect(spanUpdates).toHaveLength(1);
      expect((spanUpdates[0].body as Record<string, unknown>).output).toBe('4');
    });

    it('should handle tool error', async () => {
      const chainRunId = makeUUID();
      const toolRunId = makeUUID();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleToolStart(
        serialized('api_call'),
        '{"url": "https://example.com"}',
        toolRunId,
        chainRunId
      );

      await handler.handleToolError(new Error('Connection timeout'), toolRunId);

      const spanUpdates = client.getByType('span-update');
      const errorUpdate = spanUpdates.find(
        (e) => (e.body as Record<string, unknown>).level === 'ERROR'
      );
      expect(errorUpdate).toBeDefined();
      expect((errorUpdate!.body as Record<string, unknown>).statusMessage).toBe(
        'Connection timeout'
      );
    });
  });

  // -----------------------------------------------------------------------
  // 5. Retriever events
  // -----------------------------------------------------------------------

  describe('retriever tracing', () => {
    it('should create retriever-create events', async () => {
      const chainRunId = makeUUID();
      const retrieverRunId = makeUUID();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleRetrieverStart(
        serialized('VectorStoreRetriever'),
        'How does photosynthesis work?',
        retrieverRunId,
        chainRunId
      );

      const retrieverEvents = client.getByType('retriever-create');
      expect(retrieverEvents).toHaveLength(1);

      const retriever = retrieverEvents[0].body as Record<string, unknown>;
      expect(retriever.name).toBe('Retriever: VectorStoreRetriever');
      expect(retriever.input).toBe('How does photosynthesis work?');
    });

    it('should produce span-update with documents on retriever end', async () => {
      const chainRunId = makeUUID();
      const retrieverRunId = makeUUID();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleRetrieverStart(
        serialized('VectorStoreRetriever'),
        'photosynthesis',
        retrieverRunId,
        chainRunId
      );

      const documents: Document[] = [
        {
          pageContent: 'Photosynthesis is the process...',
          metadata: { source: 'wiki', page: 1 },
        },
        {
          pageContent: 'Plants use chlorophyll...',
          metadata: { source: 'textbook', page: 42 },
        },
      ];

      await handler.handleRetrieverEnd(documents, retrieverRunId);

      const spanUpdates = client.getByType('span-update');
      expect(spanUpdates).toHaveLength(1);

      const update = spanUpdates[0].body as Record<string, unknown>;
      const output = update.output as Array<Record<string, unknown>>;
      expect(output).toHaveLength(2);
      expect(output[0].pageContent).toBe('Photosynthesis is the process...');
      expect(output[0].metadata).toEqual({ source: 'wiki', page: 1 });

      const meta = update.metadata as Record<string, unknown>;
      expect(meta.documentCount).toBe(2);
    });
  });

  // -----------------------------------------------------------------------
  // 6. Parent-child hierarchy
  // -----------------------------------------------------------------------

  describe('parent-child hierarchy', () => {
    it('should correctly link nested observations to parent', async () => {
      const outerChainId = makeUUID();
      const innerChainId = makeUUID();
      const llmRunId = makeUUID();
      const toolRunId = makeUUID();

      // Outer chain starts
      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'orchestrate' },
        outerChainId
      );

      // Inner chain nested under outer
      await handler.handleChainStart(
        serialized('ChatPromptTemplate'),
        { input: 'format prompt' },
        innerChainId,
        outerChainId
      );

      // LLM nested under inner chain
      await handler.handleLLMStart(
        serialized('ChatOpenAI'),
        ['formatted prompt'],
        llmRunId,
        innerChainId
      );

      // Tool nested under outer chain
      await handler.handleToolStart(
        serialized('calculator'),
        '2+2',
        toolRunId,
        outerChainId
      );

      // Verify hierarchy
      const chainCreates = client.getByType('chain-create');
      const genCreates = client.getByType('generation-create');
      const toolCreates = client.getByType('tool-create');

      // All share the same traceId
      const traceId = handler.getTraceId();
      expect((chainCreates[0].body as Record<string, unknown>).traceId).toBe(
        traceId
      );
      expect((chainCreates[1].body as Record<string, unknown>).traceId).toBe(
        traceId
      );
      expect((genCreates[0].body as Record<string, unknown>).traceId).toBe(
        traceId
      );
      expect((toolCreates[0].body as Record<string, unknown>).traceId).toBe(
        traceId
      );

      // Inner chain -> parent is outer chain's observationId
      const outerChainObsId = (chainCreates[0].body as Record<string, unknown>)
        .id;
      const innerChainObsId = (chainCreates[1].body as Record<string, unknown>)
        .id;
      expect(
        (chainCreates[1].body as Record<string, unknown>).parentObservationId
      ).toBe(outerChainObsId);

      // LLM -> parent is inner chain
      expect(
        (genCreates[0].body as Record<string, unknown>).parentObservationId
      ).toBe(innerChainObsId);

      // Tool -> parent is outer chain
      expect(
        (toolCreates[0].body as Record<string, unknown>).parentObservationId
      ).toBe(outerChainObsId);
    });
  });

  // -----------------------------------------------------------------------
  // 7. Full agent pipeline simulation
  // -----------------------------------------------------------------------

  describe('full agent pipeline simulation', () => {
    it('should produce correct event sequence for a RAG agent with tools', async () => {
      const chainRunId = makeUUID();
      const agentChainId = makeUUID();
      const retrieverRunId = makeUUID();
      const llmRunId1 = makeUUID();
      const toolRunId = makeUUID();
      const llmRunId2 = makeUUID();

      // 1. Root chain starts (the main graph)
      await handler.handleChainStart(
        serialized('CompiledStateGraph'),
        {
          messages: [
            {
              role: 'user',
              content: 'What papers did OpenAI publish in 2024?',
            },
          ],
        },
        chainRunId
      );

      // 2. Agent chain starts (nested)
      await handler.handleChainStart(
        serialized('AgentExecutor'),
        { input: 'What papers did OpenAI publish in 2024?' },
        agentChainId,
        chainRunId
      );

      // 3. First LLM call (decides to use tool)
      await handler.handleLLMStart(
        serialized('ChatAnthropic'),
        ['System: You are a research assistant.\nHuman: What papers...'],
        llmRunId1,
        agentChainId,
        undefined,
        undefined,
        { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' }
      );

      await handler.handleLLMEnd(
        {
          generations: [
            [
              {
                text: "I'll search for OpenAI papers from 2024.",
                generationInfo: {},
              },
            ],
          ],
          llmOutput: {
            tokenUsage: {
              promptTokens: 200,
              completionTokens: 30,
              totalTokens: 230,
            },
            modelName: 'claude-sonnet-4-20250514',
            provider: 'anthropic',
          },
        },
        llmRunId1
      );

      // 4. Retriever runs
      await handler.handleRetrieverStart(
        serialized('VectorStoreRetriever'),
        'OpenAI papers 2024',
        retrieverRunId,
        agentChainId
      );

      await handler.handleRetrieverEnd(
        [
          {
            pageContent: 'GPT-4o: omni-model for text, vision, audio...',
            metadata: { source: 'arxiv', year: 2024 },
          },
          {
            pageContent: 'Sora: text-to-video generation model...',
            metadata: { source: 'blog', year: 2024 },
          },
        ] as Document[],
        retrieverRunId
      );

      // 5. Tool call (web search)
      await handler.handleToolStart(
        serialized('web_search'),
        '{"query": "OpenAI 2024 research papers list"}',
        toolRunId,
        agentChainId
      );

      await handler.handleToolEnd(
        JSON.stringify([
          { title: 'GPT-4o Technical Report', url: 'https://arxiv.org/...' },
          { title: 'Sora', url: 'https://openai.com/sora' },
        ]),
        toolRunId
      );

      // 6. Second LLM call (final answer)
      await handler.handleLLMStart(
        serialized('ChatAnthropic'),
        ['System: ...\nContext: ...\nHuman: What papers...'],
        llmRunId2,
        agentChainId,
        undefined,
        undefined,
        { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' }
      );

      await handler.handleLLMEnd(
        {
          generations: [
            [
              {
                text: 'OpenAI published several notable papers in 2024 including GPT-4o and Sora...',
                generationInfo: {},
              },
            ],
          ],
          llmOutput: {
            tokenUsage: {
              promptTokens: 800,
              completionTokens: 150,
              totalTokens: 950,
            },
            modelName: 'claude-sonnet-4-20250514',
            provider: 'anthropic',
          },
        },
        llmRunId2
      );

      // 7. Chains end
      await handler.handleChainEnd(
        { output: 'OpenAI published several notable papers...' },
        agentChainId
      );
      await handler.handleChainEnd(
        {
          messages: [
            {
              role: 'assistant',
              content: 'OpenAI published several notable papers...',
            },
          ],
        },
        chainRunId
      );

      // Flush
      await handler.flushAsync();
      expect(client.flushed).toBe(true);

      // --- Assertions ---

      // Uncomment to debug: client.printEvents();

      // Verify event counts
      const eventTypes = client.getEventTypes();
      const typeCounts = eventTypes.reduce(
        (acc, t) => {
          acc[t] = (acc[t] || 0) + 1;
          return acc;
        },
        {} as Record<string, number>
      );

      expect(typeCounts['trace-create']).toBe(1); // Root trace
      expect(typeCounts['chain-create']).toBe(1); // Outer chain only
      expect(typeCounts['agent-create']).toBe(1); // AgentExecutor detected as agent!
      expect(typeCounts['generation-create']).toBe(2); // Two LLM calls
      expect(typeCounts['generation-update']).toBe(2); // Two LLM ends
      expect(typeCounts['retriever-create']).toBe(1); // One retriever
      expect(typeCounts['tool-create']).toBe(1); // One tool
      expect(typeCounts['span-update']).toBe(4); // retriever end + tool end + 2 chain ends

      // Verify all events share the same traceId
      const traceId = handler.getTraceId()!;
      for (const event of client.events) {
        const body = event.body as Record<string, unknown>;
        if (body.traceId) {
          expect(body.traceId).toBe(traceId);
        }
      }

      // Verify token usage in generation-update events
      const genUpdates = client.getByType('generation-update');

      // First LLM call usage
      const firstLLMUpdate = genUpdates[0].body as Record<string, unknown>;
      const firstUsage = firstLLMUpdate.usage as Record<string, unknown>;
      expect(firstUsage.promptTokens).toBe(200);
      expect(firstUsage.completionTokens).toBe(30);
      expect(firstUsage.totalTokens).toBe(230);
      expect(firstLLMUpdate.model).toBe('claude-sonnet-4-20250514');
      expect(firstLLMUpdate.provider).toBe('anthropic');

      // Second LLM call usage
      const secondLLMUpdate = genUpdates[1].body as Record<string, unknown>;
      const secondUsage = secondLLMUpdate.usage as Record<string, unknown>;
      expect(secondUsage.promptTokens).toBe(800);
      expect(secondUsage.completionTokens).toBe(150);
      expect(secondUsage.totalTokens).toBe(950);

      // Verify total token computation
      const totalPrompt = 200 + 800;
      const totalCompletion = 30 + 150;
      const totalTokens = totalPrompt + totalCompletion;
      expect(totalPrompt).toBe(1000);
      expect(totalCompletion).toBe(180);
      expect(totalTokens).toBe(1180);

      // Total: 1180 tokens (1000 prompt + 180 completion) across 2 LLM calls
    });
  });

  // -----------------------------------------------------------------------
  // 8. Token usage edge cases
  // -----------------------------------------------------------------------

  describe('token usage edge cases', () => {
    it('should handle different token usage field names (prompt_tokens vs promptTokens)', async () => {
      const chainRunId = makeUUID();
      const llmRunId = makeUUID();

      await handler.handleChainStart(
        serialized('Chain'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleLLMStart(
        serialized('ChatOpenAI'),
        ['test'],
        llmRunId,
        chainRunId
      );

      // OpenAI style: uses prompt_tokens (snake_case)
      await handler.handleLLMEnd(
        {
          generations: [[{ text: 'response', generationInfo: {} }]],
          llmOutput: {
            tokenUsage: {
              prompt_tokens: 100,
              completion_tokens: 50,
              total_tokens: 150,
            },
          },
        },
        llmRunId
      );

      const updates = client.getByType('generation-update');
      const usage = (updates[0].body as Record<string, unknown>)
        .usage as Record<string, unknown>;

      // The handler should normalize both field naming conventions
      expect(usage.promptTokens ?? usage.prompt_tokens).toBeTruthy();
    });

    it('should handle missing token usage gracefully', async () => {
      const chainRunId = makeUUID();
      const llmRunId = makeUUID();

      await handler.handleChainStart(
        serialized('Chain'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleLLMStart(
        serialized('ChatOpenAI'),
        ['test'],
        llmRunId,
        chainRunId
      );

      // No llmOutput at all
      await handler.handleLLMEnd(
        { generations: [[{ text: 'response', generationInfo: {} }]] },
        llmRunId
      );

      const updates = client.getByType('generation-update');
      expect(updates).toHaveLength(1);
      // Should not crash, usage should be present but with undefined values
      const usage = (updates[0].body as Record<string, unknown>)
        .usage as Record<string, unknown>;
      expect(usage).toBeDefined();
    });
  });

  // -----------------------------------------------------------------------
  // 9. Flush and shutdown
  // -----------------------------------------------------------------------

  describe('flush and shutdown', () => {
    it('flushAsync should call client.flush()', async () => {
      await handler.handleChainStart(
        serialized('Chain'),
        { input: 'test' },
        makeUUID()
      );
      await handler.flushAsync();
      expect(client.flushed).toBe(true);
    });

    it('shutdown should call client.flush() (not shutdown) for external client', async () => {
      await handler.handleChainStart(
        serialized('Chain'),
        { input: 'test' },
        makeUUID()
      );
      await handler.shutdown();
      // Since we provided an external client, it should only flush, not shutdown
      expect(client.flushed).toBe(true);
      expect(client.shutdownCalled).toBe(false);
    });
  });

  // -----------------------------------------------------------------------
  // 10. Agent detection
  // -----------------------------------------------------------------------

  describe('agent detection', () => {
    it('should emit agent-create for chains with langgraph_node agent= metadata', async () => {
      const outerChainId = makeUUID();
      const agentChainId = makeUUID();

      await handler.handleChainStart(
        serialized('CompiledStateGraph'),
        { messages: [{ role: 'user', content: 'Hello' }] },
        outerChainId
      );

      // Agent chain with langgraph_node metadata
      // Cast to any to pass tags/metadata params (ts-jest resolves different .d.ts for BaseCallbackHandler)
      await (handler as any).handleChainStart(
        serialized('RunnableSequence'),
        { input: 'agent task' },
        agentChainId,
        outerChainId,
        undefined, // tags
        { langgraph_node: 'agent=research-agent' } // metadata
      );

      const agentEvents = client.getByType('agent-create');
      expect(agentEvents).toHaveLength(1);

      const agentBody = agentEvents[0].body as Record<string, unknown>;
      expect(agentBody.name).toBe('Agent: research-agent');
      expect(agentBody.traceId).toBe(handler.getTraceId());

      const meta = agentBody.metadata as Record<string, unknown>;
      expect(meta.agentId).toBe('research-agent');
      expect(meta.langchainMetadata).toEqual({
        langgraph_node: 'agent=research-agent',
      });
    });

    it('should emit agent-create for AgentExecutor chain name', async () => {
      const agentRunId = makeUUID();

      await handler.handleChainStart(
        serialized('AgentExecutor'),
        { input: 'do something' },
        agentRunId
      );

      const agentEvents = client.getByType('agent-create');
      expect(agentEvents).toHaveLength(1);
      expect((agentEvents[0].body as Record<string, unknown>).name).toBe(
        'AgentExecutor'
      );

      // Should NOT produce a chain-create
      const chainEvents = client.getByType('chain-create');
      expect(chainEvents).toHaveLength(0);
    });

    it('should emit agent-create for chain names ending in Agent', async () => {
      const agentRunId = makeUUID();

      await handler.handleChainStart(
        serialized('ReactAgent'),
        { input: 'think and act' },
        agentRunId
      );

      const agentEvents = client.getByType('agent-create');
      expect(agentEvents).toHaveLength(1);
      expect((agentEvents[0].body as Record<string, unknown>).name).toBe(
        'ReactAgent'
      );
    });

    it('should NOT emit agent-create for regular chains', async () => {
      const chainRunId = makeUUID();

      await handler.handleChainStart(
        serialized('RunnableSequence'),
        { input: 'regular chain' },
        chainRunId
      );

      const agentEvents = client.getByType('agent-create');
      expect(agentEvents).toHaveLength(0);

      const chainEvents = client.getByType('chain-create');
      expect(chainEvents).toHaveLength(1);
    });

    it('should correctly end agent observations with span-update', async () => {
      const agentRunId = makeUUID();

      await (handler as any).handleChainStart(
        serialized('RunnableSequence'),
        { input: 'agent task' },
        agentRunId,
        undefined, // parentRunId
        undefined, // tags
        { langgraph_node: 'agent=planner' } // metadata
      );

      await handler.handleChainEnd({ output: 'agent completed' }, agentRunId);

      const agentEvents = client.getByType('agent-create');
      expect(agentEvents).toHaveLength(1);

      const spanUpdates = client.getByType('span-update');
      expect(spanUpdates).toHaveLength(1);

      const update = spanUpdates[0].body as Record<string, unknown>;
      expect(update.output).toEqual({ output: 'agent completed' });
      expect(update.endTime).toBeDefined();
      // The span-update should reference the same observation ID as agent-create
      expect(update.id).toBe(
        (agentEvents[0].body as Record<string, unknown>).id
      );
    });

    it('should handle agent error with ERROR level', async () => {
      const agentRunId = makeUUID();

      await handler.handleChainStart(
        serialized('AgentExecutor'),
        { input: 'fail task' },
        agentRunId
      );

      await handler.handleChainError(
        new Error('Agent failed: tool not found'),
        agentRunId
      );

      const agentEvents = client.getByType('agent-create');
      expect(agentEvents).toHaveLength(1);

      const spanUpdates = client.getByType('span-update');
      const errorUpdate = spanUpdates.find(
        (e) => (e.body as Record<string, unknown>).level === 'ERROR'
      );
      expect(errorUpdate).toBeDefined();
      expect((errorUpdate!.body as Record<string, unknown>).statusMessage).toBe(
        'Agent failed: tool not found'
      );
    });

    it('should integrate agent-create in full pipeline with correct counts', async () => {
      const rootChainId = makeUUID();
      const agentChainId = makeUUID();
      const llmRunId = makeUUID();

      // Root chain
      await handler.handleChainStart(
        serialized('CompiledStateGraph'),
        { messages: [{ role: 'user', content: 'Hello' }] },
        rootChainId
      );

      // Agent chain (detected via metadata)
      await (handler as any).handleChainStart(
        serialized('RunnableSequence'),
        { input: 'agent work' },
        agentChainId,
        rootChainId,
        undefined, // tags
        { langgraph_node: 'agent=qa-agent' } // metadata
      );

      // LLM inside agent
      await handler.handleLLMStart(
        serialized('ChatAnthropic'),
        ['prompt'],
        llmRunId,
        agentChainId,
        undefined,
        undefined,
        { ls_model_name: 'claude-sonnet-4-20250514', ls_provider: 'anthropic' }
      );

      await handler.handleLLMEnd(
        {
          generations: [[{ text: 'answer', generationInfo: {} }]],
          llmOutput: {
            tokenUsage: {
              promptTokens: 100,
              completionTokens: 20,
              totalTokens: 120,
            },
            modelName: 'claude-sonnet-4-20250514',
          },
        },
        llmRunId
      );

      await handler.handleChainEnd({ output: 'done' }, agentChainId);
      await handler.handleChainEnd({ output: 'done' }, rootChainId);

      const typeCounts = client.getEventTypes().reduce(
        (acc, t) => {
          acc[t] = (acc[t] || 0) + 1;
          return acc;
        },
        {} as Record<string, number>
      );

      expect(typeCounts['trace-create']).toBe(1);
      expect(typeCounts['chain-create']).toBe(1); // Only root chain
      expect(typeCounts['agent-create']).toBe(1); // Agent detected!
      expect(typeCounts['generation-create']).toBe(1);
      expect(typeCounts['generation-update']).toBe(1);
      expect(typeCounts['span-update']).toBe(2); // agent end + root chain end

      // LLM parent should be the agent observation
      const agentObsId = (
        client.getByType('agent-create')[0].body as Record<string, unknown>
      ).id;
      const genCreate = client.getByType('generation-create')[0].body as Record<
        string,
        unknown
      >;
      expect(genCreate.parentObservationId).toBe(agentObsId);
    });
  });

  // -----------------------------------------------------------------------
  // 11. Event structure validation
  // -----------------------------------------------------------------------

  describe('event structure validation', () => {
    it('every event should have id, type, timestamp, and body', async () => {
      const chainRunId = makeUUID();
      const llmRunId = makeUUID();

      await handler.handleChainStart(
        serialized('Chain'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleLLMStart(
        serialized('LLM'),
        ['test'],
        llmRunId,
        chainRunId
      );
      await handler.handleLLMEnd(
        {
          generations: [[{ text: 'ok', generationInfo: {} }]],
          llmOutput: { tokenUsage: {} },
        },
        llmRunId
      );
      await handler.handleChainEnd({ output: 'done' }, chainRunId);

      for (const event of client.events) {
        expect(event.id).toBeDefined();
        expect(typeof event.id).toBe('string');
        expect(event.type).toBeDefined();
        expect(event.timestamp).toMatch(/^\d{4}-\d{2}-\d{2}T/);
        expect(event.body).toBeDefined();
        expect(typeof event.body).toBe('object');
      }
    });

    it('observation events should have traceId and observationId', async () => {
      const chainRunId = makeUUID();
      await handler.handleChainStart(
        serialized('Chain'),
        { input: 'test' },
        chainRunId
      );
      await handler.handleChainEnd({ output: 'done' }, chainRunId);

      const nonTraceEvents = client.events.filter(
        (e) => e.type !== 'trace-create'
      );
      for (const event of nonTraceEvents) {
        const body = event.body as Record<string, unknown>;
        expect(body.traceId).toBe(handler.getTraceId());
        expect(body.id).toBeDefined();
      }
    });
  });
});

// ---------------------------------------------------------------------------
// IllumaSpanProcessor unit tests (OTel integration)
// ---------------------------------------------------------------------------

describe('IllumaSpanProcessor integration', () => {
  it('should be importable from @illuma-ai/observability-otel', async () => {
    const { IllumaSpanProcessor } = await import(
      '@illuma-ai/observability-otel'
    );
    expect(IllumaSpanProcessor).toBeDefined();
    expect(typeof IllumaSpanProcessor).toBe('function');
  });
});

// ---------------------------------------------------------------------------
// Core SDK features: sampling, PII masking, observation types, prompt/dataset
// ---------------------------------------------------------------------------

describe('Core SDK features', () => {
  // Use require to bypass ts-jest module resolution issues with dynamic imports

  const {
    ObservabilityCoreClient,
    TraceClient,
    SpanClient,
    IngestionEventType,
  } = require('@illuma-ai/observability-core');

  // Create a concrete subclass at runtime to avoid TS abstract class issues
  const TestClientClass = class extends (ObservabilityCoreClient as any) {
    protected async fetchWithRetry(): Promise<any> {
      return {
        status: 200,
        statusText: 'OK',
        ok: true,
        json: async () => ({}),
        text: async () => '',
      };
    }
  };

  function createTestClient(config: Record<string, unknown> = {}): {
    client: any;
    events: any[];
  } {
    const events: any[] = [];

    const client = new (TestClientClass as any)({
      publicKey: 'pk-test',
      secretKey: 'sk-test',
      baseUrl: 'http://localhost:9999',
      flushInterval: 0,
      ...config,
    });

    // Intercept enqueue to capture events
    const origEnqueue = client.enqueue.bind(client);
    client.enqueue = (event: any) => {
      origEnqueue(event);
      events.push(event);
    };

    return { client, events };
  }

  // -----------------------------------------------------------------------
  // 12. Sampling rate
  // -----------------------------------------------------------------------

  describe('sampling rate', () => {
    it('should trace everything when sampleRate=1.0 (default)', () => {
      const { client, events } = createTestClient({ sampleRate: 1.0 });

      for (let i = 0; i < 20; i++) {
        client.trace({ name: `trace-${i}` });
      }

      // All 20 traces should be created (each produces a trace-create event queued)
      const traceEvents = events.filter((e: any) => e.type === 'trace-create');
      expect(traceEvents.length).toBe(20);
    });

    it('should trace nothing when sampleRate=0.0', () => {
      const { client, events } = createTestClient({ sampleRate: 0.0 });

      for (let i = 0; i < 20; i++) {
        const trace = client.trace({ name: `trace-${i}` });
        // Even child events should be silently dropped
        trace.generation({ name: 'gen', model: 'gpt-4o' });
        trace.span({ name: 'span' });
      }

      // No events should be enqueued at all
      expect(events.length).toBe(0);
    });

    it('should sample approximately the right percentage', () => {
      // Use a fixed seed approach: run many traces with 50% sample rate
      const { client, events } = createTestClient({ sampleRate: 0.5 });

      const iterations = 1000;
      for (let i = 0; i < iterations; i++) {
        client.trace({ name: `trace-${i}` });
      }

      const traceEvents = events.filter((e: any) => e.type === 'trace-create');
      // With 1000 iterations at 50%, expect roughly 400-600 traces
      expect(traceEvents.length).toBeGreaterThan(300);
      expect(traceEvents.length).toBeLessThan(700);
    });

    it('should include all children once a trace is sampled in', () => {
      const { client, events } = createTestClient({ sampleRate: 1.0 });

      const trace = client.trace({ name: 'sampled-trace' });
      trace.span({ name: 'child-span' });
      trace.generation({ name: 'child-gen', model: 'gpt-4o' });
      trace.agent({ name: 'child-agent' });
      trace.tool({ name: 'child-tool' });

      // 1 trace + 4 children = 5 events
      expect(events.length).toBe(5);
    });

    it('should clamp sampleRate to [0, 1] range', () => {
      // sampleRate > 1 should be clamped to 1 (trace everything)
      const { events: events1 } = createTestClient({ sampleRate: 5.0 });
      // sampleRate < 0 should be clamped to 0 (trace nothing)
      const { client: client2, events: events2 } = createTestClient({
        sampleRate: -1,
      });

      client2.trace({ name: 'should-not-appear' });
      expect(events2.length).toBe(0);
    });
  });

  // -----------------------------------------------------------------------
  // 13. PII masking
  // -----------------------------------------------------------------------

  describe('PII masking', () => {
    it('should apply mask function to event bodies', () => {
      const maskFunction = (body: Record<string, unknown>) => {
        const masked = { ...body };
        if (typeof masked.input === 'string') {
          masked.input = masked.input.replace(
            /[\w.-]+@[\w.-]+\.\w+/g,
            '[EMAIL]'
          );
        }
        return masked;
      };

      const { client } = createTestClient({ maskFunction });

      const trace = client.trace({ name: 'pii-test' });
      trace.span({
        name: 'user-query',
        input: 'My email is john@example.com and I need help' as any,
      });

      // Check the queued events - the span should have masked input
      const queue = client.getQueue();
      const spanEvent = [...queue].find((e: any) => e.type === 'span-create');
      if (spanEvent) {
        expect((spanEvent.body as any).input).toBe(
          'My email is [EMAIL] and I need help'
        );
        expect((spanEvent.body as any).input).not.toContain('john@example.com');
      }
    });

    it('should mask multiple PII patterns', () => {
      const maskFunction = (body: Record<string, unknown>) => {
        const masked = { ...body };
        const str = JSON.stringify(masked);
        const redacted = str
          .replace(/[\w.-]+@[\w.-]+\.\w+/g, '[EMAIL]')
          .replace(/\b\d{3}[-.]?\d{3}[-.]?\d{4}\b/g, '[PHONE]')
          .replace(/\b\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}\b/g, '[CARD]');
        return JSON.parse(redacted);
      };

      const { client } = createTestClient({ maskFunction });

      const trace = client.trace({
        name: 'multi-pii',
        metadata: { userEmail: 'test@corp.com', phone: '555-123-4567' },
      });

      const queue = client.getQueue();
      const traceEvent = [...queue].find((e: any) => e.type === 'trace-create');
      if (traceEvent) {
        const meta = (traceEvent.body as any).metadata;
        expect(meta.userEmail).toBe('[EMAIL]');
        expect(meta.phone).toBe('[PHONE]');
      }
    });

    it('should gracefully handle mask function errors', () => {
      const maskFunction = () => {
        throw new Error('Mask function exploded');
      };

      const { client } = createTestClient({ maskFunction });

      // Should not throw — falls back to unmasked
      expect(() => {
        client.trace({ name: 'should-not-crash' });
      }).not.toThrow();

      const queue = client.getQueue();
      expect(queue.length).toBeGreaterThan(0);
    });

    it('should not modify events when no mask function is set', () => {
      const { client } = createTestClient();

      client.trace({
        name: 'no-mask',
        metadata: { email: 'visible@example.com' },
      });

      const queue = client.getQueue();
      const traceEvent = [...queue].find((e: any) => e.type === 'trace-create');
      expect((traceEvent!.body as any).metadata.email).toBe(
        'visible@example.com'
      );
    });
  });

  // -----------------------------------------------------------------------
  // 14. All observation types (TraceClient + SpanClient)
  // -----------------------------------------------------------------------

  describe('observation types completeness', () => {
    it('TraceClient should have all observation type methods', () => {
      const enqueue = jest.fn();
      const trace = new TraceClient('trace-123', enqueue);

      // All typed observation methods should exist and return SpanClient or GenerationClient
      expect(typeof trace.span).toBe('function');
      expect(typeof trace.agent).toBe('function');
      expect(typeof trace.tool).toBe('function');
      expect(typeof trace.chain).toBe('function');
      expect(typeof trace.retriever).toBe('function');
      expect(typeof trace.guardrail).toBe('function');
      expect(typeof trace.evaluator).toBe('function');
      expect(typeof trace.embedding).toBe('function');
      expect(typeof trace.generation).toBe('function');
      expect(typeof trace.event).toBe('function');
      expect(typeof trace.score).toBe('function');
    });

    it('SpanClient should have all observation type methods', () => {
      const enqueue = jest.fn();
      const span = new SpanClient('span-123', 'trace-123', enqueue);

      expect(typeof span.span).toBe('function');
      expect(typeof span.agent).toBe('function');
      expect(typeof span.tool).toBe('function');
      expect(typeof span.chain).toBe('function');
      expect(typeof span.retriever).toBe('function');
      expect(typeof span.guardrail).toBe('function');
      expect(typeof span.evaluator).toBe('function');
      expect(typeof span.embedding).toBe('function');
      expect(typeof span.generation).toBe('function');
      expect(typeof span.event).toBe('function');
      expect(typeof span.score).toBe('function');
    });

    it('should emit correct event types for each observation method on TraceClient', () => {
      const events: any[] = [];
      const enqueue = (event: any) => events.push(event);
      const trace = new TraceClient('trace-456', enqueue);

      trace.span({ name: 'span' });
      trace.agent({ name: 'agent' });
      trace.tool({ name: 'tool' });
      trace.chain({ name: 'chain' });
      trace.retriever({ name: 'retriever' });
      trace.guardrail({ name: 'guardrail' });
      trace.evaluator({ name: 'evaluator' });
      trace.embedding({ name: 'embedding' });
      trace.generation({ name: 'generation' });
      trace.event({ name: 'event' });

      const types = events.map((e) => e.type);
      expect(types).toContain('span-create');
      expect(types).toContain('agent-create');
      expect(types).toContain('tool-create');
      expect(types).toContain('chain-create');
      expect(types).toContain('retriever-create');
      expect(types).toContain('guardrail-create');
      expect(types).toContain('evaluator-create');
      expect(types).toContain('embedding-create');
      expect(types).toContain('generation-create');
      expect(types).toContain('event-create');
    });

    it('should emit correct event types for each observation method on SpanClient', () => {
      const events: any[] = [];
      const enqueue = (event: any) => events.push(event);
      const span = new SpanClient('span-789', 'trace-789', enqueue);

      span.span({ name: 'child-span' });
      span.agent({ name: 'child-agent' });
      span.tool({ name: 'child-tool' });
      span.chain({ name: 'child-chain' });
      span.retriever({ name: 'child-retriever' });
      span.guardrail({ name: 'child-guardrail' });
      span.evaluator({ name: 'child-evaluator' });
      span.embedding({ name: 'child-embedding' });
      span.generation({ name: 'child-generation' });
      span.event({ name: 'child-event' });

      const types = events.map((e) => e.type);
      expect(types).toContain('span-create');
      expect(types).toContain('agent-create');
      expect(types).toContain('tool-create');
      expect(types).toContain('chain-create');
      expect(types).toContain('retriever-create');
      expect(types).toContain('guardrail-create');
      expect(types).toContain('evaluator-create');
      expect(types).toContain('embedding-create');
      expect(types).toContain('generation-create');
      expect(types).toContain('event-create');
    });

    it('SpanClient children should set parentObservationId automatically', () => {
      const events: any[] = [];
      const enqueue = (event: any) => events.push(event);
      const span = new SpanClient('parent-span', 'trace-abc', enqueue);

      span.guardrail({ name: 'guardrail-check' });
      span.evaluator({ name: 'eval-run' });
      span.embedding({ name: 'embed-op' });

      for (const event of events) {
        expect(event.body.parentObservationId).toBe('parent-span');
        expect(event.body.traceId).toBe('trace-abc');
      }
    });
  });

  // -----------------------------------------------------------------------
  // 15. Prompt & Dataset SDK methods
  // -----------------------------------------------------------------------

  describe('prompt and dataset SDK methods', () => {
    it('core client should expose getPrompt method', () => {
      const { client } = createTestClient();
      expect(typeof client.getPrompt).toBe('function');
    });

    it('core client should expose createPrompt method', () => {
      const { client } = createTestClient();
      expect(typeof client.createPrompt).toBe('function');
    });

    it('core client should expose getDataset method', () => {
      const { client } = createTestClient();
      expect(typeof client.getDataset).toBe('function');
    });

    it('core client should expose createDataset method', () => {
      const { client } = createTestClient();
      expect(typeof client.createDataset).toBe('function');
    });

    it('core client should expose createDatasetItem method', () => {
      const { client } = createTestClient();
      expect(typeof client.createDatasetItem).toBe('function');
    });

    it('getPrompt should return parsed response from server', async () => {
      const { client } = createTestClient();
      // Mock fetchWithRetry returns { ok: true, json: async () => ({}) }
      const result = await client.getPrompt({ name: 'my-prompt' });
      expect(result).toBeDefined();
      expect(result).toEqual({});
    });

    it('getDataset should return parsed response from server', async () => {
      const { client } = createTestClient();
      const result = await client.getDataset({ name: 'my-dataset' });
      expect(result).toBeDefined();
      expect(result).toEqual({});
    });

    it('prompt/dataset methods should return null when disabled', async () => {
      const { client } = createTestClient({ enabled: false });
      expect(await client.getPrompt({ name: 'test' })).toBeNull();
      expect(
        await client.createPrompt({ name: 'test', prompt: 'hello' })
      ).toBeNull();
      expect(await client.getDataset({ name: 'test' })).toBeNull();
      expect(await client.createDataset({ name: 'test' })).toBeNull();
      expect(
        await client.createDatasetItem({ datasetName: 'test', input: 'x' })
      ).toBeNull();
    });
  });

  // -----------------------------------------------------------------------
  // 16. Node SDK: sampling + PII masking propagation
  // -----------------------------------------------------------------------

  describe('node SDK config propagation', () => {
    const { Observability } = require('@illuma-ai/observability-node');

    it('Observability class should accept sampleRate option', async () => {
      expect(Observability).toBeDefined();

      const obs = new Observability({
        publicKey: 'pk-test',
        secretKey: 'sk-test',
        baseUrl: 'http://localhost:9999',
        sampleRate: 0.5,
        flushInterval: 0,
      });

      expect(obs).toBeDefined();
      await obs.shutdown();
    });

    it('Observability class should accept maskFunction option', async () => {
      const obs = new Observability({
        publicKey: 'pk-test',
        secretKey: 'sk-test',
        baseUrl: 'http://localhost:9999',
        maskFunction: (body: Record<string, unknown>) => ({
          ...body,
          input: '[REDACTED]',
        }),
        flushInterval: 0,
      });

      expect(obs).toBeDefined();
      await obs.shutdown();
    });
  });
});

// ---------------------------------------------------------------------------
// 17. Guardrail tracing via ObservabilityCallbackHandler
// ---------------------------------------------------------------------------

/** Re-use MockObservabilityClient shape for guardrail tests (same file, top-level scope) */
class GuardrailMockClient {
  events: {
    id: string;
    type: string;
    timestamp: string;
    body: Record<string, unknown>;
  }[] = [];
  flushed = false;
  shutdownCalled = false;
  enqueue(event: {
    id: string;
    type: string;
    timestamp: string;
    body: Record<string, unknown>;
  }): void {
    this.events.push(event);
  }
  async flush(): Promise<void> {
    this.flushed = true;
  }
  async shutdown(): Promise<void> {
    this.shutdownCalled = true;
  }
  getByType(type: string) {
    return this.events.filter((e) => e.type === type);
  }
  getEventTypes() {
    return this.events.map((e) => e.type);
  }
  printEvents(): void {
    for (const event of this.events) {
      const body = event.body;
      const usage = body.usage as Record<string, unknown> | undefined;
      console.log(
        `  ${event.type} | name=${body.name ?? '—'} | id=${body.id ?? '—'} | traceId=${body.traceId ?? '—'} | parentObsId=${body.parentObservationId ?? '—'}` +
          (usage
            ? ` | usage={prompt:${usage.promptTokens ?? '—'},completion:${usage.completionTokens ?? '—'},total:${usage.totalTokens ?? '—'}}`
            : '') +
          (body.model ? ` | model=${body.model}` : '') +
          (body.level === 'ERROR' || body.level === 'WARNING'
            ? ` | ${body.level}: ${body.statusMessage}`
            : '')
      );
    }
  }
}

const grSerialized = (name: string): Serialized => ({
  lc: 1,
  type: 'not_implemented',
  id: ['langchain', name],
});

describe('Guardrail tracing', () => {
  let grClient: GuardrailMockClient;
  let grHandler: any; // Use any to bypass ts-jest type resolution for traceGuardrail

  beforeEach(() => {
    grClient = new GuardrailMockClient();
    grHandler = new ObservabilityCallbackHandler({
      client: grClient as any,
    });
  });

  it('outcome=passed should create guardrail-create + span-update with DEFAULT level', async () => {
    const chainRunId = makeUUID();
    await grHandler.handleChainStart(
      grSerialized('RunnableSequence'),
      { input: 'hello' },
      chainRunId
    );

    const result = grHandler.traceGuardrail({
      name: 'Output Moderation',
      guardrailId: 'bedrock-guardrail-123',
      guardrailVersion: '1',
      outcome: 'passed',
      actionApplied: false,
      action: 'NONE',
      reason: 'passed',
      source: 'OUTPUT',
      input: 'What is the weather?',
      output: 'The weather is sunny.',
      violations: [],
      assessments: [{ contentPolicy: { filters: [] } }],
    });

    expect(result).not.toBeNull();

    const guardrailEvents = grClient.getByType('guardrail-create');
    expect(guardrailEvents).toHaveLength(1);

    const grBody = guardrailEvents[0].body;
    expect(grBody.name).toBe('Output Moderation');
    expect(grBody.input).toBe('What is the weather?');
    expect((grBody.metadata as Record<string, unknown>).guardrailId).toBe(
      'bedrock-guardrail-123'
    );
    expect((grBody.metadata as Record<string, unknown>).outcome).toBe('passed');
    expect((grBody.metadata as Record<string, unknown>).source).toBe('OUTPUT');

    // span-update should show PASSED
    const spanUpdates = grClient.getByType('span-update');
    const guardrailUpdate = spanUpdates.find((e: any) => e.body.id === result);
    expect(guardrailUpdate).toBeDefined();
    expect(guardrailUpdate!.body.output).toBe('The weather is sunny.');
    expect(guardrailUpdate!.body.level).toBe('DEFAULT');
    expect(guardrailUpdate!.body.statusMessage).toBe('PASSED');
  });

  it('outcome=blocked should mark with ERROR level when enforced', async () => {
    const chainRunId = makeUUID();
    await grHandler.handleChainStart(
      grSerialized('RunnableSequence'),
      { input: 'test' },
      chainRunId
    );

    grHandler.traceGuardrail({
      name: 'Input Moderation',
      guardrailId: 'bedrock-guardrail-456',
      outcome: 'blocked',
      actionApplied: true,
      action: 'GUARDRAIL_INTERVENED',
      reason: 'policy_violation',
      source: 'INPUT',
      input: 'How to hack a system?',
      output: 'Sorry, I cannot help with that.',
      violations: [
        { type: 'CONTENT_POLICY', category: 'VIOLENCE', action: 'BLOCKED' },
      ],
    });

    const spanUpdates = grClient.getByType('span-update');
    const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
    expect(guardrailUpdate.body.level).toBe('ERROR');
    expect(guardrailUpdate.body.statusMessage).toBe(
      'BLOCKED: policy_violation'
    );
    expect(
      (guardrailUpdate.body.metadata as Record<string, unknown>).outcome
    ).toBe('blocked');
    expect(
      (guardrailUpdate.body.metadata as Record<string, unknown>).actionApplied
    ).toBe(true);
    expect(
      (guardrailUpdate.body.metadata as Record<string, unknown>).violations
    ).toHaveLength(1);
  });

  it('outcome=blocked with actionApplied=false should note "not enforced"', async () => {
    const chainRunId = makeUUID();
    await grHandler.handleChainStart(
      grSerialized('RunnableSequence'),
      { input: 'test' },
      chainRunId
    );

    grHandler.traceGuardrail({
      name: 'Input Moderation',
      outcome: 'blocked',
      actionApplied: false,
      reason: 'policy_violation',
      source: 'INPUT',
      input: 'Some flagged content',
    });

    const spanUpdates = grClient.getByType('span-update');
    const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
    expect(guardrailUpdate.body.level).toBe('ERROR');
    expect(guardrailUpdate.body.statusMessage).toContain('not enforced');
    expect(
      (guardrailUpdate.body.metadata as Record<string, unknown>).actionApplied
    ).toBe(false);
  });

  it('outcome=anonymized should mark with WARNING and include modified content', async () => {
    const chainRunId = makeUUID();
    await grHandler.handleChainStart(
      grSerialized('RunnableSequence'),
      { input: 'test' },
      chainRunId
    );

    grHandler.traceGuardrail({
      name: 'Output Moderation',
      outcome: 'anonymized',
      actionApplied: true,
      action: 'GUARDRAIL_INTERVENED',
      reason: 'anonymized',
      source: 'OUTPUT',
      input: 'My email is john@example.com and SSN is 123-45-6789',
      originalContent: 'My email is john@example.com and SSN is 123-45-6789',
      modifiedContent: 'My email is [EMAIL] and SSN is [SSN]',
      violations: [
        { type: 'PII_POLICY', category: 'EMAIL', action: 'ANONYMIZED' },
        { type: 'PII_POLICY', category: 'SSN', action: 'ANONYMIZED' },
      ],
    });

    const spanUpdates = grClient.getByType('span-update');
    const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
    expect(guardrailUpdate.body.level).toBe('WARNING');
    expect(guardrailUpdate.body.statusMessage).toBe(
      'ANONYMIZED: PII detected and masked'
    );
    // output should be the modified content
    expect(guardrailUpdate.body.output).toBe(
      'My email is [EMAIL] and SSN is [SSN]'
    );
    const meta = guardrailUpdate.body.metadata as Record<string, unknown>;
    expect(meta.originalContent).toBe(
      'My email is john@example.com and SSN is 123-45-6789'
    );
    expect(meta.modifiedContent).toBe('My email is [EMAIL] and SSN is [SSN]');
    expect(meta.violations).toHaveLength(2);
  });

  it('outcome=intervened should mark with WARNING level', async () => {
    const chainRunId = makeUUID();
    await grHandler.handleChainStart(
      grSerialized('RunnableSequence'),
      { input: 'test' },
      chainRunId
    );

    grHandler.traceGuardrail({
      name: 'Output Moderation',
      outcome: 'intervened',
      actionApplied: true,
      action: 'GUARDRAIL_INTERVENED',
      reason: 'intervened_passthrough',
      source: 'OUTPUT',
      input: 'Original response text',
      modifiedContent: 'Modified response text',
    });

    const spanUpdates = grClient.getByType('span-update');
    const guardrailUpdate = spanUpdates[spanUpdates.length - 1];
    expect(guardrailUpdate.body.level).toBe('WARNING');
    expect(guardrailUpdate.body.statusMessage).toBe(
      'INTERVENED: intervened_passthrough'
    );
  });

  it('traceGuardrail() should return null when no trace exists', () => {
    const freshClient = new GuardrailMockClient();
    const freshHandler: any = new ObservabilityCallbackHandler({
      client: freshClient as any,
    });

    const result = freshHandler.traceGuardrail({
      name: 'Pre-trace Guardrail',
      outcome: 'passed',
      actionApplied: false,
      source: 'INPUT',
    });

    expect(result).toBeNull();
    expect(freshClient.getByType('guardrail-create')).toHaveLength(0);
  });

  it('traceGuardrail() should use custom startTime and endTime', async () => {
    const chainRunId = makeUUID();
    await grHandler.handleChainStart(
      grSerialized('Chain'),
      { input: 'test' },
      chainRunId
    );

    const startTime = '2024-01-01T00:00:00.000Z';
    const endTime = '2024-01-01T00:00:01.500Z';

    grHandler.traceGuardrail({
      name: 'Timed Guardrail',
      outcome: 'passed',
      actionApplied: false,
      source: 'OUTPUT',
      startTime,
      endTime,
    });

    const guardrailEvents = grClient.getByType('guardrail-create');
    expect(guardrailEvents[0].body.startTime).toBe(startTime);

    const spanUpdates = grClient.getByType('span-update');
    const lastUpdate = spanUpdates[spanUpdates.length - 1];
    expect(lastUpdate.body.endTime).toBe(endTime);
  });

  it('should trace both input and output guardrails on the same trace', async () => {
    const chainRunId = makeUUID();
    await grHandler.handleChainStart(
      grSerialized('RunnableSequence'),
      { input: 'test' },
      chainRunId
    );

    const inputId = grHandler.traceGuardrail({
      name: 'Input Moderation',
      outcome: 'passed',
      actionApplied: false,
      source: 'INPUT',
      input: 'What is 2+2?',
    });

    const outputId = grHandler.traceGuardrail({
      name: 'Output Moderation',
      outcome: 'passed',
      actionApplied: false,
      source: 'OUTPUT',
      input: 'The answer is 4.',
    });

    expect(inputId).not.toBeNull();
    expect(outputId).not.toBeNull();
    expect(inputId).not.toBe(outputId);

    const guardrailEvents = grClient.getByType('guardrail-create');
    expect(guardrailEvents).toHaveLength(2);

    const traceId = guardrailEvents[0].body.traceId;
    expect(guardrailEvents[1].body.traceId).toBe(traceId);

    expect(
      (guardrailEvents[0].body.metadata as Record<string, unknown>).source
    ).toBe('INPUT');
    expect(
      (guardrailEvents[1].body.metadata as Record<string, unknown>).source
    ).toBe('OUTPUT');
  });

  it('should include full AWS Bedrock assessments and violations in metadata', async () => {
    const chainRunId = makeUUID();
    await grHandler.handleChainStart(
      grSerialized('Chain'),
      { input: 'test' },
      chainRunId
    );

    grHandler.traceGuardrail({
      name: 'Detailed Guardrail',
      outcome: 'blocked',
      actionApplied: true,
      action: 'GUARDRAIL_INTERVENED',
      reason: 'policy_violation',
      source: 'INPUT',
      violations: [
        {
          type: 'CONTENT_POLICY',
          category: 'VIOLENCE',
          confidence: 'HIGH',
          action: 'BLOCKED',
        },
        { type: 'PII_POLICY', category: 'SSN', action: 'BLOCKED' },
      ],
      assessments: [
        {
          contentPolicy: {
            filters: [
              { type: 'VIOLENCE', confidence: 'HIGH', action: 'BLOCKED' },
            ],
          },
          sensitiveInformationPolicy: {
            piiEntities: [{ type: 'SSN', action: 'BLOCKED' }],
          },
        },
      ],
    });

    const spanUpdates = grClient.getByType('span-update');
    const lastUpdate = spanUpdates[spanUpdates.length - 1];
    const meta = lastUpdate.body.metadata as Record<string, unknown>;

    expect(meta.violations).toHaveLength(2);
    expect(meta.assessments).toHaveLength(1);
    expect((meta.violations as any[])[0]).toEqual({
      type: 'CONTENT_POLICY',
      category: 'VIOLENCE',
      confidence: 'HIGH',
      action: 'BLOCKED',
    });
  });
});

// ---------------------------------------------------------------------------
// 18. Guardrail tracing with debug logging — full pipeline with all 4 outcomes
// ---------------------------------------------------------------------------

describe('Guardrail tracing with debug output', () => {
  it('should produce debug-friendly trace output for manual verification', async () => {
    const debugClient = new GuardrailMockClient();
    const debugHandler: any = new ObservabilityCallbackHandler({
      client: debugClient as any,
      traceName: 'guardrail-debug-test',
      userId: 'user-123',
      sessionId: 'session-456',
      debug: true,
    });

    // 1. Chain starts (creates trace)
    const chainRunId = makeUUID();
    await debugHandler.handleChainStart(
      grSerialized('RunnableSequence'),
      { messages: [{ role: 'user', content: 'Explain quantum computing' }] },
      chainRunId
    );

    // 2. LLM call
    const llmRunId = makeUUID();
    await debugHandler.handleLLMStart(
      grSerialized('ChatOpenAI'),
      ['Explain quantum computing'],
      llmRunId,
      chainRunId
    );
    const llmResult: LLMResult = {
      generations: [
        [{ text: 'Quantum computing uses qubits...', generationInfo: {} }],
      ],
      llmOutput: {
        tokenUsage: { promptTokens: 10, completionTokens: 25, totalTokens: 35 },
        modelName: 'gpt-4o',
      },
    };
    await debugHandler.handleLLMEnd(llmResult, llmRunId);
    await debugHandler.handleChainEnd(
      { output: 'Quantum computing uses qubits...' },
      chainRunId
    );

    // 3. Input guardrail — passed
    const inputGuardrailId = debugHandler.traceGuardrail({
      name: 'Input Moderation (Bedrock)',
      guardrailId: 'arn:aws:bedrock:us-east-1:123456:guardrail/abc123',
      guardrailVersion: '3',
      outcome: 'passed',
      actionApplied: false,
      action: 'NONE',
      reason: 'passed',
      source: 'INPUT',
      input: 'Explain quantum computing',
      startTime: '2024-06-15T10:00:00.000Z',
      endTime: '2024-06-15T10:00:00.150Z',
    });

    // 4. Output guardrail — passed
    const outputGuardrailId = debugHandler.traceGuardrail({
      name: 'Output Moderation (Bedrock)',
      guardrailId: 'arn:aws:bedrock:us-east-1:123456:guardrail/abc123',
      guardrailVersion: '3',
      outcome: 'passed',
      actionApplied: false,
      action: 'NONE',
      reason: 'passed',
      source: 'OUTPUT',
      input: 'Quantum computing uses qubits...',
      violations: [],
      assessments: [
        { topicPolicy: { topics: [] }, contentPolicy: { filters: [] } },
      ],
      startTime: '2024-06-15T10:00:01.000Z',
      endTime: '2024-06-15T10:00:01.200Z',
    });

    expect(inputGuardrailId).not.toBeNull();
    expect(outputGuardrailId).not.toBeNull();

    console.log('\n=== SAMPLE GUARDRAIL TRACE OUTPUT ===');
    debugClient.printEvents();
    console.log('=== END TRACE ===\n');

    // Verify complete event sequence
    const types = debugClient.getEventTypes();
    expect(types).toEqual([
      'trace-create',
      'chain-create',
      'generation-create',
      'generation-update',
      'span-update', // chain end
      'guardrail-create', // input guardrail
      'span-update', // input guardrail update
      'guardrail-create', // output guardrail
      'span-update', // output guardrail update
    ]);

    // Verify trace ID consistency
    const traceId = debugClient.getByType('trace-create')[0].body.id;
    for (const evt of debugClient.events) {
      const body = evt.body;
      if (evt.type !== 'trace-create') {
        expect(body.traceId).toBe(traceId);
      }
    }
  });

  it('should trace all 4 guardrail outcomes with correct levels', async () => {
    const debugClient = new GuardrailMockClient();
    const handler: any = new ObservabilityCallbackHandler({
      client: debugClient as any,
      debug: true,
    });

    // Create trace
    const chainRunId = makeUUID();
    await handler.handleChainStart(
      grSerialized('Chain'),
      { input: 'test' },
      chainRunId
    );

    // All 4 outcomes
    handler.traceGuardrail({
      name: 'Passed',
      outcome: 'passed',
      actionApplied: false,
      source: 'INPUT',
    });
    handler.traceGuardrail({
      name: 'Blocked',
      outcome: 'blocked',
      actionApplied: true,
      reason: 'policy_violation',
      source: 'INPUT',
    });
    handler.traceGuardrail({
      name: 'Anonymized',
      outcome: 'anonymized',
      actionApplied: true,
      reason: 'anonymized',
      source: 'OUTPUT',
      modifiedContent: '[EMAIL]',
    });
    handler.traceGuardrail({
      name: 'Intervened',
      outcome: 'intervened',
      actionApplied: true,
      reason: 'intervened_passthrough',
      source: 'OUTPUT',
    });

    console.log('\n=== ALL 4 GUARDRAIL OUTCOMES ===');
    debugClient.printEvents();
    console.log('=== END ===\n');

    const guardrailEvents = debugClient.getByType('guardrail-create');
    expect(guardrailEvents).toHaveLength(4);

    const spanUpdates = debugClient.getByType('span-update');
    // 4 guardrail updates (chain not ended in this test)
    expect(spanUpdates).toHaveLength(4);

    // Check levels: passed=DEFAULT, blocked=ERROR, anonymized=WARNING, intervened=WARNING
    const guardrailUpdates = spanUpdates;
    expect(guardrailUpdates[0].body.level).toBe('DEFAULT');
    expect(guardrailUpdates[1].body.level).toBe('ERROR');
    expect(guardrailUpdates[2].body.level).toBe('WARNING');
    expect(guardrailUpdates[3].body.level).toBe('WARNING');
  });
});