/**
 * Integration test: Full Truncation Recovery Pipeline
 *
 * Tests the end-to-end flow that would catch the Bedrock buffer bug:
 *   Streaming chunks → handleToolCallChunks → StreamingToolCallBuffer → handleToolCalls → ToolNode.recoverTruncatedArgs → tool receives recovered args
 *
 * Uses real StreamingToolCallBuffer (not mocked) and real ToolNode with real tool definitions.
 * Only the graph dispatch layer is mocked (since we're not testing SSE dispatch).
 *
 * Verified against live Bedrock API output (see bedrock-live-poc.ts):
 *   START chunk: { id, name, index, args: "" }
 *   DELTA chunks: { args, index } — NO id, NO name
 */

import { z } from 'zod';
import { tool } from '@langchain/core/tools';
import { AIMessage } from '@langchain/core/messages';
import { describe, it, expect, beforeEach, jest } from '@jest/globals';
import type { ToolCallChunk } from '@langchain/core/messages/tool';
import type { StructuredToolInterface } from '@langchain/core/tools';
import type { StandardGraph } from '@/graphs';
import type * as t from '@/types';
import { StepTypes } from '@/common';
import { StreamingToolCallBuffer } from '../StreamingToolCallBuffer';
import { handleToolCallChunks } from '../handlers';
import { ToolNode } from '../ToolNode';

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

/** Creates a mock graph that delegates buffer operations to a REAL StreamingToolCallBuffer */
// eslint-disable-next-line @typescript-eslint/explicit-function-return-type
function createIntegrationGraph(buffer: StreamingToolCallBuffer) {
  let stepCounter = 0;
  return {
    getStepKey: jest.fn<() => string>().mockReturnValue('step-key'),
    getStepIdByKey: jest.fn<() => string>().mockReturnValue('prev-step-id'),
    getRunStep: jest.fn<() => t.RunStep | undefined>().mockReturnValue({
      type: StepTypes.MESSAGE_CREATION,
      id: 'prev-step-id',
      index: 0,
      stepDetails: {
        type: StepTypes.MESSAGE_CREATION,
        message_creation: { message_id: 'msg-1' },
      },
      usage: null,
    }),
    dispatchRunStep: jest
      .fn<() => Promise<string>>()
      .mockImplementation(async () => `step-${++stepCounter}`),
    dispatchRunStepDelta: jest
      .fn<() => Promise<void>>()
      .mockResolvedValue(undefined),
    toolCallStepIds: new Map<string, string>(),
    messageStepHasToolCalls: new Map<string, boolean>(),
    messageIdsByStepKey: new Map<string, string>(),
    prelimMessageIdsByStepKey: new Map<string, string>(),
    streamingToolCallBuffer: buffer, // REAL buffer, not mocked
  };
}

/** Creates a tool that captures received args for assertion */
function createCaptureTool(
  name: string,
  schema: z.ZodObject<Record<string, z.ZodTypeAny>>,
  capturedArgs: Record<string, unknown>[]
): StructuredToolInterface {
  return tool(
    async (input) => {
      capturedArgs.push({ ...input });
      return 'OK';
    },
    { name, description: `Test tool: ${name}`, schema }
  ) as unknown as StructuredToolInterface;
}

/** Simulates Bedrock streaming: START chunk with {id, name, index, args:""}, then DELTA chunks with {args, index} */
function simulateBedrockChunks(
  toolCallId: string,
  toolName: string,
  index: number,
  argsJson: string,
  chunkSize = 30
): ToolCallChunk[] {
  const chunks: ToolCallChunk[] = [];

  // START chunk: id + name + index, args="" (verified from live Bedrock API)
  chunks.push({
    id: toolCallId,
    name: toolName,
    index,
    args: '',
    type: 'tool_call_chunk',
  });

  // DELTA chunks: args + index, NO id, NO name (verified from live Bedrock API)
  for (let i = 0; i < argsJson.length; i += chunkSize) {
    chunks.push({
      id: undefined,
      name: undefined,
      index,
      args: argsJson.substring(i, i + chunkSize),
      type: 'tool_call_chunk',
    });
  }

  return chunks;
}

/** Simulates OpenAI/Anthropic streaming: every chunk has {id, args} */
function simulateOpenAIChunks(
  toolCallId: string,
  toolName: string,
  index: number,
  argsJson: string,
  chunkSize = 30
): ToolCallChunk[] {
  const chunks: ToolCallChunk[] = [];

  // First chunk: id + name + first args fragment
  const firstArgs = argsJson.substring(0, chunkSize);
  chunks.push({
    id: toolCallId,
    name: toolName,
    index,
    args: firstArgs,
    type: 'tool_call_chunk',
  });

  // Subsequent chunks: id + args (OpenAI sends id on every chunk)
  for (let i = chunkSize; i < argsJson.length; i += chunkSize) {
    chunks.push({
      id: toolCallId,
      name: undefined,
      index,
      args: argsJson.substring(i, i + chunkSize),
      type: 'tool_call_chunk',
    });
  }

  return chunks;
}

const contentToolSchema = z.object({
  action: z.string(),
  filename: z.string().optional(),
  content: z.string().optional(),
});

const codeExecutorSchema = z.object({
  code: z.string(),
  language: z.string().optional(),
});

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

describe('Truncation Recovery Pipeline (Integration)', () => {
  let buffer: StreamingToolCallBuffer;

  beforeEach(() => {
    buffer = new StreamingToolCallBuffer();
  });

  it('Bedrock truncated write — full pipeline recovery', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const graph = createIntegrationGraph(buffer);

    // Full args JSON that Bedrock would stream (including content)
    const fullArgsJson =
      '{"action":"write","filename":"app.tsx","content":"import React from \\"react\\";\\nexport default function App() {\\n  return <div>Hello World</div>;\\n}"}';

    // Simulate Bedrock streaming: START + DELTA chunks
    const chunks = simulateBedrockChunks(
      'tooluse_abc',
      'content_tool',
      0,
      fullArgsJson,
      20
    );

    // Feed all chunks through handleToolCallChunks (real buffer accumulation)
    for (const chunk of chunks) {
      await handleToolCallChunks({
        graph: graph as unknown as StandardGraph,
        stepKey: 'step-key',
        toolCallChunks: [chunk],
        metadata: { run_id: 'test' },
      });
    }

    // Verify buffer was populated
    expect(buffer.has('tooluse_abc')).toBe(true);
    expect(buffer.getRawArgs('tooluse_abc')).toBe(fullArgsJson);

    // Simulate truncation: parsePartialJson lost the content field
    const truncatedArgs = { action: 'write', filename: 'app.tsx' }; // content MISSING

    // Create ToolNode with real buffer
    const contentTool = createCaptureTool(
      'content_tool',
      contentToolSchema,
      capturedArgs
    );
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    // Invoke ToolNode with truncated args
    const aiMsg = new AIMessage({
      content: '',
      tool_calls: [
        { id: 'tooluse_abc', name: 'content_tool', args: truncatedArgs },
      ],
    });
    await toolNode.invoke({ messages: [aiMsg] });

    // Assert: tool received RECOVERED content, not truncated
    expect(capturedArgs).toHaveLength(1);
    expect(capturedArgs[0].action).toBe('write');
    expect(capturedArgs[0].filename).toBe('app.tsx');
    expect(capturedArgs[0].content).toBe(
      'import React from "react";\nexport default function App() {\n  return <div>Hello World</div>;\n}'
    );
  });

  it('OpenAI/Anthropic chunked write — full pipeline recovery', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const graph = createIntegrationGraph(buffer);

    const fullArgsJson =
      '{"action":"write","filename":"test.tsx","content":"const x = 42;\\nconst y = 100;"}';
    const chunks = simulateOpenAIChunks(
      'tc_openai',
      'content_tool',
      0,
      fullArgsJson,
      15
    );

    for (const chunk of chunks) {
      await handleToolCallChunks({
        graph: graph as unknown as StandardGraph,
        stepKey: 'step-key',
        toolCallChunks: [chunk],
        metadata: { run_id: 'test' },
      });
    }

    expect(buffer.has('tc_openai')).toBe(true);
    expect(buffer.getRawArgs('tc_openai')).toBe(fullArgsJson);

    // Simulate truncation: content missing
    const contentTool = createCaptureTool(
      'content_tool',
      contentToolSchema,
      capturedArgs
    );
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    const aiMsg = new AIMessage({
      content: '',
      tool_calls: [
        {
          id: 'tc_openai',
          name: 'content_tool',
          args: { action: 'write', filename: 'test.tsx' },
        },
      ],
    });
    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs[0].content).toBe('const x = 42;\nconst y = 100;');
  });

  it('multiple parallel tool calls — no cross-contamination', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const graph = createIntegrationGraph(buffer);

    // Two parallel tool calls with different indexes
    const args1 =
      '{"action":"write","filename":"a.tsx","content":"file A content"}';
    const args2 =
      '{"action":"write","filename":"b.tsx","content":"file B content"}';

    const chunks1 = simulateBedrockChunks('tc_1', 'content_tool', 0, args1, 15);
    const chunks2 = simulateBedrockChunks('tc_2', 'content_tool', 1, args2, 15);

    // Interleave chunks (realistic concurrent streaming)
    const interleaved: ToolCallChunk[] = [];
    const maxLen = Math.max(chunks1.length, chunks2.length);
    for (let i = 0; i < maxLen; i++) {
      if (i < chunks1.length) interleaved.push(chunks1[i]);
      if (i < chunks2.length) interleaved.push(chunks2[i]);
    }

    for (const chunk of interleaved) {
      await handleToolCallChunks({
        graph: graph as unknown as StandardGraph,
        stepKey: 'step-key',
        toolCallChunks: [chunk],
        metadata: { run_id: 'test' },
      });
    }

    expect(buffer.getRawArgs('tc_1')).toBe(args1);
    expect(buffer.getRawArgs('tc_2')).toBe(args2);

    const contentTool = createCaptureTool(
      'content_tool',
      contentToolSchema,
      capturedArgs
    );
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    // Only tc_1 is truncated (missing content), tc_2 has all fields
    const aiMsg = new AIMessage({
      content: '',
      tool_calls: [
        {
          id: 'tc_1',
          name: 'content_tool',
          args: { action: 'write', filename: 'a.tsx' },
        }, // truncated
        {
          id: 'tc_2',
          name: 'content_tool',
          args: {
            action: 'write',
            filename: 'b.tsx',
            content: 'file B content',
          },
        }, // complete
      ],
    });
    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs).toHaveLength(2);
    // tc_1: content recovered from buffer
    expect(capturedArgs[0].content).toBe('file A content');
    // tc_2: content unchanged (already present in parsed args)
    expect(capturedArgs[1].content).toBe('file B content');
  });

  it('buffer cleanup after recovery — no stale data', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const graph = createIntegrationGraph(buffer);

    const fullArgs = '{"action":"write","filename":"x.tsx","content":"test"}';
    const chunks = simulateBedrockChunks(
      'tc_cleanup',
      'content_tool',
      0,
      fullArgs
    );

    for (const chunk of chunks) {
      await handleToolCallChunks({
        graph: graph as unknown as StandardGraph,
        stepKey: 'step-key',
        toolCallChunks: [chunk],
        metadata: { run_id: 'test' },
      });
    }

    expect(buffer.has('tc_cleanup')).toBe(true);

    const contentTool = createCaptureTool(
      'content_tool',
      contentToolSchema,
      capturedArgs
    );
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    const aiMsg = new AIMessage({
      content: '',
      tool_calls: [
        { id: 'tc_cleanup', name: 'content_tool', args: { action: 'write' } },
      ],
    });
    await toolNode.invoke({ messages: [aiMsg] });

    // Buffer should be cleared after processing
    expect(buffer.has('tc_cleanup')).toBe(false);
    expect(buffer.getRawArgs('tc_cleanup')).toBeUndefined();
  });

  it('no truncation — buffer populated but not used for overwrite', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const graph = createIntegrationGraph(buffer);

    const fullArgs = '{"action":"read","filename":"existing.tsx"}';
    const chunks = simulateBedrockChunks(
      'tc_complete',
      'content_tool',
      0,
      fullArgs
    );

    for (const chunk of chunks) {
      await handleToolCallChunks({
        graph: graph as unknown as StandardGraph,
        stepKey: 'step-key',
        toolCallChunks: [chunk],
        metadata: { run_id: 'test' },
      });
    }

    const contentTool = createCaptureTool(
      'content_tool',
      contentToolSchema,
      capturedArgs
    );
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    // Complete args — nothing missing
    const aiMsg = new AIMessage({
      content: '',
      tool_calls: [
        {
          id: 'tc_complete',
          name: 'content_tool',
          args: { action: 'read', filename: 'existing.tsx' },
        },
      ],
    });
    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs[0].action).toBe('read');
    expect(capturedArgs[0].filename).toBe('existing.tsx');
    // No content field in args or buffer — should stay absent
    expect(capturedArgs[0]).not.toHaveProperty('content');
  });

  it('generic tool recovery — non-content_tool (code_executor)', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const graph = createIntegrationGraph(buffer);

    const largeCode =
      'function fibonacci(n) {\\n  if (n <= 1) return n;\\n  return fibonacci(n-1) + fibonacci(n-2);\\n}\\nconsole.log(fibonacci(10));';
    const fullArgs = `{"code":"${largeCode}","language":"javascript"}`;
    const chunks = simulateBedrockChunks(
      'tc_code',
      'execute_code',
      0,
      fullArgs,
      20
    );

    for (const chunk of chunks) {
      await handleToolCallChunks({
        graph: graph as unknown as StandardGraph,
        stepKey: 'step-key',
        toolCallChunks: [chunk],
        metadata: { run_id: 'test' },
      });
    }

    const codeTool = createCaptureTool(
      'execute_code',
      codeExecutorSchema,
      capturedArgs
    );
    const toolNode = new ToolNode({
      tools: [codeTool],
      streamingToolCallBuffer: buffer,
    });

    // Truncated: only has language, code is missing
    const aiMsg = new AIMessage({
      content: '',
      tool_calls: [
        {
          id: 'tc_code',
          name: 'execute_code',
          args: { language: 'javascript' },
        },
      ],
    });
    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs[0].code).toContain('fibonacci');
    expect(capturedArgs[0].language).toBe('javascript');
  });

  it('Bedrock large content (5000+ chars) — recovery preserves full content', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const graph = createIntegrationGraph(buffer);

    // Generate large content (simulating a real dashboard component)
    const lines: string[] = [];
    lines.push('import React from \\"react\\";');
    lines.push('import { useState, useEffect } from \\"react\\";');
    for (let i = 0; i < 100; i++) {
      lines.push(`// Component section ${i}`);
      lines.push(`const Section${i} = () => <div>Section ${i}</div>;`);
    }
    lines.push('export default function Dashboard() {');
    lines.push('  return <div>Dashboard</div>;');
    lines.push('}');
    const content = lines.join('\\n');

    const fullArgs = `{"action":"write","filename":"dashboard.tsx","content":"${content}"}`;
    const chunks = simulateBedrockChunks(
      'tc_large',
      'content_tool',
      0,
      fullArgs,
      50
    );

    for (const chunk of chunks) {
      await handleToolCallChunks({
        graph: graph as unknown as StandardGraph,
        stepKey: 'step-key',
        toolCallChunks: [chunk],
        metadata: { run_id: 'test' },
      });
    }

    const contentTool = createCaptureTool(
      'content_tool',
      contentToolSchema,
      capturedArgs
    );
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    // Completely empty args (truncation wiped everything)
    const aiMsg = new AIMessage({
      content: '',
      tool_calls: [{ id: 'tc_large', name: 'content_tool', args: {} }],
    });
    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs[0].action).toBe('write');
    expect(capturedArgs[0].filename).toBe('dashboard.tsx');
    const recoveredContent = capturedArgs[0].content as string;
    expect(recoveredContent).toContain('import React from "react"');
    expect(recoveredContent).toContain('Section 99');
    expect(recoveredContent).toContain('export default function Dashboard()');
  });
});