import { BaseChatModel } from "@langchain/core/language_models/chat_models";
import { BaseChatModelParams } from "@langchain/core/language_models/chat_models";
import { BaseMessage, AIMessage } from "@langchain/core/messages";
import { CallbackManagerForLLMRun } from "@langchain/core/callbacks/manager";

interface ChatCompletionResponse {
  id: string;
  object: string;
  created: number;
  model: string;
  choices: {
    message: {
      role: string;
      content: string;
    };
    finish_reason: string;
    index: number;
  }[];
  usage: {
    prompt_tokens: number;
    completion_tokens: number;
    total_tokens: number;
  };
}

/**
 * Deepseek模型参数接口
 */
export interface DeepseekInput extends BaseChatModelParams {
  apiKey: string;
  baseURL?: string;
  temperature?: number;
  maxTokens?: number;
  modelName?: string;
}

/**
 * Deepseek聊天模型实现
 */
export class ChatDeepseek extends BaseChatModel {
  apiKey: string;
  baseURL: string;
  temperature: number;
  maxTokens: number;
  modelName: string;

  constructor(fields?: Partial<DeepseekInput>) {
    super(fields ?? {});

    this.apiKey = fields?.apiKey ?? "";
    this.baseURL = fields?.baseURL ?? "https://api.deepseek.com/v1";
    this.temperature = fields?.temperature ?? 0.7;
    this.maxTokens = fields?.maxTokens ?? 2048;
    this.modelName = fields?.modelName ?? "deepseek-chat";
  }

  /**
   * 生成聊天补全
   */
  async _generate(
    messages: BaseMessage[],
    options: this["ParsedCallOptions"],
    runManager?: CallbackManagerForLLMRun
  ) {
    const response = await fetch(`${this.baseURL}/chat/completions`, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "Authorization": `Bearer ${this.apiKey}`,
      },
      body: JSON.stringify({
        model: this.modelName,
        messages: messages.map(message => ({
          role: message._getType() === "human" ? "user" : "assistant",
          content: message.content,
        })),
        temperature: this.temperature,
        max_tokens: this.maxTokens,
        stream: false,
      }),
    });

    if (!response.ok) {
      throw new Error(`Deepseek API 请求失败: ${response.statusText}`);
    }

    const data = await response.json() as ChatCompletionResponse;
    const text = data.choices[0].message.content;

    return {
      generations: [{
        text,
        message: new AIMessage(text),
      }],
    };
  }

  /** 获取默认参数 */
  _llmType(): string {
    return "deepseek";
  }
}