import { BedrockChatModel } from "./BedrockChatModel";

/**
 * Builds a minimal Amazon EventStream message containing the provided UTF-8 payload.
 * This helper keeps CRC fields at zero because the decoder ignores them.
 */
const buildEventStreamChunk = (payload: string): string => {
  const encoder = typeof TextEncoder !== "undefined" ? new TextEncoder() : null;
  const payloadBytes = encoder ? encoder.encode(payload) : Buffer.from(payload, "utf-8");
  const headersLength = 0;
  const totalLength = 12 + headersLength + payloadBytes.length + 4;

  const buffer = new Uint8Array(totalLength);
  const view = new DataView(buffer.buffer);

  view.setUint32(0, totalLength, false);
  view.setUint32(4, headersLength, false);
  view.setUint32(8, 0, false); // Prelude CRC (ignored by decoder)

  buffer.set(payloadBytes, 12);
  view.setUint32(totalLength - 4, 0, false); // Message CRC (ignored by decoder)

  return Buffer.from(buffer).toString("base64");
};

const createModel = (enableThinking = false): BedrockChatModel =>
  new BedrockChatModel({
    modelId: "anthropic.claude-3-haiku-20240307-v1:0",
    apiKey: "test-key",
    endpoint: "https://example.com/model/anthropic.claude-3-haiku-20240307-v1%3A0/invoke",
    streamEndpoint:
      "https://example.com/model/anthropic.claude-3-haiku-20240307-v1%3A0/invoke-with-response-stream",
    anthropicVersion: "bedrock-2023-05-31",
    enableThinking,
    fetchImplementation: jest.fn(),
  });

describe("BedrockChatModel streaming decode", () => {
  it("decodes simple base64 JSON payloads", () => {
    const payload = JSON.stringify({
      type: "content_block_delta",
      content_block_delta: {
        index: 0,
        delta: { text: "Hello there" },
      },
    });

    const base64 = Buffer.from(payload, "utf-8").toString("base64");

    const model = createModel();
    const decoded = (model as any).decodeChunkBytes(base64);

    expect(decoded).toEqual([payload]);
  });

  it("extracts payloads from Amazon EventStream encoded chunks", () => {
    const payload = JSON.stringify({
      type: "content_block_delta",
      content_block_delta: {
        index: 0,
        delta: { type: "text_delta", text: "Streaming works!" },
      },
    });

    const base64 = buildEventStreamChunk(payload);
    const model = createModel();
    const decoded = (model as any).decodeChunkBytes(base64);

    expect(decoded).toEqual([payload]);
  });

  it("produces ChatGenerationChunk entries for decoded deltas", async () => {
    const payload = JSON.stringify({
      type: "content_block_delta",
      content_block_delta: {
        index: 0,
        delta: { type: "text_delta", text: "Chunk text" },
      },
    });

    const base64 = buildEventStreamChunk(payload);
    const event = {
      type: "chunk",
      chunk: { bytes: base64 },
    };

    const model = createModel();
    const processed = await (model as any).processStreamEvent(
      event,
      undefined,
      undefined,
      undefined
    );

    expect(processed.hasText).toBe(true);
    expect(processed.deltaChunks).toHaveLength(1);
    expect(processed.deltaChunks[0]?.text).toBe("Chunk text");
  });

  describe("thinking content support", () => {
    it("buildContentItemsFromDelta recognizes thinking delta type", () => {
      const event = {
        type: "content_block_delta",
        content_block_delta: {
          index: 0,
          delta: {
            type: "thinking",
            thinking: "Let me analyze this problem...",
          },
        },
      };

      const model = createModel();
      const contentItems = (model as any).buildContentItemsFromDelta(event);

      expect(contentItems).toHaveLength(1);
      expect(contentItems[0]).toEqual({
        type: "thinking",
        thinking: "Let me analyze this problem...",
      });
    });

    it("buildContentItemsFromDelta recognizes text_delta type", () => {
      const event = {
        type: "content_block_delta",
        content_block_delta: {
          index: 0,
          delta: {
            type: "text_delta",
            text: "Based on my analysis, the answer is...",
          },
        },
      };

      const model = createModel();
      const contentItems = (model as any).buildContentItemsFromDelta(event);

      expect(contentItems).toHaveLength(1);
      expect(contentItems[0]).toEqual({
        type: "text",
        text: "Based on my analysis, the answer is...",
      });
    });

    it("processStreamEvent returns chunks with thinking content array", async () => {
      const payload = JSON.stringify({
        type: "content_block_delta",
        content_block_delta: {
          index: 0,
          delta: {
            type: "thinking",
            thinking: "Reasoning through this...",
          },
        },
      });

      const base64 = buildEventStreamChunk(payload);
      const event = {
        type: "chunk",
        chunk: { bytes: base64 },
      };

      const model = createModel();
      const processed = await (model as any).processStreamEvent(
        event,
        undefined,
        undefined,
        undefined
      );

      expect(processed.hasText).toBe(true);
      expect(processed.deltaChunks).toHaveLength(1);
      const chunk = processed.deltaChunks[0];
      expect(chunk?.text).toBe("Reasoning through this...");

      // Check that content is an array with thinking type
      expect(Array.isArray(chunk?.message.content)).toBe(true);
      const content = chunk?.message.content as any[];
      expect(content).toHaveLength(1);
      expect(content[0]).toEqual({
        type: "thinking",
        thinking: "Reasoning through this...",
      });

      // Check for OpenRouter compatibility
      expect(chunk?.message.additional_kwargs).toBeDefined();
      expect(chunk?.message.additional_kwargs?.delta).toEqual({
        reasoning: "Reasoning through this...",
      });
    });

    it("processStreamEvent returns chunks with text content array", async () => {
      const payload = JSON.stringify({
        type: "content_block_delta",
        content_block_delta: {
          index: 0,
          delta: {
            type: "text_delta",
            text: "Here is my final answer.",
          },
        },
      });

      const base64 = buildEventStreamChunk(payload);
      const event = {
        type: "chunk",
        chunk: { bytes: base64 },
      };

      const model = createModel();
      const processed = await (model as any).processStreamEvent(
        event,
        undefined,
        undefined,
        undefined
      );

      expect(processed.hasText).toBe(true);
      expect(processed.deltaChunks).toHaveLength(1);
      const chunk = processed.deltaChunks[0];
      expect(chunk?.text).toBe("Here is my final answer.");

      // Check that content is an array with text type
      expect(Array.isArray(chunk?.message.content)).toBe(true);
      const content = chunk?.message.content as any[];
      expect(content).toHaveLength(1);
      expect(content[0]).toEqual({
        type: "text",
        text: "Here is my final answer.",
      });

      // No additional_kwargs.delta for regular text
      expect(chunk?.message.additional_kwargs?.delta).toBeUndefined();
    });

    it("handles mixed thinking and text deltas correctly", async () => {
      const model = createModel();

      // First chunk: thinking
      const thinkingPayload = JSON.stringify({
        type: "content_block_delta",
        content_block_delta: {
          index: 0,
          delta: {
            type: "thinking",
            thinking: "First, I'll consider...",
          },
        },
      });

      const thinkingBase64 = buildEventStreamChunk(thinkingPayload);
      const thinkingEvent = {
        type: "chunk",
        chunk: { bytes: thinkingBase64 },
      };

      const thinkingResult = await (model as any).processStreamEvent(
        thinkingEvent,
        undefined,
        undefined,
        undefined
      );

      expect(thinkingResult.deltaChunks).toHaveLength(1);
      const thinkingChunk = thinkingResult.deltaChunks[0];
      expect((thinkingChunk?.message.content as any[])[0]?.type).toBe("thinking");

      // Second chunk: text
      const textPayload = JSON.stringify({
        type: "content_block_delta",
        content_block_delta: {
          index: 0,
          delta: {
            type: "text_delta",
            text: "Therefore, the answer is X.",
          },
        },
      });

      const textBase64 = buildEventStreamChunk(textPayload);
      const textEvent = {
        type: "chunk",
        chunk: { bytes: textBase64 },
      };

      const textResult = await (model as any).processStreamEvent(
        textEvent,
        undefined,
        undefined,
        undefined
      );

      expect(textResult.deltaChunks).toHaveLength(1);
      const textChunk = textResult.deltaChunks[0];
      expect((textChunk?.message.content as any[])[0]?.type).toBe("text");
    });

    it("extractStreamText can fallback to extract thinking content", () => {
      const event = {
        type: "content_block_delta",
        content_block_delta: {
          delta: {
            type: "thinking",
            thinking: "Fallback thinking extraction",
          },
        },
      };

      const model = createModel();
      const extracted = (model as any).extractStreamText(event);

      expect(extracted).toBe("Fallback thinking extraction");
    });

    it("handles empty thinking content gracefully", () => {
      const event = {
        type: "content_block_delta",
        content_block_delta: {
          delta: {
            type: "thinking",
            thinking: "",
          },
        },
      };

      const model = createModel();
      const contentItems = (model as any).buildContentItemsFromDelta(event);

      expect(contentItems).toHaveLength(1);
      expect(contentItems[0]).toEqual({
        type: "thinking",
        thinking: "",
      });
    });
  });

  describe("thinking mode enablement", () => {
    it("includes thinking parameter when enableThinking is true", () => {
      const model = createModel(true);
      const requestBody = (model as any).buildRequestBody([
        { role: "user", content: "test", _getType: () => "human" },
      ]);

      expect(requestBody.thinking).toEqual({
        type: "enabled",
        budget_tokens: 2048,
      });
      expect(requestBody.temperature).toBe(1);
      expect(requestBody.anthropic_version).toBe("bedrock-2023-05-31");
    });

    it("does not include thinking parameter when enableThinking is false", () => {
      const model = createModel(false);
      const requestBody = (model as any).buildRequestBody(
        [{ role: "user", content: "test", _getType: () => "human" }],
        { temperature: 0.7 }
      );

      expect(requestBody.thinking).toBeUndefined();
      expect(requestBody.temperature).toBe(0.7);
      // anthropic_version should always be present when provided (required for all Bedrock requests)
      expect(requestBody.anthropic_version).toBe("bedrock-2023-05-31");
    });

    it("respects user temperature when thinking is disabled", () => {
      const model = createModel(false);
      const requestBody = (model as any).buildRequestBody(
        [{ role: "user", content: "test", _getType: () => "human" }],
        { temperature: 0.5 }
      );

      expect(requestBody.temperature).toBe(0.5);
      expect(requestBody.thinking).toBeUndefined();
    });

    it("forces temperature to 1 when thinking is enabled", () => {
      const model = createModel(true);
      const requestBody = (model as any).buildRequestBody(
        [{ role: "user", content: "test", _getType: () => "human" }],
        { temperature: 0.5 } // User tries to set 0.5, should be overridden to 1
      );

      expect(requestBody.temperature).toBe(1);
      expect(requestBody.thinking).toBeDefined();
    });
  });

  describe("vision support", () => {
    describe("convertImageContent", () => {
      it("converts valid data URL to Claude image format", () => {
        const model = createModel();
        const dataUrl = "data:image/jpeg;base64,/9j/4AAQSkZJRg==";
        const result = (model as any).convertImageContent(dataUrl);

        expect(result).toEqual({
          type: "image",
          source: {
            type: "base64",
            media_type: "image/jpeg",
            data: "/9j/4AAQSkZJRg==",
          },
        });
      });

      it("handles PNG images", () => {
        const model = createModel();
        const dataUrl = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUg==";
        const result = (model as any).convertImageContent(dataUrl);

        expect(result).toEqual({
          type: "image",
          source: {
            type: "base64",
            media_type: "image/png",
            data: "iVBORw0KGgoAAAANSUhEUg==",
          },
        });
      });

      it("returns null for invalid data URL format", () => {
        const model = createModel();
        const invalidUrl = "not-a-data-url";
        const result = (model as any).convertImageContent(invalidUrl);

        expect(result).toBeNull();
      });

      it("returns null for non-image media type", () => {
        const model = createModel();
        const dataUrl = "data:text/plain;base64,SGVsbG8gV29ybGQ=";
        const result = (model as any).convertImageContent(dataUrl);

        expect(result).toBeNull();
      });
    });

    describe("normaliseMessageContent", () => {
      it("preserves array content with images", () => {
        const model = createModel();
        const message = {
          content: [
            { type: "text", text: "What's in this image?" },
            {
              type: "image_url",
              image_url: { url: "data:image/jpeg;base64,/9j/4AAQSkZJRg==" },
            },
          ],
          _getType: () => "human",
        };

        const result = (model as any).normaliseMessageContent(message);

        expect(Array.isArray(result)).toBe(true);
        expect(result).toHaveLength(2);
        expect(result[0]).toEqual({ type: "text", text: "What's in this image?" });
        expect(result[1]).toEqual({
          type: "image_url",
          image_url: { url: "data:image/jpeg;base64,/9j/4AAQSkZJRg==" },
        });
      });

      it("flattens array content without images to string", () => {
        const model = createModel();
        const message = {
          content: [
            { type: "text", text: "Hello " },
            { type: "text", text: "world!" },
          ],
          _getType: () => "human",
        };

        const result = (model as any).normaliseMessageContent(message);

        expect(typeof result).toBe("string");
        expect(result).toBe("Hello world!");
      });

      it("returns string content unchanged", () => {
        const model = createModel();
        const message = {
          content: "Simple text message",
          _getType: () => "human",
        };

        const result = (model as any).normaliseMessageContent(message);

        expect(result).toBe("Simple text message");
      });
    });

    describe("buildRequestBody with images", () => {
      it("includes images in request body for multimodal messages", () => {
        const model = createModel();
        const messages = [
          {
            content: [
              { type: "text", text: "What's in this image?" },
              {
                type: "image_url",
                image_url: { url: "data:image/jpeg;base64,/9j/4AAQSkZJRg==" },
              },
            ],
            _getType: () => "human",
          },
        ];

        const requestBody = (model as any).buildRequestBody(messages);

        expect(requestBody.messages).toHaveLength(1);
        expect(requestBody.messages[0].content).toHaveLength(2);

        // Check text block
        expect(requestBody.messages[0].content[0]).toEqual({
          type: "text",
          text: "What's in this image?",
        });

        // Check image block (converted to Claude format)
        expect(requestBody.messages[0].content[1]).toEqual({
          type: "image",
          source: {
            type: "base64",
            media_type: "image/jpeg",
            data: "/9j/4AAQSkZJRg==",
          },
        });
      });

      it("handles multiple images in a single message", () => {
        const model = createModel();
        const messages = [
          {
            content: [
              { type: "text", text: "Compare these images:" },
              {
                type: "image_url",
                image_url: { url: "data:image/jpeg;base64,IMAGE1DATA" },
              },
              {
                type: "image_url",
                image_url: { url: "data:image/png;base64,IMAGE2DATA" },
              },
            ],
            _getType: () => "human",
          },
        ];

        const requestBody = (model as any).buildRequestBody(messages);

        expect(requestBody.messages[0].content).toHaveLength(3);
        expect(requestBody.messages[0].content[0].type).toBe("text");
        expect(requestBody.messages[0].content[1].type).toBe("image");
        expect(requestBody.messages[0].content[1].source.media_type).toBe("image/jpeg");
        expect(requestBody.messages[0].content[2].type).toBe("image");
        expect(requestBody.messages[0].content[2].source.media_type).toBe("image/png");
      });

      it("handles text-only messages correctly", () => {
        const model = createModel();
        const messages = [
          {
            content: "Just text, no images",
            _getType: () => "human",
          },
        ];

        const requestBody = (model as any).buildRequestBody(messages);

        expect(requestBody.messages).toHaveLength(1);
        expect(requestBody.messages[0].content).toHaveLength(1);
        expect(requestBody.messages[0].content[0]).toEqual({
          type: "text",
          text: "Just text, no images",
        });
      });

      it("skips invalid images and keeps valid content", () => {
        const model = createModel();
        const messages = [
          {
            content: [
              { type: "text", text: "Valid text" },
              {
                type: "image_url",
                image_url: { url: "invalid-url" }, // Invalid - should be skipped
              },
              {
                type: "image_url",
                image_url: { url: "data:image/jpeg;base64,VALIDDATA" }, // Valid
              },
            ],
            _getType: () => "human",
          },
        ];

        const requestBody = (model as any).buildRequestBody(messages);

        // Should have text + 1 valid image (invalid one skipped)
        expect(requestBody.messages[0].content).toHaveLength(2);
        expect(requestBody.messages[0].content[0].type).toBe("text");
        expect(requestBody.messages[0].content[1].type).toBe("image");
      });
    });
  });
});
