import {
  RunnableSequence,
  Runnable,
  RunnableConfig,
  RunnableBatchOptions,
} from "@langchain/core/runnables";
import { BaseMessage } from "@langchain/core/messages";
import { IterableReadableStream } from "@langchain/core/utils/stream";

export interface ChatGrokCallOptions extends RunnableConfig {
  temperature?: number;
  maxTokens?: number;
}

export class ChatGrok implements Record<string, unknown> {
  [key: string]: unknown;

  private apiKey: string;
  private model: string;
  private temperature: number;
  private maxTokens: number;

  constructor(options: {
    apiKey: string;
    model?: string;
    temperature?: number;
    maxTokens?: number;
  }) {
    this.apiKey = options.apiKey;
    this.model = options.model || "grok-beta";
    this.temperature = options.temperature || 0.7;
    this.maxTokens = options.maxTokens || 1000;
  }

  async invoke(
    input: string | BaseMessage,
    options?: ChatGrokCallOptions
  ): Promise<string> {
    const content = typeof input === "string" ? input : input.content;

    const response = await fetch("https://api.grok.x.ai/v1/chat", {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        Authorization: `Bearer ${this.apiKey}`,
      },
      body: JSON.stringify({
        model: this.model,
        messages: [{ role: "user", content }],
        temperature: options?.temperature ?? this.temperature,
        max_tokens: options?.maxTokens ?? this.maxTokens,
      }),
    });

    const data = await response.json();
    return data.choices[0].message.content;
  }

  async batch(
    inputs: (string | BaseMessage)[],
    options?: ChatGrokCallOptions,
    batchOptions?: RunnableBatchOptions
  ): Promise<string[]> {
    return Promise.all(inputs.map((input) => this.invoke(input, options)));
  }

  async stream(
    input: string | BaseMessage,
    options?: ChatGrokCallOptions
  ): Promise<IterableReadableStream<string>> {
    throw new Error("Streaming not implemented for ChatGrok");
  }
}