import { ChatOpenAI } from "@langchain/openai";
import { HumanMessage, AIMessage, BaseMessage } from "@langchain/core/messages";
import { RojosConfig } from "./RojosConfig";

const model = new ChatOpenAI({
  model: RojosConfig.model,
  apiKey: RojosConfig.apiKey,
  configuration: {
    baseURL: RojosConfig.baseURL,
  },
});

type MessageData = { role: "user" | "assistant"; content: string };
type SessionData = { id: string; messages: MessageData[] };

const sessions = new Map<string, SessionData>();

function toBaseMessages( messages: MessageData[] ): BaseMessage[]
{
  return messages.map( m =>
    m.role === "user" ? new HumanMessage( m.content ) : new AIMessage( m.content )
  );
}

export async function getAgentStream( id: string, userMessage: string )
{
  if ( ! sessions.has( id ) )
  {
    sessions.set( id, { id, messages: [] } );
  }

  const session = sessions.get( id )!;
  session.messages.push( { role: "user", content: userMessage } );

  return model.stream( toBaseMessages( session.messages ) );
}

export function updateAgentConversation( id: string, assistantMessage: string )
{
  const session = sessions.get( id );

  if ( session )
  {
    session.messages.push( { role: "assistant", content: assistantMessage } );
  }
}