import { ChatOpenAI } from "@langchain/openai";
import { HumanMessage, AIMessage, SystemMessage, ToolMessage, BaseMessage } from "@langchain/core/messages";
import { RojoToolContext } from "./tools/RojoTool.js";

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

export type AgentEvent =
  | { type: "text";        text: string }
  | { type: "tool_call";  name: string; args: Record<string, unknown> }
  | { type: "tool_result"; name: string; result: string };

type ToolExecutor = ( name: string, args: Record<string, unknown>, ctx: RojoToolContext ) => Promise<unknown>;

type SessionData = { id: string; messages: BaseMessage[] };

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

export async function* getAgentStream(
  id:          string,
  userMessage: string,
  config:      AgentConfig,
  toolContext?: RojoToolContext,
  toolExecutor?: ToolExecutor,
): AsyncGenerator<AgentEvent>
{
  const model = new ChatOpenAI(
  {
    model:  config.model,
    apiKey: config.apiKey,
    configuration:
    {
      baseURL:        config.baseURL,
      defaultHeaders: config.headers ?? {},
    },
  } );

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

    if ( config.systemPrompt?.trim() )
    {
      messages.push( new SystemMessage( config.systemPrompt ) );
    }

    sessions.set( id, { id, messages } );
  }

  const session    = sessions.get( id )!;
  const tools      = config.tools ?? [];
  const boundModel = tools.length > 0 ? model.bindTools( tools ) : model;

  session.messages.push( new HumanMessage( userMessage ) );

  while ( true )
  {
    const stream     = await boundModel.stream( session.messages );
    let accumulated: any = null;

    for await ( const chunk of stream )
    {
      const text = extractText( chunk.content );

      if ( text ) yield { type: "text", text };

      accumulated = accumulated ? accumulated.concat( chunk ) : chunk;
    }

    const toolCalls: any[] = accumulated?.tool_calls ?? [];

    console.log( "[RojosAgent] loop end — tool_calls:", JSON.stringify( toolCalls ),
      "| text length:", extractText( accumulated?.content ).length,
      "| accumulated null:", accumulated === null || accumulated === undefined,
      "| content type:", typeof accumulated?.content,
      "| content:", JSON.stringify( accumulated?.content )?.slice( 0, 200 ) );

    if ( toolCalls.length > 0 && toolContext && toolExecutor )
    {
      session.messages.push( new AIMessage( { content: extractText( accumulated?.content ), tool_calls: toolCalls } ) );

      for ( const tc of toolCalls )
      {
        yield { type: "tool_call", name: tc.name, args: tc.args };

        let result: unknown;

        try
        {
          result = await toolExecutor( tc.name, tc.args, toolContext );
        }
        catch ( err )
        {
          result = { error: String( err ) };
        }

        const resultStr = typeof result === "string" ? result : JSON.stringify( result );

        yield { type: "tool_result", name: tc.name, result: resultStr };

        session.messages.push( new ToolMessage( { content: resultStr, tool_call_id: tc.id } ) );
      }
    }
    else
    {
      const text = extractText( accumulated?.content );

      if ( text ) session.messages.push( new AIMessage( text ) );

      break;
    }
  }
}

function extractText( content: unknown ): string
{
  if ( typeof content === "string" ) return content;

  if ( Array.isArray( content ) )
  {
    return content
      .map( c => typeof c === "string" ? c : ( ( c as any )?.text ?? "" ) )
      .join( "" );
  }

  return "";
}