import { z } from 'zod';
import { tool } from '@langchain/core/tools';
import { AIMessage } from '@langchain/core/messages';
import { describe, it, expect } from '@jest/globals';
import type { StructuredToolInterface } from '@langchain/core/tools';
import { ToolNode } from '../ToolNode';
import { StreamingToolCallBuffer } from '../StreamingToolCallBuffer';

/**
 * Creates a mock content_tool that captures the args it receives.
 * Used to verify that truncated args are recovered from the streaming buffer.
 */
function createMockContentTool(
  capturedArgs: Record<string, unknown>[]
): StructuredToolInterface {
  return tool(
    async (input) => {
      capturedArgs.push({ ...input });
      return `Wrote ${(input as Record<string, unknown>).filename ?? 'file'}`;
    },
    {
      name: 'content_tool',
      description: 'Read, write, and manage files',
      schema: z.object({
        action: z.string(),
        filename: z.string().optional(),
        content: z.string().optional(),
      }),
    }
  ) as unknown as StructuredToolInterface;
}

function createAIMessageWithContentTool(
  callId: string,
  args: Record<string, unknown>
): AIMessage {
  return new AIMessage({
    content: '',
    tool_calls: [
      {
        id: callId,
        name: 'content_tool',
        args,
      },
    ],
  });
}

describe('ToolNode truncated arg recovery', () => {
  it('recovers missing content field from streaming buffer', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const buffer = new StreamingToolCallBuffer();

    // Simulate streaming: buffer accumulated the full raw args
    buffer.append(
      'tc_1',
      '{"action":"write","filename":"app.tsx","content":"function App() {\\n  return <div>Hello</div>;\\n}"}'
    );
    buffer.setToolName('tc_1', 'content_tool');

    const contentTool = createMockContentTool(capturedArgs);
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    // Simulate truncated args: parsePartialJson lost the content field
    const aiMsg = createAIMessageWithContentTool('tc_1', {
      action: 'write',
      filename: 'app.tsx',
      // content is MISSING — truncated by max_tokens
    });

    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs).toHaveLength(1);
    expect(capturedArgs[0].action).toBe('write');
    expect(capturedArgs[0].filename).toBe('app.tsx');
    expect(capturedArgs[0].content).toBe(
      'function App() {\n  return <div>Hello</div>;\n}'
    );
  });

  it('recovers completely empty args from streaming buffer', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const buffer = new StreamingToolCallBuffer();

    // Buffer has full args, but parsePartialJson returned empty object
    buffer.append(
      'tc_1',
      '{"action":"write","filename":"test.js","content":"const x = 1;"}'
    );
    buffer.setToolName('tc_1', 'content_tool');

    const contentTool = createMockContentTool(capturedArgs);
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    // Empty args — JSON parsing aborted completely
    const aiMsg = createAIMessageWithContentTool('tc_1', {});

    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs).toHaveLength(1);
    expect(capturedArgs[0].action).toBe('write');
    expect(capturedArgs[0].filename).toBe('test.js');
    expect(capturedArgs[0].content).toBe('const x = 1;');
  });

  it('does not modify args when buffer has no entry', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const buffer = new StreamingToolCallBuffer();
    // No buffer entry for tc_1

    const contentTool = createMockContentTool(capturedArgs);
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    const aiMsg = createAIMessageWithContentTool('tc_1', {
      action: 'read',
      filename: 'test.js',
    });

    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs).toHaveLength(1);
    expect(capturedArgs[0].action).toBe('read');
    expect(capturedArgs[0].filename).toBe('test.js');
    expect(capturedArgs[0].content).toBeUndefined();
  });

  it('does not overwrite existing parsed values', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const buffer = new StreamingToolCallBuffer();

    // Buffer has different values than parsed
    buffer.append(
      'tc_1',
      '{"action":"write","filename":"old.tsx","content":"old content"}'
    );

    const contentTool = createMockContentTool(capturedArgs);
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    // Parsed args have action and filename already
    const aiMsg = createAIMessageWithContentTool('tc_1', {
      action: 'write',
      filename: 'new.tsx',
      // content is missing — should be recovered
    });

    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs).toHaveLength(1);
    expect(capturedArgs[0].action).toBe('write');
    expect(capturedArgs[0].filename).toBe('new.tsx'); // Keeps parsed value, not buffer value
    expect(capturedArgs[0].content).toBe('old content'); // Recovered from buffer
  });

  it('clears buffer after recovery', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const buffer = new StreamingToolCallBuffer();

    buffer.append('tc_1', '{"action":"write","content":"test"}');

    const contentTool = createMockContentTool(capturedArgs);
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    const aiMsg = createAIMessageWithContentTool('tc_1', { action: 'write' });
    await toolNode.invoke({ messages: [aiMsg] });

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

  it('works with tools other than content_tool', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const buffer = new StreamingToolCallBuffer();

    buffer.append('tc_1', '{"query":"SELECT * FROM users","database":"prod"}');
    buffer.setToolName('tc_1', 'sql_query');

    const sqlTool = tool(
      async (input) => {
        capturedArgs.push({ ...input });
        return 'Results: 42 rows';
      },
      {
        name: 'sql_query',
        description: 'Execute SQL query',
        schema: z.object({
          query: z.string(),
          database: z.string().optional(),
        }),
      }
    ) as unknown as StructuredToolInterface;

    const toolNode = new ToolNode({
      tools: [sqlTool],
      streamingToolCallBuffer: buffer,
    });

    // Truncated: only has query, missing database
    const aiMsg = new AIMessage({
      content: '',
      tool_calls: [
        {
          id: 'tc_1',
          name: 'sql_query',
          args: { query: 'SELECT * FROM users' },
        },
      ],
    });

    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs).toHaveLength(1);
    expect(capturedArgs[0].query).toBe('SELECT * FROM users');
    expect(capturedArgs[0].database).toBe('prod');
  });

  it('recovers truncated content with newlines and special chars', async () => {
    const capturedArgs: Record<string, unknown>[] = [];
    const buffer = new StreamingToolCallBuffer();

    const largeContent =
      'import React from \\"react\\";\\n\\n' +
      'export default function Dashboard() {\\n' +
      '  const [count, setCount] = React.useState(0);\\n' +
      '  return (\\n' +
      '    <div className=\\"dashboard\\">\\n' +
      '      <h1>Dashboard</h1>\\n' +
      '      <button onClick={() => setCount(c => c + 1)}>\\n' +
      '        Count: {count}\\n' +
      '      </button>\\n' +
      '    </div>\\n' +
      '  );\\n' +
      '}';

    buffer.append(
      'tc_1',
      `{"action":"write","filename":"Dashboard.tsx","content":"${largeContent}"}`
    );

    const contentTool = createMockContentTool(capturedArgs);
    const toolNode = new ToolNode({
      tools: [contentTool],
      streamingToolCallBuffer: buffer,
    });

    // content_tool with missing content
    const aiMsg = createAIMessageWithContentTool('tc_1', {
      action: 'write',
      filename: 'Dashboard.tsx',
    });

    await toolNode.invoke({ messages: [aiMsg] });

    expect(capturedArgs).toHaveLength(1);
    expect(capturedArgs[0].content).toContain('import React from "react"');
    expect(capturedArgs[0].content).toContain(
      'export default function Dashboard()'
    );
    expect(capturedArgs[0].content).toContain('Count: {count}');
  });
});