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

export interface AgentConfig
{
  baseURL: string;
  model: string;
  apiKey: string;
  headers?: Record<string, string>;
  systemPrompt?: string;
}

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

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

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

export async function getAgentStream( id: string, userMessage: string, config: AgentConfig )
{
  const model = new ChatOpenAI(
  {
    model: config.model,
    apiKey: config.apiKey,
    configuration:
    {
      baseURL: config.baseURL,
      defaultHeaders: config.headers ?? {},
    },
  } );

  if ( !sessions.has( id ) )
  {
    const messages: MessageData[] = [];

    if ( config.systemPrompt?.trim() )
    {
      messages.push( { role: "system", content: config.systemPrompt } );
    }

    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 ): void
{
  const session = sessions.get( id );

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